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

@@ -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)
+        g8, amax_g = quantize(g2, sg.reciprocal(), fmt)
-        x8, _ = quantize(x, sx.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()

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(
+    rows = a.size(2) if trans_a else a.size(1)
-        (a.size(1) if trans_a else a.size(0), b.size(0) if trans_b else b.size(1)),
+    cols = b.size(1) if trans_b else b.size(2)
-        device=a.device,
+    batches = [t.size(0) for t in (a, b) if t.dim() == 3]
-        dtype=dtype,
+    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()
+    aa = a.float().transpose(-2, -1) if trans_a else a.float()
-    bb = b.float().t() if trans_b else b.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
+    ``a``/``b`` must be FP8 tensors of the same format, 2D or 3D (batched,
-    combined dequantization scale. The result is BF16; FP8 output is a separate
+    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

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

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 <typename Traits, typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
+template <typename Traits, typename LayoutA = RowMajor, typename LayoutB = RowMajor, int kRasterGroup = 0,
-        bool kBStaged = true, bool kLeanRing = false>
+        bool kBStaged = true, bool kLeanRing = false, bool kStreamOut = false>
 __global__ void __launch_bounds__(Traits::kCtaThreads,
                                   Fp8GemmSmem<Traits, LayoutA, LayoutB,
                                               kBStaged, kLeanRing>::kMinCtas)
@@ -390,9 +390,14 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
         reinterpret_cast<T8*>(fp8_gemm_smem + kARing * kAStageBytes);
     T8* const b_canon = b_base + kStB * kBStageBytes;  // staged B only
 
-    const auto* a = reinterpret_cast<const T8*>(p.a_ptr);
+    // Batch slice (grid.z): broadcast operands carry a 0 stride, so the
-    const auto* b = reinterpret_cast<const T8*>(p.b_ptr);
+    // same pointer serves every batch.
-    auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr);
+    const auto* a = reinterpret_cast<const T8*>(p.a_ptr) +
+                    (int64_t)blockIdx.z * p.a_batch_stride;
+    const auto* b = reinterpret_cast<const T8*>(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
+    // Tile scheduler: the linear CTA id maps to (block_m, block_n) in
-    // the linear block id so consecutive CTAs cover a group of kGroupM M-tiles
+    // grouped (L2-friendly, CUTLASS-style) or plain raster order — the
-    // before advancing along N. All CTAs of one group share the same B column
+    // grouped order makes consecutive CTAs cover a group of kRasterGroup
-    // stripe, so B tiles stay hot in L2 across the wave (the default
+    // M-tiles before advancing along N, so all CTAs of one group share the
-    // N-fastest order makes each wave touch every B tile instead).
+    // same B column stripe and B tiles stay hot in L2 across the wave (the
-    // kGroupRaster is a template knob (the launcher defaults it to the
+    // plain N-fastest order makes each wave touch every B tile instead;
-    // measured best per layout: grouped for A-crosswise (dW) and for the
+    // kRasterGroup=0 selects plain, the measured best for dX's crosswise-B
-    // congruous NT forward — whose big B operand gains the most from the
+    // layouts where grouping measured neutral).
-    // shared stripe — plain for dX's crosswise-B layouts, where it measured
+    // Persistent schedules (static round-robin and an atomic ticket
-    // neutral).
+    // dispenser, grid capped at the resident CTAs) were both measured and
-    constexpr int kGroupM = 8;
+    // 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;
@@ -644,6 +653,10 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
         // 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
@@ -733,7 +746,16 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
         const uint4 v = *reinterpret_cast<const uint4*>(out_chunk(r, c));
         auto* dst = out_bf16 + row * n + col;
         if (col + 8 <= n && (reinterpret_cast<uintptr_t>(dst) & 15) == 0) {
+            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<uint4*>(dst), v);
+            } else {
                 *reinterpret_cast<uint4*>(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
+// Crosswise-B staging+transpose vs the synchronous direct load: measured on
-// the contract dim is long enough that B streams from DRAM (dX-class GEMMs,
+// the current kernel generation, direct wins everywhere probed — contract k
-// k = N_ffn); short-K crosswise GEMMs (dW: k = M tokens) read L2-resident
+// 2048..32768 including B operands (128/256MB) that stream from DRAM past L2
-// operands, where the staging round trip costs more shared-memory traffic
+// (direct 171-181 TF vs staged 140-163 TF; the staging round trip costs more
-// than the latency it hides (measured: dW ~37 TF direct vs ~29 TF staged,
+// shared-memory traffic than the latency it hides). The old "stage past k=
-// dX ~39 TF staged vs ~38 direct).
+// 8192" rule reflected a pre-direct-path kernel; staging is now disabled.
-constexpr int64_t kCrossStageMinK = 8192;
+// 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
+// csrc/tests/fp8_sweep.cu): small outputs take 64x64 CTAs of 32x32 warps
-// CTAs on a 24-SM part — take 64x64 CTAs of 32x32 warps with a lean
+// with a lean (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and
-// (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and the extra
+// the extra blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs).
-// blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs). Everything
+// The large-output path takes the 128x128 CTA (8 warps x 64x32) with the
-// larger takes the 128x128 CTA (8 warps x 64x32) with the kStages+1 ring —
+// kStages+1 ring — one __syncthreads per k-tile and ~200 TF at scale.
-// one __syncthreads per k-tile. m <= 64 keeps the 64x128 CTA so a 128-row
+// Crossover (congruous NT, k=2048): 96 tiles small +14%, 112 tie, 135 big
-// tile never wastes half its MMA work on predicated-off rows.
+// +16% — threshold at ~2.3 waves of the resident (2/SM) 128x128 CTAs.
-constexpr int64_t kSmallShapeMaxTiles = 48;
+// 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 <FP8Format Fmt, typename LayoutA = RowMajor,
           typename LayoutB = RowMajor, int kK = 64, int Stages = 2,
-          bool GroupRaster = std::is_same_v<LayoutA, ColMajor> ||
+          int GroupRaster = (std::is_same_v<LayoutA, ColMajor> ||
-                             std::is_same_v<LayoutB, ColMajor>>
+                             std::is_same_v<LayoutB, ColMajor>)
+                                ? 8
+                                : 0>
 void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
-    const bool b_staged = p.k >= kCrossStageMinK;
+    // Staging is disabled (see kCrossStageMinK); the flag stays so the
-    if (p.m <= 64) {
+    // staged template instantiations below keep compiling for the sweep.
-        using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
+    const bool b_staged = false;
-        dim3 grid((p.n + 127) / 128, (p.m + 63) / 64);
+    // m <= 64 and small total outputs share the 64x64 small CTA; the
-        if (b_staged)
+    // threshold counts batch x per-matrix tiles (see small_shape_max_tiles).
-            launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
-                                             GroupRaster, true, true>>(
-                Fp8GemmSmem<Traits, LayoutA, LayoutB, true, true>::kBytes,
-                grid, dim3(Traits::kCtaThreads), stream, p);
-        else
-            launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
-                                             GroupRaster, false, true>>(
-                Fp8GemmSmem<Traits, LayoutA, LayoutB, false, true>::kBytes,
-                grid, dim3(Traits::kCtaThreads), stream, p);
-        return;
-    }
     const int64_t tiles_128 =
-        ((p.m + 127) / 128) * ((p.n + 127) / 128);
+        (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128);
-    if (tiles_128 < kSmallShapeMaxTiles) {
+    if (p.m <= 64 || tiles_128 < small_shape_max_tiles()) {
         using Traits = Fp8GemmTraits<Fmt, 64, 64, kK, 3, 32, 32>;
-        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<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
                                              GroupRaster, true, true>>(
@@ -835,7 +877,7 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
         return;
     }
     using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
-    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<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
                                          GroupRaster, true, false>>(

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<torch::Tensor, torch::Tensor> 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();
+    // Batched operands follow matmul broadcast rules: 2D acts as a batch
-    const bool ta = trans_a != 0;
+    // of 1; a size-1 batch broadcasts across the other side (stride 0).
-    const bool tb = trans_b != 0;
+    const int64_t batch_a = a.dim() == 3 ? a.size(0) : 1;
-    const int64_t a_ld = a_c.size(1);
+    const int64_t batch_b = b.dim() == 3 ? b.size(0) : 1;
-    const int64_t b_ld = b_c.size(1);
+    TORCH_CHECK(batch_a == batch_b || batch_a == 1 || batch_b == 1,
-    const int64_t m = ta ? a_c.size(1) : a_c.size(0);
+                "batch dim mismatch (got ", batch_a, " and ", batch_b, ")");
-    const int64_t k = ta ? a_c.size(0) : a_c.size(1);
+    const int64_t batch = std::max(batch_a, batch_b);
-    const int64_t n = tb ? b_c.size(0) : b_c.size(1);
+    TORCH_CHECK(batch <= 65535, "batch dim exceeds the grid.z launch limit");
-    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));
+    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,
+    pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
-              k, a_ld, b_ld);
+              m, n, k, a_ld, b_ld);
+    p.batch = static_cast<int>(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<FP8Format::E4M3>(p, stream.stream(), ta, tb);
+        dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
     else
-        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), ta, tb);
+        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
     C10_CUDA_CHECK(cudaGetLastError());
     return output;
 }

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 /

scripts/tools/bench_fp8_gemm.py

@@ -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()

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