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

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