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

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)