perf: speed up fp8 gemm across small and large shapes

- parameterize warp tile (WarpM/WarpN) in Fp8GemmTraits; MMA loops, fragment arrays and epilogue scale with kMt/kNt instead of the fixed 64x32/4x4, enabling cuBLAS-style 64x64 CTAs of 32x32 warps
- dispatch by output tiling (grid-searched via csrc/tests/fp8_sweep.cu): fewer than 48 output tiles take 64x64/32x32 with a lean ring (4 CTAs/SM fill the wave-quantization gap: 512^3 goes 16 -> 64 CTAs); larger shapes keep 128x128 with the kStages+1 ring
- kStages+1 canonic ring rotation drops the post-compute barrier on the congruous path (one __syncthreads per k-tile); LeanRing keeps the kStages ring for the small CTA; direct-crosswise operands always rotate kStages+1 (their prefetch issues right after barrier 1 and would race a lean ring - caught by the pure C layout suite)
- stage the bf16 epilogue through the reclaimed operand smem: swizzled scatter + barrier + coalesced 16B copy-out replaces 8 disjoint 16B per-warp segments (~50% write efficiency before)
- hoist per-lane ldmatrix swizzle offsets out of the mainloop (stage-relative table + ring-base add) so the innermost loop stops recomputing IMAD/LOP3 address chains
- bypass the torch.library dispatch for real CUDA tensors in quantize/mm_fp8 wrappers (~5us/call, ~40% of a 512-wide call's wall time); fake/subclass tensors keep the custom_op route

vs the previous kernel + python path, wall clock on NT squares: 512^3 52 -> 13us (4.0x, 5.2 -> 20.5 TF, now 1.36x cuBLAS _scaled_mm), 1024^3 1.05x, 2048^3 1.02x (46.9 -> 48.2 TF kernel-only); correctness: 4 layouts x 6 shapes pure C suite PASS, 588 pytest PASS

scripts/tools/bench_fp8_gemm.py

@@ -0,0 +1,88 @@
+#!/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()