From 4d6a24409334988ef3673dbbb7de268d14fc5624 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 26 Aug 2026 06:45:14 +0800 Subject: [PATCH] perf: fp8 batched gemm and measured dispatch table - mm_fp8 accepts 3D operands through the same signature: grid.z slices by batch strides, size-1 batches broadcast (stride 0), inner .t() views fold into the layout tag at zero copy - fix _LinearFp8 backward crash on 3D [B,L,d] training inputs (flatten before mm_fp8, reduce grad_b over leading dims) - expose kRasterGroup/kStreamOut as template knobs; drop the 64x128 mid CTA and staged crosswise-B path from dispatch (direct wins everywhere re-measured, including DRAM-streamed B) - dispatch thresholds grounded in fresh sweeps: m<=64 -> 64x64 CTA (+27% at 64x8192x2048), small-CTA crossover at SM*14/3 total tiles (+13% at 96 tiles), threshold counts batch x per-matrix tiles (+31% at 64x512^3 bmm, +25% at 8x1024x2048) - remove scripts/tools/bench_fp8_gemm.py (superseded by csrc/tests/fp8_sweep.cu for kernel-level tuning) Benchmark: NVIDIA L20, E4M3, NT pre-quantized, median of 100-200 iters - 64x8192x2048: 29.1 -> 22.8 us (94 TF/s) - 1024x1536x2048: 67.4 -> 59.6 us (108 TF/s) - bmm 64x512^3: 139.8 -> 106.7 us; bmm 8x1024x2048: 186 TF/s - regression-free: 4096^3 192 TF/s, 8192^3 200 TF/s, 512^3 unchanged --- astrai/extension/fp8.py | 15 ++- astrai/extension/ops/fp8.py | 21 ++-- csrc/kernels/fp8/common.h | 7 ++ csrc/kernels/fp8/gemm.cuh | 190 +++++++++++++++++++------------- csrc/kernels/fp8/ops.cu | 79 ++++++++++--- docs/developer/cuda_kernels.md | 6 + scripts/tools/bench_fp8_gemm.py | 88 --------------- tests/extension/test_fp8_mma.py | 74 +++++++++++++ 8 files changed, 287 insertions(+), 193 deletions(-) delete mode 100644 scripts/tools/bench_fp8_gemm.py diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index 57c401c..b4e9c94 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -408,22 +408,25 @@ class _LinearFp8(torch.autograd.Function): def backward(ctx, g): x, w, _sx_fwd, _sw_fwd = ctx.saved_tensors fmt = ctx.fmt_bwd + # Flatten leading dims (the forward GEMMs ran on [-1, N] / [-1, K] + # views; the kernels only accept 2D operands). + g2 = g.reshape(-1, g.size(-1)) if ctx.is_dynamic: - sg = _dynamic_scale(g, ctx.recipe, fmt) + sg = _dynamic_scale(g2, ctx.recipe, fmt) sw = _dynamic_scale(w, ctx.recipe, fmt) sx = _dynamic_scale(x, ctx.recipe, fmt) else: meta = ctx.meta if not meta.g.initialized: - meta.g.seed(g, fmt) + meta.g.seed(g2, fmt) sg = meta.g.scale.clone() sw, sx = _sw_fwd, _sx_fwd - g8, amax_g = quantize(g, sg.reciprocal(), fmt) - x8, _ = quantize(x, sx.reciprocal(), fmt) + g8, amax_g = quantize(g2, sg.reciprocal(), fmt) + x8, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt) w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0] - grad_x = mm_fp8(g8, w8, sg * sw) # g8[m,n] @ w8[n,k] natural + grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape) # g8[m,n] @ w8[n,k] grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True) # g8.T @ x8 - grad_b = g.sum(0).to(torch.bfloat16) + grad_b = g2.sum(0).to(torch.bfloat16) if not ctx.is_dynamic: meta.g.update(amax_g, fmt) meta.g.advance() diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index d0e55e1..992c4e7 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -87,6 +87,7 @@ def fp8_gemm( ) -> torch.Tensor: """FP8 GEMM: ``a @ b * scale`` with FP32 accumulation. + 2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules). The result is always BF16; FP8 output is a separate quantize operation. """ @@ -94,11 +95,11 @@ def fp8_gemm( @fp8_gemm.register_fake def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0): dtype = torch.bfloat16 - return torch.empty( - (a.size(1) if trans_a else a.size(0), b.size(0) if trans_b else b.size(1)), - device=a.device, - dtype=dtype, - ) + rows = a.size(2) if trans_a else a.size(1) + cols = b.size(1) if trans_b else b.size(2) + batches = [t.size(0) for t in (a, b) if t.dim() == 3] + shape = (max(batches), rows, cols) if batches else (rows, cols) + return torch.empty(shape, device=a.device, dtype=dtype) @fp8_gemm.register_kernel("cuda") @@ -112,8 +113,8 @@ def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0): @fp8_gemm.register_kernel("cpu") def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0): - aa = a.float().t() if trans_a else a.float() - bb = b.float().t() if trans_b else b.float() + aa = a.float().transpose(-2, -1) if trans_a else a.float() + bb = b.float().transpose(-2, -1) if trans_b else b.float() acc = aa @ bb * scale return acc.to(torch.bfloat16) @@ -151,8 +152,10 @@ def mm_fp8( ) -> torch.Tensor: """Pre-quantized FP8 GEMM: ``a @ b * scale``. - ``a``/``b`` must be FP8 tensors of the same format. ``scale`` is their - combined dequantization scale. The result is BF16; FP8 output is a separate + ``a``/``b`` must be FP8 tensors of the same format, 2D or 3D (batched, + matmul-style broadcast on the batch dim). Inner-transposed views (e.g. + ``x.t()``) fold into the layout at zero copy. ``scale`` is their combined + dequantization scale. The result is BF16; FP8 output is a separate quantize operation. """ # Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 0163bc7..6d1ba0d 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -112,6 +112,13 @@ struct FP8Params { // to int64 for all pointer arithmetic. int m, n, k; + // Batched (bmm) geometry: grid.z slices step the operand/output pointers + // by these element strides (0 broadcasts the operand across batches). + int batch = 1; + int64_t a_batch_stride = 0; + int64_t b_batch_stride = 0; + int64_t out_batch_stride = 0; + // Physical leading dimensions (column count, i.e. row stride) of A and // B. For a non-transposed operand the stride equals the contract dim; // for a transposed operand it is the operand's own column count. The diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index 2dc604f..52cdfe6 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -339,8 +339,8 @@ struct Fp8GemmSmem { // (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the // launcher dispatches to it there (see launch_fp8_gemm). -template +template __global__ void __launch_bounds__(Traits::kCtaThreads, Fp8GemmSmem::kMinCtas) @@ -390,9 +390,14 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, reinterpret_cast(fp8_gemm_smem + kARing * kAStageBytes); T8* const b_canon = b_base + kStB * kBStageBytes; // staged B only - const auto* a = reinterpret_cast(p.a_ptr); - const auto* b = reinterpret_cast(p.b_ptr); - auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr); + // Batch slice (grid.z): broadcast operands carry a 0 stride, so the + // same pointer serves every batch. + const auto* a = reinterpret_cast(p.a_ptr) + + (int64_t)blockIdx.z * p.a_batch_stride; + const auto* b = reinterpret_cast(p.b_ptr) + + (int64_t)blockIdx.z * p.b_batch_stride; + auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr) + + (int64_t)blockIdx.z * p.out_batch_stride; const int64_t m = p.m, n = p.n, k = p.k; const int64_t a_ld = p.a_ld, b_ld = p.b_ld; @@ -401,19 +406,23 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, const int lane = tid & 31; const int group = lane >> 2; const int thread_in_group = lane & 3; - // L2-friendly rasterization (CUTLASS-style grouped launch order): remap - // the linear block id so consecutive CTAs cover a group of kGroupM M-tiles - // before advancing along N. All CTAs of one group share the same B column - // stripe, so B tiles stay hot in L2 across the wave (the default - // N-fastest order makes each wave touch every B tile instead). - // kGroupRaster is a template knob (the launcher defaults it to the - // measured best per layout: grouped for A-crosswise (dW) and for the - // congruous NT forward — whose big B operand gains the most from the - // shared stripe — plain for dX's crosswise-B layouts, where it measured - // neutral). - constexpr int kGroupM = 8; + // Tile scheduler: the linear CTA id maps to (block_m, block_n) in + // grouped (L2-friendly, CUTLASS-style) or plain raster order — the + // grouped order makes consecutive CTAs cover a group of kRasterGroup + // M-tiles before advancing along N, so all CTAs of one group share the + // same B column stripe and B tiles stay hot in L2 across the wave (the + // plain N-fastest order makes each wave touch every B tile instead; + // kRasterGroup=0 selects plain, the measured best for dX's crosswise-B + // layouts where grouping measured neutral). + // Persistent schedules (static round-robin and an atomic ticket + // dispenser, grid capped at the resident CTAs) were both measured and + // rejected on L20: the stride desynchronizes the in-flight window + // (-4..-8%), and the ticket variant recovers the L2 locality but lands + // within noise of plain waves (its loop-head barrier costs what the + // CTA-restart overlap saves). Keep the classic retiring-wave launch. int block_m, block_n; - if constexpr (kGroupRaster) { + if constexpr (kRasterGroup > 0) { + constexpr int kGroupM = kRasterGroup; const int blocks_m = gridDim.y; const int bid = blockIdx.y * gridDim.x + blockIdx.x; const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM; @@ -640,27 +649,31 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, if (k_seg + 1 < kSegs) transpose_tile(tile_index, k_seg + 1); } - // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 - // is issued before the MMAs consuming row mt, so the LDS fixed - // latency hides behind tensor-pipe work (cuts the `wait` stall, - // ~2.3 cycles/issue before this). Costs 4 extra registers. - unsigned a_frag[kMt + 1][4]; - astrai::ldmatrix_x4_lane(a_frag[0], a_base_addr + a_off[k_seg][0]); + // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 + // is issued before the MMAs consuming row mt, so the LDS fixed + // latency hides behind tensor-pipe work (cuts the `wait` stall, + // ~2.3 cycles/issue before this). Costs 4 extra registers. + // (Cross-k_seg prefetch of row 0 was tried and reverted: the + // register handoff broke ptxas's software pipelining — 171T → 95T + // at 2048³; the tensor pipe is issue-bound and the seg-start LDS + // already hides behind the b-fragment issue order.) + unsigned a_frag[kMt + 1][4]; + astrai::ldmatrix_x4_lane(a_frag[0], a_base_addr + a_off[k_seg][0]); #pragma unroll - for (int mt = 0; mt < kMt; ++mt) { - if (mt + 1 < kMt) - astrai::ldmatrix_x4_lane( - a_frag[mt + 1], a_base_addr + a_off[k_seg][mt + 1]); + for (int mt = 0; mt < kMt; ++mt) { + if (mt + 1 < kMt) + astrai::ldmatrix_x4_lane( + a_frag[mt + 1], a_base_addr + a_off[k_seg][mt + 1]); #pragma unroll - for (int nt = 0; nt < kNt; ++nt) - astrai::mma_sync(acc[nt][mt], a_frag[mt], - b_frag[bcur][nt], acc[nt][mt]); - } - // Barrier 3: region k_seg+1's transposes complete and become - // visible before the next k_seg reads them. - if constexpr (kBStagePath) { - if (k_seg + 1 < kSegs) __syncthreads(); - } + for (int nt = 0; nt < kNt; ++nt) + astrai::mma_sync(acc[nt][mt], a_frag[mt], + b_frag[bcur][nt], acc[nt][mt]); + } + // Barrier 3: region k_seg+1's transposes complete and become + // visible before the next k_seg reads them. + if constexpr (kBStagePath) { + if (k_seg + 1 < kSegs) __syncthreads(); + } } // Barrier 4 (staged-B / lean-ring only): every thread finished // reading this stage's tiles before the prefetch for the @@ -733,7 +746,16 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, const uint4 v = *reinterpret_cast(out_chunk(r, c)); auto* dst = out_bf16 + row * n + col; if (col + 8 <= n && (reinterpret_cast(dst) & 15) == 0) { - *reinterpret_cast(dst) = v; + if constexpr (kStreamOut) { + // Evict-first streaming store knob. Measured neutral on + // L20 squares and -3..4% on rects (the evict-first policy + // hurts more than the L2 B-tile protection helps at these + // sizes); kept as a template knob for other SKUs. Default + // off. + __stcs(reinterpret_cast(dst), v); + } else { + *reinterpret_cast(dst) = v; + } } else { // N-tail chunk or an odd-n row base: spill the elements that // survive the row edge (and stay aligned). @@ -748,6 +770,27 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, // Launchers — pure CUDA (no torch), usable from the binding and pure C tests. // --------------------------------------------------------------------------- +// SM count of the current device (cached per device; benign init race — +// every writer stores the same value). Host-side only: feeds the +// device-adaptive dispatch thresholds. +inline int device_sm_count() { + static int cached[64] = {}; + int dev = 0; + cudaGetDevice(&dev); + if (dev < 0 || dev >= 64) { + int sms = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + return sms > 0 ? sms : 1; + } + if (!cached[dev]) { + int sms = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + cached[dev] = sms > 0 ? sms : 1; + } + return cached[dev]; +} + + // Launch one kernel instantiation with its shared-memory budget: stages live // in dynamic smem, so budgets beyond the 48KB static limit opt in once per // instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared @@ -778,50 +821,49 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block, // per LayoutA (grouped for A-crosswise, plain for A-congruous). m <= 64 // dispatches to the 64x128 CTA — a 128-row CTA would waste half its MMA work // on predicated-off rows. -// Crosswise B takes the asynchronous staging+transpose pipeline only when -// the contract dim is long enough that B streams from DRAM (dX-class GEMMs, -// k = N_ffn); short-K crosswise GEMMs (dW: k = M tokens) read L2-resident -// operands, where the staging round trip costs more shared-memory traffic -// than the latency it hides (measured: dW ~37 TF direct vs ~29 TF staged, -// dX ~39 TF staged vs ~38 direct). -constexpr int64_t kCrossStageMinK = 8192; +// Crosswise-B staging+transpose vs the synchronous direct load: measured on +// the current kernel generation, direct wins everywhere probed — contract k +// 2048..32768 including B operands (128/256MB) that stream from DRAM past L2 +// (direct 171-181 TF vs staged 140-163 TF; the staging round trip costs more +// shared-memory traffic than the latency it hides). The old "stage past k= +// 8192" rule reflected a pre-direct-path kernel; staging is now disabled. +// The staged kernel template remains for csrc/tests/fp8_sweep.cu A/B runs. +constexpr int64_t kCrossStageMinK = (int64_t)1 << 62; // unreachable: never stage // Shape-based tile dispatch (grid-searched on the production shapes, see -// csrc/tests/fp8_sweep.cu): small outputs — fewer than ~2 waves of 128x128 -// CTAs on a 24-SM part — take 64x64 CTAs of 32x32 warps with a lean -// (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and the extra -// blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs). Everything -// larger takes the 128x128 CTA (8 warps x 64x32) with the kStages+1 ring — -// one __syncthreads per k-tile. m <= 64 keeps the 64x128 CTA so a 128-row -// tile never wastes half its MMA work on predicated-off rows. -constexpr int64_t kSmallShapeMaxTiles = 48; +// csrc/tests/fp8_sweep.cu): small outputs take 64x64 CTAs of 32x32 warps +// with a lean (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and +// the extra blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs). +// The large-output path takes the 128x128 CTA (8 warps x 64x32) with the +// kStages+1 ring — one __syncthreads per k-tile and ~200 TF at scale. +// Crossover (congruous NT, k=2048): 96 tiles small +14%, 112 tie, 135 big +// +16% — threshold at ~2.3 waves of the resident (2/SM) 128x128 CTAs. +// The threshold applies to the TOTAL tile count (batch x per-matrix tiles): +// batched runs keep full per-matrix CTA efficiency once the aggregate grid +// saturates the device (measured 64x512^3: big 160 vs small 123 TF — a +// per-matrix-only threshold lost 30%). m <= 64 always takes the small CTA: +// a 128-row CTA would waste half its MMA work on predicated-off rows. +inline int64_t small_shape_max_tiles() { + return (int64_t)device_sm_count() * 14 / 3; // 112 tiles on a 24-SM part +} template || - std::is_same_v> + int GroupRaster = (std::is_same_v || + std::is_same_v) + ? 8 + : 0> void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { - const bool b_staged = p.k >= kCrossStageMinK; - if (p.m <= 64) { - using Traits = Fp8GemmTraits; - dim3 grid((p.n + 127) / 128, (p.m + 63) / 64); - if (b_staged) - launch_with_smem>( - Fp8GemmSmem::kBytes, - grid, dim3(Traits::kCtaThreads), stream, p); - else - launch_with_smem>( - Fp8GemmSmem::kBytes, - grid, dim3(Traits::kCtaThreads), stream, p); - return; - } + // Staging is disabled (see kCrossStageMinK); the flag stays so the + // staged template instantiations below keep compiling for the sweep. + const bool b_staged = false; + // m <= 64 and small total outputs share the 64x64 small CTA; the + // threshold counts batch x per-matrix tiles (see small_shape_max_tiles). const int64_t tiles_128 = - ((p.m + 127) / 128) * ((p.n + 127) / 128); - if (tiles_128 < kSmallShapeMaxTiles) { + (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128); + if (p.m <= 64 || tiles_128 < small_shape_max_tiles()) { using Traits = Fp8GemmTraits; - dim3 grid((p.n + 63) / 64, (p.m + 63) / 64); + dim3 grid((p.n + 63) / 64, (p.m + 63) / 64, p.batch); if (b_staged) launch_with_smem>( @@ -835,7 +877,7 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { return; } using Traits = Fp8GemmTraits; - dim3 grid((p.n + 127) / 128, (p.m + 127) / 128); + dim3 grid((p.n + 127) / 128, (p.m + 127) / 128, p.batch); if (b_staged) launch_with_smem>( diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index 0f1c59f..7fcfc4e 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -90,6 +90,31 @@ void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a, } } +// Inner-layout resolution for one GEMM operand. The user flag names the +// math (0 = tensor's last two dims are [rows][contract], 1 = transposed); +// the storage may independently be a col-major view (.t() of a contiguous +// buffer), which folds into the returned dispatch flag at zero copy — the +// kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the +// user flag only; the fold never swaps them (see the layout table in +// gemm.cuh). Tensors whose inner dims are neither natural layout fall back +// to .contiguous(). +bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld, + int64_t& batch_stride, torch::Tensor& storage) { + torch::Tensor t = t_in; + bool col_major = false; + if (t.stride(-1) != 1) { + if (t.stride(-2) == 1) { + col_major = true; + } else { + t = t.contiguous(); + } + } + storage = t; + ld = col_major ? t.stride(-1) : t.stride(-2); + batch_stride = t.dim() == 3 ? t.stride(0) : 0; + return flag ^ col_major; +} + } // namespace std::tuple quantize(torch::Tensor x, @@ -145,30 +170,52 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale, a.scalar_type() == torch::kFloat8_e5m2, "a and b must be fp8"); TORCH_CHECK(a.scalar_type() == b.scalar_type(), "a and b must share format"); - TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D"); + TORCH_CHECK((a.dim() == 2 || a.dim() == 3) && + (b.dim() == 2 || b.dim() == 3), + "a and b must be 2D or 3D (batched)"); TORCH_CHECK(a.device() == b.device(), "a and b must share device"); check_scale(scale, a); check_fp8_device(a); const at::cuda::OptionalCUDAGuard guard(a.device()); auto stream = at::cuda::getCurrentCUDAStream(); - auto a_c = a.contiguous(); - auto b_c = b.contiguous(); - const bool ta = trans_a != 0; - const bool tb = trans_b != 0; - const int64_t a_ld = a_c.size(1); - const int64_t b_ld = b_c.size(1); - const int64_t m = ta ? a_c.size(1) : a_c.size(0); - const int64_t k = ta ? a_c.size(0) : a_c.size(1); - const int64_t n = tb ? b_c.size(0) : b_c.size(1); - TORCH_CHECK(k == (tb ? b_c.size(1) : b_c.size(0)), "inner dim mismatch"); - auto output = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16)); + + // Batched operands follow matmul broadcast rules: 2D acts as a batch + // of 1; a size-1 batch broadcasts across the other side (stride 0). + const int64_t batch_a = a.dim() == 3 ? a.size(0) : 1; + const int64_t batch_b = b.dim() == 3 ? b.size(0) : 1; + TORCH_CHECK(batch_a == batch_b || batch_a == 1 || batch_b == 1, + "batch dim mismatch (got ", batch_a, " and ", batch_b, ")"); + const int64_t batch = std::max(batch_a, batch_b); + TORCH_CHECK(batch <= 65535, "batch dim exceeds the grid.z launch limit"); + + torch::Tensor a_st, b_st; + int64_t a_ld, b_ld, a_bstride, b_bstride; + const bool tag_a = + resolve_operand(a, trans_a != 0, a_ld, a_bstride, a_st); + const bool tag_b = + resolve_operand(b, trans_b != 0, b_ld, b_bstride, b_st); + // GEMM dims from the user flags; storage layout never swaps them. + const int64_t m = trans_a ? a.size(-1) : a.size(-2); + const int64_t k = trans_a ? a.size(-2) : a.size(-1); + const int64_t n = trans_b ? b.size(-2) : b.size(-1); + TORCH_CHECK(k == (trans_b ? b.size(-1) : b.size(-2)), "inner dim mismatch"); + + const bool batched_out = a.dim() == 3 || b.dim() == 3; + torch::Tensor output = + batched_out + ? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16)) + : torch::empty({m, n}, a.options().dtype(torch::kBFloat16)); FP8Params p; - pack_gemm(p, a_c.data_ptr(), b_c.data_ptr(), output.data_ptr(), scale, m, n, - k, a_ld, b_ld); + pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale, + m, n, k, a_ld, b_ld); + p.batch = static_cast(batch); + p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride; + p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride; + p.out_batch_stride = m * n; if (a.scalar_type() == torch::kFloat8_e4m3fn) - dispatch_gemm(p, stream.stream(), ta, tb); + dispatch_gemm(p, stream.stream(), tag_a, tag_b); else - dispatch_gemm(p, stream.stream(), ta, tb); + dispatch_gemm(p, stream.stream(), tag_a, tag_b); C10_CUDA_CHECK(cudaGetLastError()); return output; } diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 8c4cf3f..b7cadaf 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -56,6 +56,12 @@ takes the combined dequant scale (`sa * sb`); the strategy layer passes `scale.reciprocal()` / `sa * sb` respectively. `amax` is always returned in the original input domain. +`mm_fp8` also accepts 3D (batched) operands through the same signature: +`grid.z` slices the operands by their batch strides, a size-1 batch +broadcasts (stride 0), and inner-transposed views (e.g. `x.t()`) fold into +the kernel's layout tag at zero copy — only genuinely strided operands pay +a `.contiguous()` copy. + Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless primitives (`quantize` / `mm_fp8`) via `torch.library.custom_op`, and `astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed / diff --git a/scripts/tools/bench_fp8_gemm.py b/scripts/tools/bench_fp8_gemm.py deleted file mode 100644 index 79a246f..0000000 --- a/scripts/tools/bench_fp8_gemm.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -"""FP8 GEMM micro-benchmark: AstrAI kernel vs cuBLAS (torch._scaled_mm). - -Sizes 512-8192, forward (NT) layout: x8[M,K] @ w8[N,K]^T -> bf16. -cuBLAS reference uses the same pre-quantized fp8 operands and the same -combined scale, so the comparison isolates the GEMM loop itself. - - python scripts/tools/bench_fp8_gemm.py --sizes 512 1024 2048 -""" - -import argparse -import sys - -import torch - -sys.path.insert(0, ".") - -from astrai.extension import loader - - -def bench(fn, iters=50, warmup=10): - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - fn() - end.record() - torch.cuda.synchronize() - return start.elapsed_time(end) / iters # ms - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--sizes", type=int, nargs="+", default=[512, 1024, 2048, 4096, 8192] - ) - parser.add_argument( - "--rect", action="store_true", help="also bench M=4096xN=1024 style rectangles" - ) - parser.add_argument("--iters", type=int, default=50) - args = parser.parse_args() - - assert loader.is_available("fp8_ops"), "fp8_ops extension not built" - from astrai.extension.ops.fp8 import mm_fp8 - - dev = torch.device("cuda") - torch.manual_seed(0) - - shapes = [(s, s, s) for s in args.sizes] - if args.rect: - shapes += [(4096, 1024, 4096), (8192, 4096, 8192), (2048, 8192, 2048)] - - print( - f"{'M':>6} {'N':>6} {'K':>6} | {'ours(ms)':>9} {'TFLOPs':>7} | " - f"{'cublas(ms)':>10} {'TFLOPs':>7} | {'ratio':>6}" - ) - print("-" * 72) - for m, n, k in shapes: - x8 = (torch.randn(m, k, device=dev) * 0.05).to(torch.float8_e4m3fn) - w8 = (torch.randn(n, k, device=dev) * 0.05).to(torch.float8_e4m3fn) - scale = torch.ones(1, device=dev, dtype=torch.float32) - - # ours: NT (x8 @ w8^T, LayoutB=ColMajor = weight layout) - t_ours = bench(lambda: mm_fp8(x8, w8, scale, trans_b=True), iters=args.iters) - # cuBLAS: _scaled_mm needs A row-major, B column-major (= w8.t()) - wt = w8.t() - sa = torch.ones(1, device=dev) - sb = torch.ones(1, device=dev) - - def cublas(): - return torch._scaled_mm(x8, wt, sa, sb, out_dtype=torch.bfloat16) - - t_cublas = bench(cublas, iters=args.iters) - flops = 2.0 * m * n * k - tf_ours = flops / (t_ours * 1e-3) / 1e12 - tf_cublas = flops / (t_cublas * 1e-3) / 1e12 - print( - f"{m:>6} {n:>6} {k:>6} | {t_ours:>9.3f} {tf_ours:>7.1f} | " - f"{t_cublas:>10.3f} {tf_cublas:>7.1f} | " - f"{tf_ours / tf_cublas:>5.0%}" - ) - - -if __name__ == "__main__": - main() diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index 4d1ff52..afba3a1 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -116,6 +116,80 @@ def test_mm_fp8_transposed_operands(trans_a, trans_b): torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) +@skip_no_fp8 +@pytest.mark.parametrize("trans_a", [False, True]) +@pytest.mark.parametrize("trans_b", [False, True]) +def test_mm_fp8_batched(trans_a, trans_b): + """3D operands run as one bmm launch: all four layouts, odd shapes.""" + torch.manual_seed(23) + batch, m, n, k = 4, 19, 13, 37 + a = torch.randn(batch, m, k, device="cuda", dtype=torch.bfloat16) + b = torch.randn(batch, n, k, device="cuda", dtype=torch.bfloat16) + sa, sb = _scale(a), _scale(b) + a8, _ = quantize(a, sa.reciprocal(), "e4m3") + b8, _ = quantize(b, sb.reciprocal(), "e4m3") + a_op = a8.transpose(-2, -1).contiguous() if trans_a else a8 + b_op = b8 if trans_b else b8.transpose(-2, -1).contiguous() + + out = mm_fp8(a_op, b_op, sa * sb, trans_a=trans_a, trans_b=trans_b) + assert out.shape == (batch, m, n) + # flags + transposed buffers reconstruct the original operands: the math + # is always A_orig @ B_orig^T regardless of the layout combination. + expected = (_quantize(a, sa) @ _quantize(b, sb).transpose(-2, -1) * sa * sb).to( + torch.bfloat16 + ) + torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) + + +@skip_no_fp8 +def test_mm_fp8_batched_broadcast(): + """A size-1 batch broadcasts across the other operand (matmul rules), + and a 2D operand broadcasts across a 3D one.""" + torch.manual_seed(29) + batch, m, n, k = 3, 16, 8, 32 + a = torch.randn(batch, m, k, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, n, k, device="cuda", dtype=torch.bfloat16) + sa, sb = _scale(a), _scale(b) + a8, _ = quantize(a, sa.reciprocal(), "e4m3") + b8, _ = quantize(b, sb.reciprocal(), "e4m3") + + out = mm_fp8(a8, b8, sa * sb, trans_b=True) + assert out.shape == (batch, m, n) + expected = (_quantize(a, sa) @ _quantize(b, sb).transpose(-2, -1) * sa * sb).to( + torch.bfloat16 + ) + torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) + + # 2D weight broadcast over 3D activations + w8 = b8[0] + out2 = mm_fp8(a8, w8, sa * sb, trans_b=True) + assert out2.shape == (batch, m, n) + torch.testing.assert_close(out2, expected, atol=0.125, rtol=0.01) + + +@skip_no_fp8 +def test_mm_fp8_col_major_view_zero_copy(): + """An inner-transposed view (.t() of a contiguous buffer) folds into the + layout tag with no device copy — the only allocation is the output.""" + torch.manual_seed(31) + m, n, k = 64, 64, 64 + a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + sa, sb = _scale(a), _scale(b) + a8, _ = quantize(a, sa.reciprocal(), "e4m3") + b8, _ = quantize(b, sb.reciprocal(), "e4m3") + + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + out = mm_fp8(a8.t(), b8, sa * sb, trans_a=True, trans_b=True) + torch.cuda.synchronize() + grew = torch.cuda.memory_allocated() - before + assert grew == out.numel() * out.element_size() # no operand copy + + expected = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16) + torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) + + @skip_no_fp8 def test_delayed_scaling_forward_uses_snapshot_scale(): """The delayed scale for step N is computed from amax(steps < N); the