feat: fp8 backward reuses pre-quantized operands

- g/x/w may each be bf16 or pre-quantized fp8 matching fmt; a pre-quantized operand skips its quantize kernel
- snapshot sx/sw/sg before the ring finalize overwrites the aliased scale slot so the gemm dequantizes with the quantize scale
- forward carries its scale to backward so gradients reuse the forward's scale
- grad_input/grad_weight forced bf16; a pre-quantized g dequantizes before the bias-sum
- regression test: two delayed steps with a changing amax must not leak the scale ratio

astrai/extension/fp8.py

@@ -344,25 +344,31 @@ def fp8_linear_forward(
         sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
         sw = _dynamic_scale(w, cfg.recipe, fmt)
         out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt)
-        return out
+        return out, sx, sw
 
     meta = state.get_weight_meta(w)
     if not meta.w.initialized:
         meta.w.seed(w, fmt)
     if not meta.x.initialized:
         meta.x.seed(x, fmt)
-    # The kernel finalizes each ring in-kernel (overwriting the scale slot), so
+    # The quantize kernels finalize each ring in-kernel and overwrite the ring's
-    # the w/x scales are taken from the ring before the quantize.
+    # scale slot, which ALIASES meta.*.scale (a view into the state buffer).
+    # Snapshot the scales first so the GEMM dequantizes with the SAME scale the
+    # operands were quantized with, and so the backward can reuse this step's
+    # scale (gradient consistency with the forward). The ring finalize still
+    # publishes the next step's scale into the original slot.
     if w.dtype is not torch.bfloat16:  # static pre-quantized weight
         w_arg, sw_arg, w_ring = w, meta.w.scale, None
     else:
         w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state
+    sx = meta.x.scale.clone()
+    sw = sw_arg.clone()
     out, _x8, _w8, _ax, _aw = linear_forward_fp8(
         x,
         w_arg,
         bias,
-        meta.x.scale,
+        sx,
-        sw_arg,
+        sw,
         fmt,
         None,
         meta.x.state,
@@ -375,7 +381,7 @@ def fp8_linear_forward(
     meta.x.advance()
     if w_ring is not None:
         meta.w.advance()
-    return out
+    return out, sx, sw
 
 
 class _LinearFp8(torch.autograd.Function):
@@ -391,8 +397,8 @@ class _LinearFp8(torch.autograd.Function):
     @staticmethod
     def forward(ctx, x, w, bias):
         cfg = _current_config()
-        out = fp8_linear_forward(x, w, bias, cfg)
+        out, sx, sw = fp8_linear_forward(x, w, bias, cfg)
-        ctx.save_for_backward(x, w)
+        ctx.save_for_backward(x, w, sx, sw)
         ctx.fmt_bwd = cfg.fp8_format.bwd()
         ctx.recipe = cfg.recipe
         ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
@@ -402,7 +408,7 @@ class _LinearFp8(torch.autograd.Function):
     @staticmethod
     @torch.autograd.function.once_differentiable
     def backward(ctx, g):
-        x, w = ctx.saved_tensors
+        x, w, _sx_fwd, _sw_fwd = ctx.saved_tensors
         fmt = ctx.fmt_bwd
         # Per-recipe scale/ring selection; both branches share one call below.
         if ctx.is_dynamic:
@@ -414,8 +420,12 @@ class _LinearFp8(torch.autograd.Function):
             meta = ctx.meta
             if not meta.g.initialized:
                 meta.g.seed(g, fmt)
-            sg, ring, idx = meta.g.scale, meta.g.state, meta.g.idx
+            # Snapshot the g scale before its ring finalize overwrites the slot
-            sw, sx = meta.w.scale, meta.x.scale
+            # (same aliasing as the forward); reuse the forward's w/x scales so
+            # the backward quantizes with the scale the forward actually used.
+            sg = meta.g.scale.clone()
+            ring, idx = meta.g.state, meta.g.idx
+            sw, sx = _sw_fwd, _sx_fwd
         grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8(
             g,
             x,

astrai/extension/ops/fp8.py

@@ -212,20 +212,24 @@ def linear_backward_fp8(
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``.
 
-    The gradient (and the transposed w/x operands) are quantized to ``fmt``
+    ``g``/``x``/``w`` may each be bf16 (quantized to ``fmt`` here) or already
-    (default E5M2 — larger dynamic range for gradients) and the two GEMMs run
+    pre-quantized fp8 matching ``fmt`` — a pre-quantized operand skips its
-    as FP8 tensor-core products sharing a single gradient quantization.
+    quantize kernel and is read directly by the GEMM (the ``cached_cast``
-    ``g_ring`` (delayed scaling) is a ``[hist | scale | counter]`` buffer the
+    analog for the backward, symmetric with :func:`linear_forward_fp8`'s
-    g quantize kernel finalizes in-kernel (see :func:`linear_forward_fp8`).
+    pre-quantized weight path). ``fmt`` defaults to E5M2 (larger dynamic range
+    for gradients); the two GEMMs run as FP8 tensor-core products sharing a
+    single gradient quantization. ``g_ring`` (delayed scaling) is a
+    ``[hist | scale | counter]`` buffer the g quantize kernel finalizes
+    in-kernel (see :func:`linear_forward_fp8`); a pre-quantized ``g`` does not
+    finalize it and reports ``amax_g = 0``.
     """
-    if not (
+    f8 = _fmt_dtype(fmt)
-        g.dtype == torch.bfloat16
+    for name, t in (("g", g), ("x", x), ("w", w)):
-        and x.dtype == torch.bfloat16
+        if t.dtype not in (torch.bfloat16, f8):
-        and w.dtype == torch.bfloat16
+            raise TypeError(
-    ):
+                f"fp8 backward requires bf16 or pre-quantized {fmt} inputs, "
-        raise TypeError(
+                f"got {name}={t.dtype}"
-            f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
+            )
-        )
     return get_module("fp8_ops").linear_backward_fp8(
         g,
         x,

csrc/kernels/fp8/ops.cu

@@ -357,10 +357,19 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     // the w/x quantizes for dX / dW never touch rings — each operand's ring
     // is finalized exactly once per step (by the forward or this kernel).
     TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required");
-    TORCH_CHECK(g.scalar_type() == torch::kBFloat16 &&
+    // Each operand may be bf16 (quantized here) or already fp8 matching fmt
-                    x.scalar_type() == torch::kBFloat16 &&
+    // (reused from the forward — the fp8 cached_cast analog, symmetric with
-                    w.scalar_type() == torch::kBFloat16,
+    // the forward's pre-quantized w path). Pre-quantized operands skip their
-                "g, x, and w must be bf16");
+    // quantize kernel; their scale is still passed for the GEMM dequant.
+    const auto f8opt = fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn;
+    const bool g_prequant = g.scalar_type() == f8opt;
+    const bool x_prequant = x.scalar_type() == f8opt;
+    const bool w_prequant = w.scalar_type() == f8opt;
+    TORCH_CHECK(
+        (g.scalar_type() == torch::kBFloat16 || g_prequant) &&
+            (x.scalar_type() == torch::kBFloat16 || x_prequant) &&
+            (w.scalar_type() == torch::kBFloat16 || w_prequant),
+        "g, x, and w must be bf16 or pre-quantized fp8 matching fmt");
     TORCH_CHECK(g.device() == x.device() && g.device() == w.device(),
                 "g, x, and w must be on the same device");
     TORCH_CHECK(masks.size() == 3, "masks must contain three values");
@@ -375,8 +384,10 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     TORCH_CHECK(x_c.size(0) == m && x_c.size(1) == k && g_c.size(1) == n,
                 "backward shape mismatch");
 
-    auto grad_input = torch::empty_like(x);
+    auto grad_input =
-    auto grad_weight = torch::empty_like(w);
+        torch::empty_like(x, x.options().dtype(torch::kBFloat16));
+    auto grad_weight =
+        torch::empty_like(w, w.options().dtype(torch::kBFloat16));
     auto grad_bias = torch::empty({0}, g.options());
     // amax_g feeds a cross-block atomic_max in the g quantize kernel; zero it
     // on the stream (empty + memset, not torch::zeros — see the forward).
@@ -385,8 +396,6 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     // is harmless.
     auto amax_g = torch::empty({1}, g.options().dtype(torch::kFloat32));
     cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), stream.stream());
-    auto f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2)
-                     : g.options().dtype(torch::kFloat8_e4m3fn);
 
     auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst,
                         const torch::Tensor& scale, torch::Tensor* amax,
@@ -420,14 +429,23 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
 
     torch::Tensor g8;
     if (masks[0] || masks[1]) {
-        g8 = torch::empty({m, n}, f8opt);
+        if (g_prequant) {
-        quantize(g_c, g8, sg, &amax_g, g_ring, g_ring_idx, g_ring_margin);
+            g8 = g_c;
+        } else {
+            g8 = torch::empty({m, n}, g.options().dtype(f8opt));
+            quantize(g_c, g8, sg, &amax_g, g_ring, g_ring_idx, g_ring_margin);
+        }
     }
     // dX = g @ w: A = g8 [M,N] (contract over N), B = w8 [N,K] read transposed
     // (b[p*b_ld + n] = w[p,n]); out = [M,K], a_ld = N, b_ld = K, contract = N.
     if (masks[0]) {
-        auto w8 = torch::empty({n, k}, f8opt);
+        torch::Tensor w8;
-        quantize(w_c, w8, sw, nullptr, c10::nullopt, 0, 0);
+        if (w_prequant) {
+            w8 = w_c;
+        } else {
+            w8 = torch::empty({n, k}, g.options().dtype(f8opt));
+            quantize(w_c, w8, sw, nullptr, c10::nullopt, 0, 0);
+        }
         auto grad_input_2d = grad_input.reshape({m, k});
         FP8Params gp;
         pack_gemm_params(gp, g8.data_ptr(), w8.data_ptr(),
@@ -439,19 +457,31 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     // x8 [M,K] read transposed (b[p*b_ld + n] = x[p,n]); out = [N,K], a_ld = N,
     // b_ld = K, contract = M.
     if (masks[1]) {
-        auto x8 = torch::empty({m, k}, f8opt);
+        torch::Tensor x8;
-        quantize(x_c, x8, sx, nullptr, c10::nullopt, 0, 0);
+        if (x_prequant) {
+            x8 = x_c;
+        } else {
+            x8 = torch::empty({m, k}, g.options().dtype(f8opt));
+            quantize(x_c, x8, sx, nullptr, c10::nullopt, 0, 0);
+        }
         FP8Params gp;
         pack_gemm_params(gp, g8.data_ptr(), x8.data_ptr(),
                          grad_weight.data_ptr(), sg, sx, nullptr, nullptr,
                          nullptr, n, k, m, n, k);
         run_bwd_gemm(gp, true, false);
     }
-    if (!masks[0] && !masks[1]) {
+    if (!masks[0] && !masks[1] && !g_prequant) {
         amax_g.copy_(g_c.abs().amax().to(torch::kFloat32));
     }
     C10_CUDA_CHECK(cudaGetLastError());
-    if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type());
+    if (masks[2]) {
+        // A pre-quantized g has no bf16 source to reduce; dequantize with its
+        // scale before the batch-sum so grad_bias stays in the true gradient
+        // domain (sg * sum(g8)).
+        grad_bias = g_prequant
+                        ? (g_c.to(torch::kFloat32) * sg).sum(0).to(torch::kBFloat16)
+                        : g_c.sum(0).to(g.scalar_type());
+    }
     return {grad_input, grad_weight, grad_bias, amax_g};
 }
 

tests/extension/test_fp8_mma.py

@@ -136,6 +136,45 @@ def test_quantize_ring_in_kernel_finalize():
         idx = (idx + 1) % 4
 
 
+@skip_no_fp8
+def test_delayed_scaling_forward_uses_snapshot_scale():
+    """Regression: the in-kernel ring finalize overwrites the scale slot, which
+    aliases the scale the GEMM must dequantize with. The forward must snapshot
+    the delayed scale first, so a changing amax across steps does not leak the
+    next-step scale into the output (otherwise out is off by
+    scale_next / scale_current)."""
+    torch.manual_seed(11)
+    dev = torch.device("cuda")
+    state = f8mod.fp8_state()
+    state.reset()
+    state.default_recipe = DelayedScaling(history_len=1, margin=0)
+    state.default_format = FP8Format.E4M3
+    try:
+        m, n, k = 32, 16, 64
+        x1 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.5
+        # Smaller amax than x1: the delayed scale (amax(x1)/448) still covers
+        # x2 without fp8 saturation, while the next-step scale would differ —
+        # exactly the condition that exposed the overwrite bug.
+        x2 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.35
+        w = torch.randn(n, k, device=dev, dtype=torch.bfloat16) * 0.5
+        bias = torch.zeros(n, device=dev, dtype=torch.bfloat16)
+
+        f8mod.fp8_linear_forward(x1, w, bias)  # step 1: seeds the rings
+        out2, _, _ = f8mod.fp8_linear_forward(x2, w, bias)  # amax changes
+        torch.cuda.synchronize()
+
+        # The delayed scale for step 2 is amax(x1)/448 (history_len=1); the
+        # GEMM must use that same scale for dequant as the quantize used.
+        sx = _scale(x1)
+        sw = _scale(w)
+        qx = _quantize(x2, sx)
+        qw = _quantize(w, sw)
+        expected = (qx @ qw.t() * sx * sw + bias).to(torch.bfloat16)
+        torch.testing.assert_close(out2, expected, atol=0.125, rtol=0.01)
+    finally:
+        state.reset()
+
+
 @skip_no_fp8
 def test_fp8_linear_forward_and_backward():
     torch.manual_seed(7)