feat: add fp8 training via cublasLt dispatch

- fp8_mm kernel (csrc): cublasLt fp8 e4m3 gemm, TN layout mapped zero-copy
- custom::fp8_mm custom op: meta/cuda/cpu kernels + scale-corrected bf16 autograd
- aten::linear and linear_backward dispatch on CUDA key, zero model changes
- per-tensor scale or raw cast; single-GPU smoke loss matches bf16

astrai/extension/fp8_ops.py

@@ -0,0 +1,88 @@
+"""FP8 matrix-multiply op (torch.library custom_op) and FP8 linear replacement.
+
+Dispatch table:
+- Meta (register_fake): shapes only, for torch.compile / dynamic shapes
+- CUDA: csrc fp8_mm kernel (cuBLASLt TN fp8 GEMM, e4m3 in, fp32 acc/out)
+- CPU: fp32 fallback (testing)
+- AutogradCUDA (register_autograd): bf16 backward, scale-corrected
+"""
+
+import torch
+from torch.library import custom_op
+
+from astrai.extension.loader import get_module, is_available
+
+
+@custom_op("custom::fp8_mm", mutates_args=())
+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] -> fp32[M,N], scales applied by the caller.
+
+    a/b arrive pre-scaled (divided by sx/sw) fp8 tensors; the op returns the
+    unscaled fp32 result so scale math stays in autograd-land.
+    """
+
+
+@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.float32)
+
+
+@fp8_mm.register_kernel("cuda")
+def _fp8_mm_cuda(a, b, sx, sw):
+    if not is_available("fp8_mm"):
+        raise RuntimeError(
+            "CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true."
+        )
+    return get_module("fp8_mm").fp8_mm(a, b)
+
+
+@fp8_mm.register_kernel("cpu")
+def _fp8_mm_cpu(a, b, sx, sw):
+    return torch.mm(a.float(), b.float().t())
+
+
+def _fp8_mm_setup_context(ctx, inputs, output):
+    ctx.save_for_backward(*inputs)
+
+
+def _fp8_mm_backward(ctx, g):
+    """Scale-corrected straight-through gradients.
+
+    out = F(a, b) * (sx * sw) with F(a, b) = a @ b^T, a = x/sx, b = w/sw:
+      dx = g * sw @ b          (dout/dx = dF/da * 1/sx * sx*sw)
+      dW = (g * sx)^T @ a      (dout/dw = dF/db * 1/sw * sx*sw)
+    bf16 GEMMs keep gradients in range (e4m3 saturates at 448).
+    """
+    a, b, sx, sw = ctx.saved_tensors
+    ga = torch.mm(g * sw, b.float())
+    gb = torch.mm((g * sx).t(), a.float())
+    return ga.to(torch.bfloat16), gb.to(torch.bfloat16), None, None
+
+
+fp8_mm.register_autograd(_fp8_mm_backward, setup_context=_fp8_mm_setup_context)
+
+
+def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
+    """FP8 replacement for F.linear(x, w, bias).
+
+    x: [..., K] bf16 (any leading dims), w: [N,K] bf16 (in_dim=K).
+    The kernel computes a @ b^T with zero-copy col-major mapping, so w is
+    passed as-is (no transpose).
+    """
+    orig_shape = x.shape
+    x2d = x.reshape(-1, w.size(1))
+    sx = x2d.abs().amax() / 448.0
+    sw = w.abs().amax() / 448.0
+    x8 = (x2d / sx).to(torch.float8_e4m3fn)
+    w8 = (w / sw).to(torch.float8_e4m3fn)
+    out = torch.ops.custom.fp8_mm(x8, w8, sx, sw)
+    out = out * (sx * sw)
+    if bias is not None:
+        out = out + bias
+    return out.reshape(*orig_shape[:-1], -1)
+
+
+def fp8_available() -> bool:
+    return is_available("fp8_mm")