perf: optimize fused FP8 GEMM kernel

astrai/extension/ops/fp8.py

@@ -27,17 +27,19 @@ def _mod():
 def fp8_mm(
     a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor
 ) -> torch.Tensor:
-    """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)."""
+    """BF16 inputs, fused FP8 GEMM with FP32 accumulation and BF16 output."""
 
 
 @fp8_mm.register_fake
 def _fp8_mm_fake(a, b, sx, sw):
-    return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.bfloat16)
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
 
 
 @fp8_mm.register_kernel("cuda")
 def _fp8_mm_cuda(a, b, sx, sw):
-    return _mod().fp8_mm(a, b)
+    if not (a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16):
+        raise TypeError(f"bf16 GEMM requires bf16 inputs, got {a.dtype}/{b.dtype}")
+    return _mod().fp8_mm(a, b, sx, sw)
 
 
 @fp8_mm.register_kernel("cpu")
@@ -46,10 +48,10 @@ def _fp8_mm_cpu(a, b, sx, sw):
 
 
 def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w):
-    """Quantize x/w with per-tensor scales + cuBLASLt GEMM + bias -> bf16.
+    """Quantize BF16 inputs to FP8, accumulate in FP32, and return BF16.
 
-    x/w: [..., K] / [N, K] bf16; sx/sw: f32 scale tensors (device scalars);
-    sx_inv/sw_inv: 1/scale; amax_x/amax_w: f32 buffers receiving max-abs.
+    x/w: [..., K] / [N, K] bf16; sx/sw and their inverses control the fused
+    E4M3 conversion; amax_x/amax_w receive the input max-abs values.
     """
     if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16):
         raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")