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

astrai/extension/ops/fp8.py

@@ -127,6 +127,18 @@ def quantize(
     ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
     E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
     """
+    # Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a
+    # 512-wide GEMM): real CUDA tensors of a supported dtype go straight to
+    # the extension. Fake/subclass tensors and non-CUDA inputs keep the
+    # custom_op route so torch.compile / meta / fake-tensor tracing and the
+    # CPU fallback behave exactly as before.
+    if (
+        type(x) is torch.Tensor
+        and x.is_cuda
+        and x.dtype in _QUANT_INPUT_DTYPES
+        and fmt in _FMT_TO_INT
+    ):
+        return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt])
     return fp8_quantize(x, scale, _fmt_int(fmt))
 
 
@@ -143,4 +155,12 @@ def mm_fp8(
     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
+    # validation identical on the direct route.
+    if (
+        type(a) is torch.Tensor
+        and a.is_cuda
+        and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+    ):
+        return get_module("fp8_ops").mm_fp8(a, b, scale, int(trans_a), int(trans_b))
     return fp8_gemm(a, b, scale, trans_a, trans_b)