fix: own fp8 linear backward via autograd Function

- backward used to read the global fp8 flag at loss.backward() time, so calling it outside fp8_autocast silently fell back to bf16 mm (953 ms cublas per step, 49.9% of the model step)
- _LinearFp8(torch.autograd.Function) now owns the fwd/bwd pair: forward captures fmt/recipe/meta on ctx inside the autocast region, backward reads only ctx (scales from the meta rings, masks from ctx.needs_input_grad), so backward is fp8 wherever it runs
- register the aten::linear impl on AutogradCUDA (replaces torch's generated linear formula that calls aten::linear_backward into the bf16 fallback) and keep the CUDA key for inference_mode
- drop the aten::linear_backward override and fp8_linear_backward (dead paths)
- regression test asserts the fp8 backward fires outside the autocast region and grads match the bf16 reference by direction/norm (E5M2 noise)
- model step (0.67B, CE loss, batch 4x1024): backward GEMMs 953 -> 618 ms (1.54x), full step ~1.2x

astrai/extension/fp8.py

@@ -15,9 +15,11 @@ Usage::
 
     with fp8_autocast(enabled=True, fp8_format="hybrid"):
         logits = model(input_ids)
-        loss.backward()
+    loss.backward()  # fp8 backward runs wherever it is called: the
+    # forward captures the fmt/recipe/meta on the autograd node
 
-Importing this module registers the aten::linear CUDA implementation.
+Importing this module registers the aten::linear CUDA and AutogradCUDA
+implementations.
 
 Format defaults follow the ecosystem consensus: E4M3 for the forward pass,
 E5M2 for the backward (gradient) pass ("hybrid"); every operand's scale is a
@@ -235,7 +237,7 @@ def fp8_autocast(
 
         with fp8_autocast(enabled=True, fp8_format="hybrid"):
             logits = model(input_ids)   # aten::linear -> fp8 path
-            loss.backward()
+        loss.backward()  # fp8 backward; state was captured at forward time
 
     Args:
         enabled: toggle fp8 dispatch for aten::linear.
@@ -298,29 +300,48 @@ def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
     return out
 
 
-def fp8_linear_backward(g: torch.Tensor, x: torch.Tensor, w: torch.Tensor, masks):
+class _LinearFp8(torch.autograd.Function):
-    """Scaled fp8 linear backward (called from aten::linear_backward).
+    """The fp8 linear forward/backward pair (standard Function style).
 
-    The gradient is quantized to the backward format (E5M2 in hybrid mode)
+    The forward runs inside the ``fp8_autocast`` region and captures the
-    and the dX / dW GEMMs share that single quantization.
+    active fmt/recipe/meta on ``ctx``; the backward reads only the captured
+    state, so ``loss.backward()`` may run after the context exits. The
+    gradient is quantized once (E5M2 in hybrid mode) and the dX / dW GEMMs
+    share that quantization; output masks come from ``needs_input_grad``.
     """
-    state = fp8_state()
+
-    fmt = state.fp8_format.bwd()
+    @staticmethod
-    meta = state.get_weight_meta(w)
+    def forward(ctx, x, w, bias):
-    if isinstance(state.recipe, DynamicScaling):
+        out = fp8_linear_forward(x, w, bias)
-        sg = _dynamic_scale(g, state.recipe, fmt)
+        state = fp8_state()
-        sw = _dynamic_scale(w, state.recipe, fmt)
+        ctx.save_for_backward(x, w)
-        sx = _dynamic_scale(x, state.recipe, fmt)
+        ctx.fmt_bwd = state.fp8_format.bwd()
-    else:
+        ctx.recipe = state.recipe
-        if not meta.g_init:
+        ctx.meta = state.get_weight_meta(w)
-            meta.init_g(g, fmt)
+        ctx.is_dynamic = isinstance(state.recipe, DynamicScaling)
-        sg, sw, sx = meta.g_scale, meta.w_scale, meta.x_scale
+        return out
-    grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
+
-        g, x, w, masks, sg, sw, sx, fmt
+    @staticmethod
-    )
+    @torch.autograd.function.once_differentiable
-    if not isinstance(state.recipe, DynamicScaling):
+    def backward(ctx, g):
-        meta.update_g(amax_g, fmt)
+        x, w = ctx.saved_tensors
-    return grad_x, grad_w, grad_b
+        fmt = ctx.fmt_bwd
+        if ctx.is_dynamic:
+            sg = _dynamic_scale(g, ctx.recipe, fmt)
+            sw = _dynamic_scale(w, ctx.recipe, fmt)
+            sx = _dynamic_scale(x, ctx.recipe, fmt)
+        else:
+            meta = ctx.meta
+            if not meta.g_init:
+                meta.init_g(g, fmt)
+            sg, sw, sx = meta.g_scale, meta.w_scale, meta.x_scale
+        masks = list(ctx.needs_input_grad)
+        grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
+            g, x, w, masks, sg, sw, sx, fmt
+        )
+        if not ctx.is_dynamic:
+            ctx.meta.update_g(amax_g, fmt)
+        return grad_x, grad_w, grad_b if masks[2] else None
 
 
 # ---------------------------------------------------------------------------
@@ -356,7 +377,7 @@ def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
         and w.dtype == torch.bfloat16
         and _fp8_supported(x, w)
     ):
-        return fp8_linear_forward(x, w, bias)
+        return _LinearFp8.apply(x, w, bias)
     return torch.ops.aten.linear.default.redispatch(
         torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd),
         x,
@@ -365,43 +386,12 @@ def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
     )
 
 
-def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask):
-    # Backward dim contract: grad_output is [..., N], weight is [N, K], so
-    # the contraction check is grad_output.size(-1) == weight.size(0) (not the
-    # forward's x.size(-1) == w.size(1) — that would silently skip fp8 for
-    # every non-square layer).
-    if (
-        fp8_linear_enabled()
-        and weight.dtype == torch.bfloat16
-        and grad_output.dim() >= 2
-        and weight.dim() == 2
-        and grad_output.size(-1) == weight.size(0)
-        and input_tensor.dim() >= 2
-        and input_tensor.size(-1) == weight.size(1)
-    ):
-        return fp8_linear_backward(grad_output, input_tensor, weight, list(output_mask))
-    compute_dtype = weight.dtype
-    grad = grad_output.to(compute_dtype)
-    grad_2d = grad.reshape(-1, weight.size(0))
-    input_2d = input_tensor.reshape(-1, input_tensor.size(-1)).to(compute_dtype)
-    # Unneeded grads come back full-shape-but-uninitialized (mirroring the
-    # fp8 binding), so reshape_as can never hit an empty tensor.
-    grad_input = (
-        torch.mm(grad_2d, weight).reshape_as(input_tensor)
-        if output_mask[0]
-        else torch.empty_like(input_tensor)
-    )
-    grad_weight = (
-        torch.mm(grad_2d.t(), input_2d) if output_mask[1] else torch.empty_like(weight)
-    )
-    grad_bias = (
-        grad.sum(dim=0)
-        if output_mask[2]
-        else torch.empty(0, device=grad.device, dtype=grad.dtype)
-    )
-    return grad_input, grad_weight, grad_bias
-
-
 _lib = Library("aten", "IMPL", "CUDA")
 _lib.impl("linear", _linear_cuda_impl)
-_lib.impl("linear_backward", _linear_backward_cuda_impl)
+# Also replace torch's generated linear autograd formula (which would call
+# aten::linear_backward after the fp8_autocast region exits). The fp8
+# backward is owned by _LinearFp8 with its state captured at forward time,
+# so loss.backward() works wherever it is called; the same CUDA registration
+# still covers inference_mode, where autograd keys are skipped entirely.
+_lib_autograd = Library("aten", "IMPL", "AutogradCUDA")
+_lib_autograd.impl("linear", _linear_cuda_impl)

tests/extension/test_fp8_mma.py

@@ -146,6 +146,57 @@ def test_linear_backward_e5m2_gradients():
     torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
 
 
+@skip_no_fp8
+def test_fp8_linear_backward_outside_autocast():
+    """aten::linear records an fp8 autograd node inside fp8_autocast; the
+    backward runs fp8 kernels even after the context exits (loss.backward()
+    placement is free), instead of falling back to bf16 mm."""
+    import torch.nn.functional as F
+
+    import astrai.extension.fp8 as f8mod
+
+    torch.manual_seed(5)
+    x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    weight = torch.randn(
+        96, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True
+    )
+    bias = torch.randn(96, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    xr, wr, br = (t.detach().clone().requires_grad_() for t in (x, weight, bias))
+
+    calls = {"bwd": 0}
+    orig = f8mod.linear_backward_fp8
+
+    def spy(g, xx, ww, masks, sg, sw, sx, fmt="e5m2"):
+        calls["bwd"] += 1
+        return orig(g, xx, ww, masks, sg, sw, sx, fmt)
+
+    f8mod.linear_backward_fp8 = spy
+    try:
+        with fp8_autocast(enabled=True):
+            out = F.linear(x, weight, bias)
+        assert type(out.grad_fn).__name__ == "_LinearFp8Backward"
+        out.float().pow(2).sum().backward()  # outside the autocast region
+    finally:
+        f8mod.linear_backward_fp8 = orig
+        f8mod.fp8_state().reset()
+
+    assert calls["bwd"] == 1  # fp8 kernels, not the bf16 fallback
+    ref = F.linear(xr, wr, br)
+    ref.float().pow(2).sum().backward()
+
+    # E5M2 backward quantization noise: compare directions/norms (the
+    # torchao/TE style) rather than elementwise against the bf16 reference.
+    def _direction(a, b):
+        cos = torch.nn.functional.cosine_similarity(
+            a.float().flatten(), b.float().flatten(), dim=0
+        )
+        return cos > 0.99 and 0.9 < a.float().norm() / b.float().norm() < 1.1
+
+    assert _direction(x.grad, xr.grad)
+    assert _direction(weight.grad, wr.grad)
+    assert _direction(bias.grad, br.grad)
+
+
 @skip_no_fp8
 def test_mm_fp8_matches_scaled_mm():
     torch.manual_seed(11)