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_dispatch.py

@@ -0,0 +1,67 @@
+"""FP8 linear dispatch: replace aten::linear on the CUDA key, no model changes.
+
+``F.linear`` -> ``aten::linear`` -> dispatcher -> this CUDA impl (fp8 when
+enabled) or the original composite implementation via ``redispatch``.
+Enabling is per-thread; model code stays untouched.
+"""
+
+import threading
+
+import torch
+from torch.library import Library
+
+from astrai.extension.fp8_ops import fp8_linear_forward
+
+_state = threading.local()
+
+
+def fp8_linear_enable(enabled: bool = True) -> None:
+    """Toggle fp8 dispatch for aten::linear on this thread."""
+    _state.enabled = enabled
+
+
+def fp8_linear_enabled() -> bool:
+    return getattr(_state, "enabled", False)
+
+
+def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
+    if fp8_linear_enabled() and x.dtype in (torch.bfloat16, torch.float32):
+        return fp8_linear_forward(x, w, bias)
+    return torch.ops.aten.linear.default.redispatch(
+        torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd),
+        x,
+        w,
+        bias,
+    )
+
+
+def _linear_backward_cuda_impl(input, grad_output, weight, output_mask):
+    # VariableType wraps aten::linear; its backward runs aten::linear_backward
+    # with schema (self, grad_output, weight, mask). Implement the bf16
+    # gradient math directly (no redispatch), supporting [..., K] inputs:
+    #   dX = g @ W, dW = g^T @ X, dB = sum(g, dim=0)
+    g = grad_output.to(torch.bfloat16)
+    g2d = g.reshape(-1, weight.size(0))
+    x2d = input.reshape(-1, input.size(-1)).to(torch.bfloat16)
+    dX = (
+        torch.mm(g2d, weight)
+        if output_mask[0]
+        else torch.empty(0, device=input.device, dtype=input.dtype)
+    )
+    dX = dX.reshape_as(input)
+    dW = (
+        torch.mm(g2d.t(), x2d)
+        if output_mask[1]
+        else torch.empty(0, device=input.device, dtype=input.dtype)
+    )
+    dB = (
+        g.sum(dim=0)
+        if output_mask[2]
+        else torch.empty(0, device=input.device, dtype=input.dtype)
+    )
+    return dX, dW, dB
+
+
+_lib = Library("aten", "IMPL", "CUDA")
+_lib.impl("linear", _linear_cuda_impl)
+_lib.impl("linear_backward", _linear_backward_cuda_impl)

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

astrai/extension/loader.py

@@ -17,6 +17,7 @@ KERNEL_NAMES = [
     "attn_paged_decode",
     "attn_paged_prefill",
     "rotary_emb",
+    "fp8_mm",
 ]
 
 _available: dict[str, bool] = {}

csrc/CMakeLists.txt

@@ -48,7 +48,7 @@ set(TORCH_LIBS
 
 set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
 
-set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb)
+set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb fp8_mm)
 
 foreach(name ${KERNELS})
     add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu")
@@ -61,6 +61,9 @@ foreach(name ${KERNELS})
         "${PYTHON_INCLUDE_DIR}")
 
     target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
+    if(${name} STREQUAL "fp8_mm")
+        target_link_libraries(${name} PRIVATE CUDA::cublasLt)
+    endif()
     target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
 
     target_compile_options(${name} PRIVATE

csrc/kernels/fp8_mm.cu

@@ -0,0 +1,108 @@
+// FP8 e4m3 matrix multiply via cuBLASLt (sm89 TN layout).
+//
+// cuBLASLt exposes fp8 kernels only for op(A)=T, op(B)=N on Ada; we exploit
+// the identity: row-major a[M,K] == A^T as col-major [K,M] (zero copy), and
+// row-major wT[N,K] == B as col-major [K,N] (zero copy). The col-major
+// result D[M,N] is C^T in row-major terms, so we transpose the output once.
+//
+// Inputs arrive pre-scaled fp8 e4m3 tensors; output is unscaled fp32.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cublasLt.h>
+#include <cstdint>
+
+static cublasLtHandle_t g_handle = nullptr;
+static cublasLtMatmulDesc_t g_desc = nullptr;
+static cublasLtMatrixLayout_t g_layout_a = nullptr;
+static cublasLtMatrixLayout_t g_layout_b = nullptr;
+static cublasLtMatrixLayout_t g_layout_c = nullptr;
+static cublasLtMatmulPreference_t g_pref = nullptr;
+static void* g_workspace = nullptr;
+static size_t g_ws_size = 0;
+
+static void ensure_cublas_lt() {
+    if (g_handle) {
+        return;
+    }
+    TORCH_CHECK(cublasLtCreate(&g_handle) == CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatmulDescCreate(&g_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F) ==
+                CUBLAS_STATUS_SUCCESS);
+    cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
+    cublasLtMatmulDescSetAttribute(g_desc, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta));
+    cublasLtMatmulDescSetAttribute(g_desc, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb));
+    TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_a, CUDA_R_8F_E4M3, 1, 1, 1) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_b, CUDA_R_8F_E4M3, 1, 1, 1) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_c, CUDA_R_32F, 1, 1, 1) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatmulPreferenceCreate(&g_pref) == CUBLAS_STATUS_SUCCESS);
+    size_t ws = 16 * 1024 * 1024;
+    TORCH_CHECK(cublasLtMatmulPreferenceSetAttribute(
+                   g_pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws, sizeof(ws)) ==
+               CUBLAS_STATUS_SUCCESS);
+}
+
+static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols,
+                       int64_t ld) {
+    TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_ROWS,
+                                                 &rows, sizeof(rows)) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_COLS,
+                                                 &cols, sizeof(cols)) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_LD,
+                                                 &ld, sizeof(ld)) ==
+                CUBLAS_STATUS_SUCCESS);
+}
+
+torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) {
+    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn, "a must be float8_e4m3fn");
+    TORCH_CHECK(b.scalar_type() == torch::kFloat8_e4m3fn, "b must be float8_e4m3fn");
+    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "2D tensors required");
+    const at::cuda::OptionalCUDAGuard guard(a.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto a_c = a.contiguous();
+    auto b_c = b.contiguous();
+    int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
+    TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch");
+
+    auto buf = torch::empty({n, m}, a_c.options().dtype(torch::kFloat32));
+
+    ensure_cublas_lt();
+    set_layout(g_layout_a, k, m, k);  // A col-major [K,M] (a row-major, op=T)
+    set_layout(g_layout_b, k, n, k);  // B col-major [K,N] (wT row-major, op=N)
+    set_layout(g_layout_c, m, n, m);  // C col-major [M,N]
+
+    float alpha = 1.0f, beta = 0.0f;
+    cublasLtMatmulHeuristicResult_t heur;
+    int returned = 0;
+    cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
+        g_handle, g_desc, g_layout_a, g_layout_b, g_layout_c, g_layout_c, g_pref, 1,
+        &heur, &returned);
+    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
+                "cublasLtMatmulAlgoGetHeuristic failed: ", cublasLtGetStatusName(st));
+    if (heur.workspaceSize > g_ws_size) {
+        if (g_workspace) {
+            cudaFree(g_workspace);
+        }
+        TORCH_CHECK(cudaMalloc(&g_workspace, heur.workspaceSize) == cudaSuccess);
+        g_ws_size = heur.workspaceSize;
+    }
+    st = cublasLtMatmul(
+        g_handle, g_desc, &alpha, a_c.data_ptr(), g_layout_a, b_c.data_ptr(),
+        g_layout_b, &beta, buf.data_ptr(), g_layout_c, buf.data_ptr(), g_layout_c,
+        &heur.algo, g_workspace, g_ws_size, stream.stream());
+    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
+                "cublasLtMatmul failed: ", cublasLtGetStatusName(st));
+    return buf.transpose(0, 1).contiguous();
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"),
+          "FP8 e4m3 GEMM: a[M,K] x b[N,K] -> fp32[M,N] (pre-scaled inputs)");
+}