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