refactor: split quantize into multi-type primitive

- split quantize into quantize.cuh, templated on input type (bf16/fp16/fp32)
- rename pybind entry quantize_bf16 to quantize; validate the fmt enum
- fix fp8x2 packing: one 32-bit word packs two pairs (halves were dropped)
- drop the dead OutFp8 template param; GEMM output is always bf16
- fp8_state.reset() restores recipe/format defaults too (test state leak)
- rewrite tests for the two-primitive API with fp32-domain amax references

astrai/extension/fp8.py

@@ -36,10 +36,7 @@ from typing import Dict, List, Optional
 import torch
 from torch.library import Library
 
-from astrai.extension.ops.fp8 import (
+from astrai.extension.ops.fp8 import mm_fp8, quantize
-    linear_backward_fp8,
-    linear_forward_fp8,
-)
 
 # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
 FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -95,11 +92,10 @@ class DynamicScaling(FP8Recipe):
 
 class _ScaleRing:
     """One operand's delayed-scaling state: a float32 buffer
-    ``[hist[n] | scale | counter]`` (views). The quantize kernel's last-finishing
+    ``[hist[n] | scale | counter]`` (views). ``update`` folds the amax
-    block records the measured amax into ``hist[idx]``, reduces the window and
+    returned by the quantize primitive into ``hist[idx]`` and publishes the
-    publishes the next scale entirely on device — the Python-side write/max/write
+    next scale from the window; ``idx`` advances host-side each step. The
-    chain is gone. The counter slot stays int32-zero (float bits) between
+    trailing slot is a legacy counter kept for state-buffer compatibility.
-    launches; ``idx`` advances host-side each step.
     """
 
     __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -114,7 +110,7 @@ class _ScaleRing:
         self.initialized = False
 
     def advance(self) -> None:
-        """Rotate to the next history slot after an in-kernel finalize."""
+        """Rotate to the next history slot after metadata update."""
         self.idx = (self.idx + 1) % self.hist.numel()
 
     def seed(self, t: torch.Tensor, fmt: str) -> None:
@@ -123,12 +119,15 @@ class _ScaleRing:
         self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
         self.initialized = True
 
+    def update(self, amax: torch.Tensor, fmt: str) -> None:
+        self.hist[self.idx].copy_(amax.reshape(()))
+        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
+
 
 class FP8TensorMeta:
-    """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/
+    """Per-weight delayed-scaling state for ``w``, ``x`` and ``g``.
-    ``g``). Fused kernels record amax while quantizing, so the scale used at step
+
-    N reflects amax from steps < N. DynamicScaling never allocates a meta — it
+    DynamicScaling never allocates a meta; it measures the current amax inline.
-    measures the current amax inline.
     """
 
     __slots__ = ("w", "x", "g")
@@ -213,7 +212,11 @@ class FP8State:
         return meta
 
     def reset(self) -> None:
+        """Restore construction defaults (switch, recipe, format) and drop all
+        per-weight metas — a full state reset for tests / reconfiguration."""
         self.default_enabled = False
+        self.default_recipe = DelayedScaling()
+        self.default_format = FP8Format.HYBRID
         self._metas.clear()
 
 
@@ -307,6 +310,11 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor
     return recipe.scale_from_history(amax, fmt)
 
 
+def _is_fp8(dtype: torch.dtype) -> bool:
+    """A pre-quantized weight takes the GEMM directly (no re-quantize)."""
+    return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+
+
 _zero_bias: Dict[Optional[int], torch.Tensor] = {}
 
 
@@ -326,9 +334,9 @@ def fp8_linear_forward(
 ):
     """Scaled fp8 linear forward (called from the aten::linear impl).
 
-    Pure FP8 path for both recipes: quantize x/w with the active scales, run the
+    Composed from the two stateless primitives: quantize x/w with the active
-    pre-quantized GEMM. Delayed scaling finalizes the rings inside the quantize
+    scales, run the pre-quantized GEMM, add the bias. Delayed scaling folds
-    kernels (amax folded into the window, next scale published on device);
+    the returned amax into the history ring and publishes the next scale;
     dynamic scaling measures the current amax itself. Training quantizes the
     weight every step (the optimizer bumps its version, so there is no cast
     cache, matching ``cached_cast``-less behavior).
@@ -337,49 +345,39 @@ def fp8_linear_forward(
     if cfg is None:
         cfg = _current_config()
     fmt = cfg.fp8_format.fwd()
-    margin = cfg.recipe.margin
     if bias is None:
         bias = _empty_bias(x)
-    if isinstance(cfg.recipe, DynamicScaling):  # measure-then-quantize, no state
+    if isinstance(cfg.recipe, DynamicScaling):
         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)
+        x8, _ = quantize(x, sx.reciprocal(), fmt)
-        return out, sx, sw
+        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
+        out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape(
+            *x.shape[:-1], w.size(0)
+        )
+        return (out + bias if bias.numel() else 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 quantize kernels finalize each ring in-kernel and overwrite the ring's
+    sx, sw = meta.x.scale.clone(), meta.w.scale.clone()
-    # scale slot, which ALIASES meta.*.scale (a view into the state buffer).
+    x8, amax_x = quantize(x, sx.reciprocal(), fmt)
-    # Snapshot the scales first so the GEMM dequantizes with the SAME scale the
+    if _is_fp8(w.dtype):
-    # operands were quantized with, and so the backward can reuse this step's
+        w8, amax_w = w, None
-    # 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
+        w8, amax_w = quantize(w, sw.reciprocal(), fmt)
-    sx = meta.x.scale.clone()
+    out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape(
-    sw = sw_arg.clone()
+        *x.shape[:-1], w.size(0)
-    out, _x8, _w8, _ax, _aw = linear_forward_fp8(
-        x,
-        w_arg,
-        bias,
-        sx,
-        sw,
-        fmt,
-        None,
-        meta.x.state,
-        meta.x.idx,
-        margin,
-        w_ring,
-        meta.w.idx,
-        margin,
     )
+    if bias.numel():
+        out = out + bias
+    meta.x.update(amax_x, fmt)
+    if amax_w is not None:
+        meta.w.update(amax_w, fmt)
     meta.x.advance()
-    if w_ring is not None:
+    if amax_w is not None:
         meta.w.advance()
     return out, sx, sw
 
@@ -410,37 +408,25 @@ class _LinearFp8(torch.autograd.Function):
     def backward(ctx, g):
         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:
             sg = _dynamic_scale(g, ctx.recipe, fmt)
             sw = _dynamic_scale(w, ctx.recipe, fmt)
             sx = _dynamic_scale(x, ctx.recipe, fmt)
-            ring, idx = None, 0
         else:
             meta = ctx.meta
             if not meta.g.initialized:
                 meta.g.seed(g, fmt)
-            # Snapshot the g scale before its ring finalize overwrites the slot
-            # (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(
+        g8, amax_g = quantize(g, sg.reciprocal(), fmt)
-            g,
+        x8, _ = quantize(x, sx.reciprocal(), fmt)
-            x,
+        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
-            w,
+        grad_x = mm_fp8(g8, w8, sg * sw)  # g8[m,n] @ w8[n,k] natural
-            list(ctx.needs_input_grad),
+        grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True)  # g8.T @ x8
-            sg,
+        grad_b = g.sum(0).to(torch.bfloat16)
-            sw,
-            sx,
-            fmt,
-            ring,
-            idx,
-            ctx.recipe.margin,
-        )
         if not ctx.is_dynamic:
-            meta.g.advance()  # the g quantize kernel finalized the ring in-kernel
+            meta.g.update(amax_g, fmt)
+            meta.g.advance()
         return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
 
 

astrai/extension/ops/fp8.py

@@ -2,10 +2,8 @@
 
 Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
 
-- ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → FP8 with fused amax
+- ``quantize(x, scale, fmt) -> (x8, amax)`` — BF16/FP16/FP32 → FP8 with fused amax
 - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
-- ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, x8, w8, amax_x, amax_w)``
-- ``linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt) -> (gx, gw, gb, amax_g)``
 
 Scale semantics: scales are *quantization steps* — the value divided out when
 quantizing (``x8 = x / scale``). Every primitive computes its own inverse
@@ -16,7 +14,7 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
 this module is stateless.
 """
 
-from typing import List, Optional, Tuple
+from typing import Tuple
 
 import torch
 from torch.library import custom_op
@@ -34,6 +32,14 @@ def _fmt_int(fmt: str) -> int:
         raise ValueError(f"unsupported fp8 format {fmt!r} (expected 'e4m3' or 'e5m2')")
 
 
+def _fmt_name(fmt: int) -> str:
+    if fmt == 0:
+        return "e4m3"
+    if fmt == 1:
+        return "e5m2"
+    raise ValueError(f"unsupported quantization type {fmt!r}")
+
+
 def _fmt_dtype(fmt: str) -> torch.dtype:
     return torch.float8_e5m2 if _fmt_int(fmt) else torch.float8_e4m3fn
 
@@ -42,28 +48,31 @@ def _fmt_dtype(fmt: str) -> torch.dtype:
 def fp8_quantize(
     x: torch.Tensor, scale: torch.Tensor, fmt: int
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``."""
+    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; ``scale`` is a multiplier."""
 
 
 @fp8_quantize.register_fake
 def _fp8_quantize_fake(x, scale, fmt):
-    dtype = torch.float8_e5m2 if fmt else torch.float8_e4m3fn
+    dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
     return (
         torch.empty(x.shape, device=x.device, dtype=dtype),
         torch.empty(1, device=x.device, dtype=torch.float32),
     )
 
 
+_QUANT_INPUT_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
+
+
 @fp8_quantize.register_kernel("cuda")
 def _fp8_quantize_cuda(x, scale, fmt):
-    if x.dtype != torch.bfloat16:
+    if x.dtype not in _QUANT_INPUT_DTYPES:
-        raise TypeError(f"fp8 quantize requires bf16 input, got {x.dtype}")
+        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
-    return get_module("fp8_ops").quantize_bf16(x, scale, int(fmt))
+    return get_module("fp8_ops").quantize(x, scale, int(fmt))
 
 
 @fp8_quantize.register_kernel("cpu")
 def _fp8_quantize_cpu(x, scale, fmt):
-    x8 = (x.float() / scale).to(_fmt_dtype("e5m2" if fmt else "e4m3"))
+    x8 = (x.float() * scale).to(_fmt_dtype(_fmt_name(fmt)))
     amax = x.abs().amax().float().reshape(1).clamp_min(1e-12)
     return x8, amax
 
@@ -72,50 +81,51 @@ def _fp8_quantize_cpu(x, scale, fmt):
 def fp8_gemm(
     a: torch.Tensor,
     b: torch.Tensor,
-    sa: torch.Tensor,
+    scale: torch.Tensor,
-    sb: torch.Tensor,
+    trans_a: int = 0,
-    out_dtype: int = 0,
+    trans_b: int = 0,
-    out_scale: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation.
+    """FP8 GEMM: ``a @ b * scale`` with FP32 accumulation.
 
-    ``out_dtype``: 0 = BF16 (default), 1 = FP8 E4M3 (requires ``out_scale``,
+    The result is always BF16; FP8 output is a separate quantize operation.
-    the quantization step for the output — mirrors ``torch._scaled_mm``).
     """
 
 
 @fp8_gemm.register_fake
-def _fp8_gemm_fake(a, b, sa, sb, out_dtype=0, out_scale=None):
+def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0):
-    dtype = torch.float8_e4m3fn if out_dtype else torch.bfloat16
+    dtype = torch.bfloat16
-    return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=dtype)
+    return torch.empty(
+        (a.size(1) if trans_a else a.size(0), b.size(0) if trans_b else b.size(1)),
+        device=a.device,
+        dtype=dtype,
+    )
 
 
 @fp8_gemm.register_kernel("cuda")
-def _fp8_gemm_cuda(a, b, sa, sb, out_dtype=0, out_scale=None):
+def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0):
     if a.dtype != b.dtype or a.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
         raise TypeError(
             f"fp8 GEMM requires matching fp8 inputs, got {a.dtype}/{b.dtype}"
         )
-    return get_module("fp8_ops").mm_fp8(a, b, sa, sb, int(out_dtype), out_scale)
+    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b)
 
 
 @fp8_gemm.register_kernel("cpu")
-def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None):
+def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0):
-    acc = a.float() @ b.float() * sa * sb
+    aa = a.float().t() if trans_a else a.float()
-    if out_dtype:
+    bb = b.float().t() if trans_b else b.float()
-        os_ = 1.0 if out_scale is None else out_scale
+    acc = aa @ bb * scale
-        return (acc * os_).to(torch.float8_e4m3fn)
     return acc.to(torch.bfloat16)
 
 
-def quantize_bf16(
+def quantize(
     x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``.
+    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; returns
+    ``(x8, amax)``.
 
-    ``scale`` is the quantization step (device scalar); ``fmt`` selects
+    ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
-    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor — the caller
+    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
-    never clears it.
     """
     return fp8_quantize(x, scale, _fmt_int(fmt))
 
@@ -123,123 +133,14 @@ def quantize_bf16(
 def mm_fp8(
     a: torch.Tensor,
     b: torch.Tensor,
-    sa: torch.Tensor,
+    scale: torch.Tensor,
-    sb: torch.Tensor,
+    trans_a: bool = False,
-    out_dtype: str = "bf16",
+    trans_b: bool = False,
-    out_scale: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``.
+    """Pre-quantized FP8 GEMM: ``a @ b * scale``.
 
-    ``a``/``b`` must be FP8 tensors of the same format (E4M3 or E5M2);
+    ``a``/``b`` must be FP8 tensors of the same format. ``scale`` is their
-    ``sa``/``sb`` are their quantization steps. ``out_dtype`` is ``"bf16"``
+    combined dequantization scale. The result is BF16; FP8 output is a separate
-    (default) or ``"e4m3"`` — FP8 output for layer-to-layer pipelines, which
+    quantize operation.
-    requires ``out_scale`` (the output quantization step).
     """
-    if out_dtype not in ("bf16", "e4m3"):
+    return fp8_gemm(a, b, scale, trans_a, trans_b)
-        raise ValueError(
-            f"unsupported out_dtype {out_dtype!r} (expected 'bf16' or 'e4m3')"
-        )
-    return fp8_gemm(a, b, sa, sb, int(out_dtype == "e4m3"), out_scale)
-
-
-def linear_forward_fp8(
-    x: torch.Tensor,
-    w: torch.Tensor,
-    bias: Optional[torch.Tensor],
-    sx: torch.Tensor,
-    sw: torch.Tensor,
-    fmt: str = "e4m3",
-    bias_scale: Optional[torch.Tensor] = None,
-    x_ring: Optional[torch.Tensor] = None,
-    x_ring_idx: int = 0,
-    x_ring_margin: int = 0,
-    w_ring: Optional[torch.Tensor] = None,
-    w_ring_idx: int = 0,
-    w_ring_margin: int = 0,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Pure FP8 linear forward: quantize x/w to ``fmt``, pre-quantized GEMM.
-
-    Returns ``(out, x8, w8, amax_x, amax_w)`` — the quantized operands are
-    handed back so the policy layer can cache the weight quantization while
-    the weight tensor is unchanged (torch autocast's cached_cast analog).
-    ``x8`` is ``[M, K]`` and ``w8`` is ``[N, K]`` (the passed-in ``w`` itself
-    on the pre-quantized path). ``bias`` may be ``None``. For static fp8
-    inference, ``w`` and ``bias`` may arrive pre-quantized to ``fmt``
-    (produced by :func:`quantize_bf16` with their scales as ``sw`` /
-    ``bias_scale``); a pre-quantized ``bias`` requires ``bias_scale``, and
-    its ``amax_w`` comes back 0. The bias is fused into the GEMM epilogue.
-    ``x_ring`` / ``w_ring`` (delayed scaling) are ``[hist | scale | counter]``
-    float32 buffers the quantize kernels finalize in-kernel: the measured
-    amax lands in ``hist[idx]`` and the next step's scale is published on
-    device, replacing the eager hist/max/scale update chain.
-    """
-    fmt8 = _fmt_dtype(fmt)
-    if x.dtype != torch.bfloat16 or w.dtype not in (torch.bfloat16, fmt8):
-        raise TypeError(
-            f"fp8 forward requires bf16 x and bf16-or-{fmt} w, got {x.dtype}/{w.dtype}"
-        )
-    if bias is None:
-        bias = torch.empty(0, device=x.device, dtype=x.dtype)
-    return get_module("fp8_ops").linear_forward_fp8(
-        x,
-        w,
-        bias,
-        sx,
-        sw,
-        _fmt_int(fmt),
-        bias_scale,
-        x_ring,
-        x_ring_idx,
-        x_ring_margin,
-        w_ring,
-        w_ring_idx,
-        w_ring_margin,
-    )
-
-
-def linear_backward_fp8(
-    g: torch.Tensor,
-    x: torch.Tensor,
-    w: torch.Tensor,
-    masks: List[bool],
-    sg: torch.Tensor,
-    sw: torch.Tensor,
-    sx: torch.Tensor,
-    fmt: str = "e5m2",
-    g_ring: Optional[torch.Tensor] = None,
-    g_ring_idx: int = 0,
-    g_ring_margin: int = 0,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``.
-
-    ``g``/``x``/``w`` may each be bf16 (quantized to ``fmt`` here) or already
-    pre-quantized fp8 matching ``fmt`` — a pre-quantized operand skips its
-    quantize kernel and is read directly by the GEMM (the ``cached_cast``
-    analog for the backward, symmetric with :func:`linear_forward_fp8`'s
-    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``.
-    """
-    f8 = _fmt_dtype(fmt)
-    for name, t in (("g", g), ("x", x), ("w", w)):
-        if t.dtype not in (torch.bfloat16, f8):
-            raise TypeError(
-                f"fp8 backward requires bf16 or pre-quantized {fmt} inputs, "
-                f"got {name}={t.dtype}"
-            )
-    return get_module("fp8_ops").linear_backward_fp8(
-        g,
-        x,
-        w,
-        list(masks),
-        sg,
-        sw,
-        sx,
-        _fmt_int(fmt),
-        g_ring,
-        g_ring_idx,
-        g_ring_margin,
-    )

csrc/kernels/fp8/common.h

@@ -69,32 +69,17 @@ struct Fp8GemmTraits {
     static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
 };
 
-// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional
+// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
-// delayed-scaling ring finalization. Separate from FP8Params so each
+// with fused amax.
-// operator owns exactly the fields it touches (the GEMM never reads amax /
-// ring state). Same NSDMI rationale: amax / ring_state gate optional paths
-// via null checks. Still an aggregate, still trivially copyable.
 struct FP8QuantizeParams {
-    // BF16 input and FP8 output buffers; scale_a is the quantization step
+    // Float input and FP8 output buffers; scale is the quantization
-    // (device scalar). amax_a (may be null) is zero-initialized by the
+    // multiplier (device scalar). amax (may be null) is zero-initialized by
-    // binding and receives the raw-domain absolute maximum.
+    // the binding and receives the raw-domain absolute maximum.
-    const void* __restrict__ a_ptr = nullptr;
+    const void* __restrict__ input_ptr = nullptr;
-    void* __restrict__ out_ptr = nullptr;
+    void* __restrict__ output_ptr = nullptr;
-    const float* __restrict__ scale_a = nullptr;
-    float* __restrict__ amax_a = nullptr;
 
-    // Optional delayed-scaling ring finalization. ring_state packs
+    const float* __restrict__ scale = nullptr;
-    // [hist[ring_len] | scale | counter] with ring_len = numel - 2. When
+    float* __restrict__ amax = nullptr;
-    // non-null and amax_a is set, the last-finishing block records the
-    // measured amax into hist[ring_idx], reduces the window and publishes
-    // the next step's scale (max(hist) / fp8_max / 2^ring_margin) — the
-    // fused replacement for the eager hist-write / max / scale-write chain,
-    // at zero extra launches. The counter slot is a persistent zero-armed
-    // int32 (float bits) electing the last block each launch.
-    float* ring_state = nullptr;
-    int ring_len = 0;
-    int ring_idx = 0;
-    int ring_margin = 0;
 
     // Element count (only the elementwise quantize kernel uses it).
     int total = 0;
@@ -103,24 +88,15 @@ struct FP8QuantizeParams {
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
 // through the pre-quantized GEMM kernels. Each kernel touches only the
 // fields it needs; buffers are raw pointers packed by the torch binding.
-// Pointer members default to null (same NSDMI rationale as AttentionParams:
+// Pointer members default to null so optional paths cannot hold garbage.
-// bias / out_scale gate optional paths via null checks, so a partially
-// packed struct must never hold garbage non-null pointers). Still an
-// aggregate, still trivially copyable.
 struct FP8Params {
     // Inputs: a/b are FP8 for the pre-quantized path. Scales are
     // quantization steps (device scalars).
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
-    const void* __restrict__ bias = nullptr;
-    const float* __restrict__ scale_a = nullptr;
-    const float* __restrict__ scale_b = nullptr;
-    const float* __restrict__ bias_scale = nullptr;
-    // Output: BF16 or FP8 (E4M3). out_scale is the output quantization step
-    // (FP8 output only).
     void* __restrict__ out_ptr = nullptr;
-    const float* __restrict__ out_scale = nullptr;
 
+    const float* __restrict__ scale = nullptr;
     // Shapes. `int` covers every realistic LLM shape; the kernels promote
     // to int64 for all pointer arithmetic.
     int m, n, k;

csrc/kernels/fp8/gemm.cuh

@@ -2,7 +2,8 @@
 // FP8 GEMM device code — pure CUDA, no torch. Mirrors the attention kernel
 // layout (attn_*_mma.cuh): kernels take the FP8Params POD, tile shape and
 // FP8 format ride on compile-time template parameters, and launchers are
-// plain functions usable from both the torch binding and pure C tests.
+// plain functions usable from both the torch binding and pure C tests. The
+// quantize kernel lives in quantize.cuh.
 
 #include <cuda_bf16.h>
 #include <cuda_fp8.h>
@@ -46,128 +47,11 @@ struct fp8_input<FP8Format::E5M2> {
 // FP8 MMA lives in the shared astrai::mma_sync template (common/mma.cuh);
 // instantiate it with fp8_input<Fmt>::type. Accumulates in-place: callers
 // pass the same accumulator array as both `d` and `c`.
-// warp_reduce_max / atomic_max_float (quantize amax) live in
+// warp_reduce_sum / group_reduce_sum (GEMM) live in common/reduce.cuh; the
-// common/reduce.cuh; the cp.async pipeline primitives (predicated 16-byte
+// cp.async pipeline primitives (predicated 16-byte copy, commit_group,
-// copy, commit_group, wait_group + runtime dispatch) in common/cp_async.cuh.
+// wait_group + runtime dispatch) in common/cp_async.cuh.
 
 // ---------------------------------------------------------------------------
-// Quantize kernel: BF16 -> FP8 (E4M3 or E5M2), fused amax over raw values.
-// ---------------------------------------------------------------------------
-
-// Convert one packed bf16 pair to one packed fp8 pair. amax sees the *raw*
-// (unscaled) values; the stored bytes see value * inv. Bit-identical to the
-// scalar __nv_fp8_*(q) constructor path (round-nearest-even + satfinite).
-template <FP8Format Fmt>
-__device__ __forceinline__ unsigned quantize2(unsigned pair, float inv,
-                                              float& amax) {
-    const float lo = __bfloat162float(__ushort_as_bfloat16(pair & 0xffffu));
-    const float hi = __bfloat162float(__ushort_as_bfloat16(pair >> 16));
-    amax = fmaxf(amax, fmaxf(fabsf(lo), fabsf(hi)));
-    constexpr __nv_fp8_interpretation_t kFmt =
-        Fmt == FP8Format::E5M2 ? __NV_E5M2 : __NV_E4M3;
-    return static_cast<unsigned>(__nv_cvt_float2_to_fp8x2(
-        make_float2(lo * inv, hi * inv), __NV_SATFINITE, kFmt));
-}
-
-template <FP8Format Fmt>
-__global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
-    const float inv = 1.0f / *p.scale_a;
-    const auto* x = reinterpret_cast<const __nv_bfloat16*>(p.a_ptr);
-    void* x8 = p.out_ptr;
-    float* amax = p.amax_a;
-    float local_amax = 0.0f;
-    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
-
-    // Vectorized body: 8 bf16 (16B load) -> 8 fp8 (8B store) per step. Torch
-    // allocations are >=16B aligned and the binding passes freshly allocated
-    // contiguous buffers, so element 0 keeps the uint4/uint2 accesses
-    // natural; a misaligned base (contiguous view with an odd storage
-    // offset) falls back to the scalar loop below via total_vec = 0.
-    const bool aligned =
-        ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(x8)) & 15) ==
-        0;
-    const int64_t total_vec = aligned ? p.total / 8 : 0;
-    const uint4* xv = reinterpret_cast<const uint4*>(x);
-    uint2* o8 = reinterpret_cast<uint2*>(x8);
-    for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total_vec;
-         i += stride) {
-        const uint4 v = xv[i];
-        const unsigned pair[4] = {v.x, v.y, v.z, v.w};
-        unsigned packed[2] = {0u, 0u};
-#pragma unroll
-        for (int j = 0; j < 4; ++j)
-            packed[j >> 1] |= quantize2<Fmt>(pair[j], inv, local_amax)
-                              << (16 * (j & 1));
-        o8[i] = make_uint2(packed[0], packed[1]);
-    }
-    // Scalar tail (and full fallback for misaligned bases).
-    for (int64_t i = total_vec * 8 + blockIdx.x * blockDim.x + threadIdx.x;
-         i < p.total; i += stride) {
-        const float f = __bfloat162float(x[i]);
-        local_amax = fmaxf(local_amax, fabsf(f));
-        if constexpr (Fmt == FP8Format::E5M2) {
-            reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] = __nv_fp8_e5m2(f * inv);
-        } else {
-            reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] = __nv_fp8_e4m3(f * inv);
-        }
-    }
-    if (amax) {
-        local_amax = warp_reduce_max(local_amax);
-        __shared__ float slots[32];
-        if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax;
-        __syncthreads();
-        if (threadIdx.x == 0) {
-            float v = 0.0f;
-            for (int w = 0; w < (blockDim.x >> 5); ++w) v = fmaxf(v, slots[w]);
-            atomic_max_float(amax, v);
-        }
-    }
-    if (p.ring_state && amax) {
-        // Delayed-scaling ring finalization as a last-block epilogue (the
-        // CUDA threadFenceReduction pattern): the fence + counter elect the
-        // final block once every block's atomic_max above is visible; warp 0
-        // folds the fresh amax into the window, reduces it and publishes the
-        // next step's scale, then re-arms the counter for the next launch.
-        // __fdiv_rn / ldexpf keep the scale bit-identical to the eager
-        // (peak / fp8_max) / 2^margin fp32 chain despite --use_fast_math.
-        __threadfence();
-        __shared__ bool ring_last;
-        if (threadIdx.x == 0)
-            ring_last = atomicAdd(reinterpret_cast<int*>(p.ring_state +
-                                                        p.ring_len + 1),
-                                  1) == gridDim.x - 1;
-        __syncthreads();
-        if (ring_last && threadIdx.x < 32) {
-            float* hist = p.ring_state;
-            const int lane = threadIdx.x;
-            float v = 0.0f;
-            if (lane < p.ring_len) v = hist[lane];
-            if (lane == p.ring_idx) {
-                v = *amax;  // the global amax is final now
-                hist[lane] = v;
-            }
-            // Windows longer than one warp (atypical) fold the tail.
-            for (int i = lane + 32; i < p.ring_len; i += 32) {
-                float h = hist[i];
-                if (i == p.ring_idx) {
-                    h = *amax;
-                    hist[i] = h;
-                }
-                v = fmaxf(v, h);
-            }
-            const float peak = warp_reduce_max(v);
-            if (lane == 0) {
-                constexpr float kFmtMax =
-                    Fmt == FP8Format::E5M2 ? 57344.0f : 448.0f;
-                p.ring_state[p.ring_len] = fmaxf(
-                    ldexpf(__fdiv_rn(peak, kFmtMax), -p.ring_margin), 1e-12f);
-                __threadfence();
-                // Re-arm the counter (0.0f bits == int32 0).
-                p.ring_state[p.ring_len + 1] = 0.0f;
-            }
-        }
-    }
-}
 
 // Swizzled address inside a flat [rows * K] staging tile: the 16-byte chunk
 // index is XORed with a row-dependent slice so a warp's fragment load (8
@@ -446,8 +330,7 @@ struct Fp8GemmSmem {
 // (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and
 // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the
 // launcher dispatches to it there (see launch_fp8_gemm).
-template <typename Traits, bool OutFp8 = false, 
+template <typename Traits, typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
-        typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
         bool kBStaged = true>
 __global__ void __launch_bounds__(Traits::kCtaThreads,
                                   Fp8GemmSmem<Traits, LayoutA, LayoutB,
@@ -495,7 +378,6 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
     const auto* a = reinterpret_cast<const T8*>(p.a_ptr);
     const auto* b = reinterpret_cast<const T8*>(p.b_ptr);
     auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr);
-    auto* out_fp8 = reinterpret_cast<__nv_fp8_e4m3*>(p.out_ptr);
     const int64_t m = p.m, n = p.n, k = p.k;
     const int64_t a_ld = p.a_ld, b_ld = p.b_ld;
 
@@ -537,8 +419,7 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
         (int64_t)block_n * kBlockN + warp_n * 32 + thread_in_group * 2;
     const int a_row0 = warp_m * 64;  // + mt * 16 in the loop
     const int b_row0 = warp_n * 32;  // + nt * 8
-    const float sa = *p.scale_a;
+    const float scale = *p.scale;
-    const float sb = *p.scale_b;
     float acc[4][4][4] = {};  // [nt][mt][acc]
 
     // Both operands end up in the canonical [M][kK] / [N][kK] shared tiles
@@ -742,49 +623,25 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
         }
     }
 
-    const float output_scale = sa * sb;
+    const float output_scale = scale;
-    // Fused bias: BF16 raw values, or FP8 storage dequantized by its own
-    // scale (bias_scale != null selects the FP8 path; the format follows the
-    // kernel's Traits). Added in real units after the operand dequantization
-    // and before any output quantization.
-    const auto* bias16 = static_cast<const __nv_bfloat16*>(p.bias);
-    const auto* bias8 = static_cast<const T8*>(p.bias);
-    auto bias_val = [&](int64_t col) -> float {
-        if (p.bias == nullptr || col >= n) return 0.0f;
-        if (p.bias_scale == nullptr) return __bfloat162float(bias16[col]);
-        return __half2float(__half(bias8[col])) * *p.bias_scale;
-    };
 #pragma unroll
     for (int nt = 0; nt < 4; ++nt) {
         const int64_t col = output_col + nt * 8;
-        const float b0 = bias_val(col);
-        const float b1 = bias_val(col + 1);
         // Per-row store: FP8 packs two adjacent columns into one 16-bit
         // write, BF16 into one 32-bit __nv_bfloat162 (single cvt+pack
         // instruction); boundary or unaligned columns fall back to scalar
         // converts so a pack never crosses the row edge or misaligns.
         auto store_out = [&](int64_t row, float v0, float v1) {
             if (row >= m) return;
-            const float r0 = v0 * output_scale + b0;
+            const float r0 = v0 * output_scale;
-            const float r1 = v1 * output_scale + b1;
+            const float r1 = v1 * output_scale;
-            if constexpr (OutFp8) {
+            auto* dst = out_bf16 + row * n + col;
-                if (col + 1 < n) {
+            if (col + 1 < n && (reinterpret_cast<uintptr_t>(dst) & 3) == 0) {
-                    *reinterpret_cast<unsigned short*>(out_fp8 + row * n + col) =
+                *reinterpret_cast<__nv_bfloat162*>(dst) =
-                        static_cast<unsigned short>(__nv_cvt_float2_to_fp8x2(
+                    __floats2bfloat162_rn(r0, r1);
-                            make_float2(r0 * *p.out_scale, r1 * *p.out_scale),
-                            __NV_SATFINITE, __NV_E4M3));
-                } else {
-                    out_fp8[row * n + col] = __nv_fp8_e4m3(r0 * *p.out_scale);
-                }
             } else {
-                auto* dst = out_bf16 + row * n + col;
+                dst[0] = __float2bfloat16(r0);
-                if (col + 1 < n && (reinterpret_cast<uintptr_t>(dst) & 3) == 0) {
+                if (col + 1 < n) dst[1] = __float2bfloat16(r1);
-                    *reinterpret_cast<__nv_bfloat162*>(dst) =
-                        __floats2bfloat162_rn(r0, r1);
-                } else {
-                    dst[0] = __float2bfloat16(r0);
-                    if (col + 1 < n) dst[1] = __float2bfloat16(r1);
-                }
             }
         };
 #pragma unroll
@@ -803,16 +660,6 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
 // Launchers — pure CUDA (no torch), usable from the binding and pure C tests.
 // ---------------------------------------------------------------------------
 
-template <FP8Format Fmt>
-void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
-    constexpr int kThreads = 256;
-    // One block per 256 vectors (8 elements each); at least one block so the
-    // scalar tail of a tiny / misaligned tensor is still covered.
-    int64_t blocks = (p.total / 8 + kThreads - 1) / kThreads;
-    if (blocks < 1) blocks = 1;
-    fp8_quantize_kernel<Fmt><<<blocks, kThreads, 0, stream>>>(p);
-}
-
 // Launch one kernel instantiation with its shared-memory budget: stages live
 // in dynamic smem, so budgets beyond the 48KB static limit opt in once per
 // instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
@@ -851,7 +698,7 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
 // dX ~39 TF staged vs ~38 direct).
 constexpr int64_t kCrossStageMinK = 8192;
 
-template <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor,
+template <FP8Format Fmt, typename LayoutA = RowMajor,
           typename LayoutB = RowMajor, int kK = 64, int Stages = 2,
           bool GroupRaster = std::is_same_v<LayoutA, ColMajor> || std::is_same_v<LayoutB, ColMajor>>
 void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
@@ -860,24 +707,24 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
     if (p.m <= 64) {
         using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
         if (b_staged)
-            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+            launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
                                              GroupRaster, true>>(
                 Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
                 dim3(Traits::kCtaThreads), stream, p);
         else
-            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+            launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
                                              GroupRaster, false>>(
                 Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
                 dim3(Traits::kCtaThreads), stream, p);
     } else {
         using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
         if (b_staged)
-            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+            launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
                                              GroupRaster, true>>(
                 Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
                 dim3(Traits::kCtaThreads), stream, p);
         else
-            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+            launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
                                              GroupRaster, false>>(
                 Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
                 dim3(Traits::kCtaThreads), stream, p);

csrc/kernels/fp8/ops.cu

@@ -1,76 +1,68 @@
-// FP8 GEMM torch binding: tensor validation, FP8Params packing, template
+// CUDA bindings for the two stateless FP8 primitives.
-// dispatch and pybind. Device code lives in gemm.cuh (pure CUDA) —
-// mirroring the attn_*.cu / attn_*_mma.cuh split of the attention kernels.
 
-#include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
-#include <cuda_bf16.h>
+#include <torch/extension.h>
+
 #include <cstdint>
 #include <mutex>
 #include <tuple>
 #include <unordered_map>
 
-#include "gemm.cuh"
 #include "../common/device.cuh"
+#include "gemm.cuh"
+#include "quantize.cuh"
 
 using namespace astrai::fp8;
 
 namespace {
 
-// FP8Format / FP8Params and the launchers live in astrai::fp8 (common.h /
-// gemm.cuh); this TU opens the using-directive above so the binding reads
-// them unqualified.
-
 void check_fp8_device(const torch::Tensor& tensor) {
     static std::mutex mutex;
     static std::unordered_map<int, bool> supported;
     const int device = tensor.device().index();
     {
         std::lock_guard<std::mutex> lock(mutex);
-        auto cached = supported.find(device);
+        auto it = supported.find(device);
-        if (cached != supported.end()) {
+        if (it != supported.end()) {
-            TORCH_CHECK(cached->second,
+            TORCH_CHECK(it->second, "FP8 MMA requires compute capability 8.9+");
-                        "fused FP8 MMA requires compute capability 8.9 or newer");
             return;
         }
     }
-
     const auto* properties = at::cuda::getDeviceProperties(device);
-    const bool is_supported =
+    const bool ok = astrai::sm_at_least(
-        astrai::sm_at_least(properties->major, properties->minor,
+        properties->major, properties->minor, astrai::kMinSmForFp8Major,
-                            astrai::kMinSmForFp8Major,
+        astrai::kMinSmForFp8Minor);
-                            astrai::kMinSmForFp8Minor);
     {
         std::lock_guard<std::mutex> lock(mutex);
-        supported.emplace(device, is_supported);
+        supported.emplace(device, ok);
     }
-    TORCH_CHECK(is_supported,
+    TORCH_CHECK(ok, "FP8 MMA requires compute capability 8.9+");
-                "fused FP8 MMA requires compute capability 8.9 or newer");
 }
 
-void check_scale(const torch::Tensor& scale, const torch::Tensor& input,
+void check_scale(const torch::Tensor& scale, const torch::Tensor& input) {
-                 const char* name) {
     TORCH_CHECK(scale.is_cuda() && scale.device() == input.device() &&
                     scale.scalar_type() == torch::kFloat32 && scale.numel() == 1,
-                name, " must be a CUDA float32 scalar on the input device");
+                "scale must be a CUDA float32 scalar on the input device");
 }
 
-// ---- FP8Params packing (mirrors attention/entry_utils.cuh pack_* helpers) ----
+void pack_quantize(FP8QuantizeParams& p, const void* input, void* output,
+                   const torch::Tensor& scale, torch::Tensor& amax,
+                   int64_t total) {
+    p.input_ptr = input;
+    p.output_ptr = output;
+    p.scale = scale.data_ptr<float>();
+    p.amax = amax.data_ptr<float>();
+    p.total = static_cast<int>(total);
+}
 
-void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
+void pack_gemm(FP8Params& p, const void* a, const void* b, void* output,
-                      const torch::Tensor& sa, const torch::Tensor& sb,
+               const torch::Tensor& scale, int64_t m, int64_t n, int64_t k,
-                      const torch::Tensor* out_scale, const void* bias,
+               int64_t a_ld, int64_t b_ld) {
-                      const torch::Tensor* bias_scale, int64_t m, int64_t n,
-                      int64_t k, int64_t a_ld, int64_t b_ld) {
     p.a_ptr = a;
     p.b_ptr = b;
-    p.out_ptr = out;
+    p.out_ptr = output;
-    p.scale_a = sa.data_ptr<float>();
+    p.scale = scale.data_ptr<float>();
-    p.scale_b = sb.data_ptr<float>();
-    p.out_scale = out_scale ? out_scale->data_ptr<float>() : nullptr;
-    p.bias = bias;
-    p.bias_scale = bias_scale ? bias_scale->data_ptr<float>() : nullptr;
     p.m = static_cast<int>(m);
     p.n = static_cast<int>(n);
     p.k = static_cast<int>(k);
@@ -78,443 +70,112 @@ void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
     p.b_ld = static_cast<int>(b_ld);
 }
 
-// Pack the quantize params, optionally wiring the delayed-scaling ring.
-// ring (may be null) packs [hist[len] | scale | counter]; len/margin come
-// from the active recipe and idx is the caller's slot for this step.
-void pack_quantize_params(FP8QuantizeParams& p, const void* x, void* x8,
-                          const torch::Tensor& scale, torch::Tensor* amax,
-                          const torch::Tensor* ring, int64_t ring_idx,
-                          int64_t ring_margin, int64_t total) {
-    p.a_ptr = x;
-    p.out_ptr = x8;
-    p.scale_a = scale.data_ptr<float>();
-    p.amax_a = amax ? amax->data_ptr<float>() : nullptr;
-    if (ring && ring->defined()) {
-        TORCH_CHECK(ring->is_cuda() && ring->scalar_type() == torch::kFloat32 &&
-                        ring->numel() >= 3 && ring->is_contiguous(),
-                    "ring must be a contiguous CUDA float32 tensor packing "
-                    "[hist | scale | counter]");
-        p.ring_state = ring->data_ptr<float>();
-        p.ring_len = static_cast<int>(ring->numel() - 2);
-        p.ring_idx = static_cast<int>(ring_idx);
-        p.ring_margin = static_cast<int>(ring_margin);
-    }
-    p.total = static_cast<int>(total);
-}
-
-// ---- GEMM launch dispatch (runtime flags -> compile-time kernel variants) ----
-
 template <FP8Format Fmt, int Variant>
-void launch_gemm_variant(const FP8Params& p, cudaStream_t stream) {
+void launch_variant(const FP8Params& p, cudaStream_t stream) {
-    static_assert(Variant >= 0 && Variant < 8,
-                  "invalid FP8 GEMM dispatch variant");
-    constexpr bool out_fp8 = (Variant & 4) != 0;
-    // Variant bits 1/0 = trans_a/trans_b -> CUTLASS-style layout tags
-    // (trans_a ? A ColMajor : RowMajor, same for B; see common.h).
     using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>;
     using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>;
-    launch_fp8_gemm<Fmt, out_fp8, LayoutA, LayoutB>(p, stream);
+    launch_fp8_gemm<Fmt, LayoutA, LayoutB>(p, stream);
 }
 
 template <FP8Format Fmt>
-void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool out_fp8,
+void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a,
-                   bool trans_a, bool trans_b) {
+                   bool trans_b) {
-    // Encode the runtime flags as [output FP8, transpose A, transpose B].
+    const int variant = (static_cast<int>(trans_a) << 1) |
-    const int variant = (static_cast<int>(out_fp8) << 2) |
-                        (static_cast<int>(trans_a) << 1) |
                         static_cast<int>(trans_b);
     switch (variant) {
-        case 0: launch_gemm_variant<Fmt, 0>(p, stream); break;
+        case 0: launch_variant<Fmt, 0>(p, stream); break;
-        case 1: launch_gemm_variant<Fmt, 1>(p, stream); break;
+        case 1: launch_variant<Fmt, 1>(p, stream); break;
-        case 2: launch_gemm_variant<Fmt, 2>(p, stream); break;
+        case 2: launch_variant<Fmt, 2>(p, stream); break;
-        case 3: launch_gemm_variant<Fmt, 3>(p, stream); break;
+        case 3: launch_variant<Fmt, 3>(p, stream); break;
-        case 4: launch_gemm_variant<Fmt, 4>(p, stream); break;
-        case 5: launch_gemm_variant<Fmt, 5>(p, stream); break;
-        case 6: launch_gemm_variant<Fmt, 6>(p, stream); break;
-        case 7: launch_gemm_variant<Fmt, 7>(p, stream); break;
     }
 }
 
 }  // namespace
 
-// ---------------------------------------------------------------------------
+std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
-// Entry points
+                                                  torch::Tensor scale,
-// ---------------------------------------------------------------------------
+                                                  int64_t fmt) {
-
+    TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
-std::tuple<torch::Tensor, torch::Tensor> quantize_bf16(torch::Tensor x,
+    TORCH_CHECK(x.scalar_type() == torch::kBFloat16 ||
-                                                       torch::Tensor scale,
+                    x.scalar_type() == torch::kHalf ||
-                                                       int64_t fmt) {
+                    x.scalar_type() == torch::kFloat32,
-    // BF16 -> FP8 quantize with fused amax. fmt: 0 = E4M3, 1 = E5M2.
+                "x must be bf16, fp16 or fp32");
-    // Returns (x8, amax); the caller never clears amax (zero-initialized here).
+    TORCH_CHECK(fmt == static_cast<int64_t>(FP8Format::E4M3) ||
-    TORCH_CHECK(x.is_cuda() && scale.is_cuda(), "CUDA tensors required");
+                    fmt == static_cast<int64_t>(FP8Format::E5M2),
-    TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
+                "unsupported quantization type: expected E4M3 (0) or E5M2 (1)");
-    check_scale(scale, x, "scale");
+    check_scale(scale, x);
     check_fp8_device(x);
     const at::cuda::OptionalCUDAGuard guard(x.device());
     auto stream = at::cuda::getCurrentCUDAStream();
-
+    auto input = x.contiguous();
-    auto x_c = x.contiguous();
+    auto output = torch::empty_like(
-    auto x8 = torch::empty_like(
+        input, input.options().dtype(fmt ? torch::kFloat8_e5m2
-        x_c, x_c.options().dtype(fmt ? torch::kFloat8_e5m2
+                                          : torch::kFloat8_e4m3fn));
-                                      : torch::kFloat8_e4m3fn));
+    auto amax = torch::zeros({1}, input.options().dtype(torch::kFloat32));
-    // amax feeds the cross-block atomic_max; zero it on the stream (empty +
-    // memset, not torch::zeros — the latter routes through a fill_ dispatcher).
-    auto amax = torch::empty({1}, x_c.options().dtype(torch::kFloat32));
-    cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream());
     FP8QuantizeParams p;
-    pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax,
+    pack_quantize(p, input.data_ptr(), output.data_ptr(), scale, amax,
-                         nullptr, 0, 0, x_c.numel());
+                  input.numel());
-    if (fmt) {
+    const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
-        launch_fp8_quantize<FP8Format::E5M2>(p, stream.stream());
+    if (x.scalar_type() == torch::kHalf) {
+        if (e5m2)
+            launch_fp8_quantize<FP8Format::E5M2, __half>(p, stream.stream());
+        else
+            launch_fp8_quantize<FP8Format::E4M3, __half>(p, stream.stream());
+    } else if (x.scalar_type() == torch::kFloat32) {
+        if (e5m2)
+            launch_fp8_quantize<FP8Format::E5M2, float>(p, stream.stream());
+        else
+            launch_fp8_quantize<FP8Format::E4M3, float>(p, stream.stream());
     } else {
-        launch_fp8_quantize<FP8Format::E4M3>(p, stream.stream());
+        if (e5m2)
+            launch_fp8_quantize<FP8Format::E5M2, __nv_bfloat16>(
+                p, stream.stream());
+        else
+            launch_fp8_quantize<FP8Format::E4M3, __nv_bfloat16>(
+                p, stream.stream());
     }
     C10_CUDA_CHECK(cudaGetLastError());
-    return {x8, amax};
+    return {output, amax};
 }
 
-torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
+torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
-                     torch::Tensor sb, int64_t out_dtype,
+                     int64_t trans_a, int64_t trans_b) {
-                     c10::optional<torch::Tensor> out_scale, int64_t trans_a,
-                     int64_t trans_b) {
-    // Pre-quantized FP8 GEMM: out = op(a) @ op(b)^T * (sa * sb), FP32 accum.
-    // trans_a / trans_b select the operand layout (0 = stored [M,K]/[K,N],
-    // 1 = transposed [K,M]/[N,K]); the default (0/0) is the plain a @ b.
-    // out_dtype: 0 = BF16 (default), 1 = FP8 E4M3 (requires out_scale, the
-    // output quantization step). Both operands share one format.
     TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
     TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn ||
                     a.scalar_type() == torch::kFloat8_e5m2,
-                "a and b must be fp8 (e4m3fn or e5m2)");
+                "a and b must be fp8");
-    TORCH_CHECK(a.scalar_type() == b.scalar_type(),
+    TORCH_CHECK(a.scalar_type() == b.scalar_type(), "a and b must share format");
-                "a and b must share the same fp8 format");
     TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D");
-    TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device");
+    TORCH_CHECK(a.device() == b.device(), "a and b must share device");
-    check_scale(sa, a, "sa");
+    check_scale(scale, a);
-    check_scale(sb, a, "sb");
     check_fp8_device(a);
     const at::cuda::OptionalCUDAGuard guard(a.device());
     auto stream = at::cuda::getCurrentCUDAStream();
-
     auto a_c = a.contiguous();
     auto b_c = b.contiguous();
-    const bool ta = (trans_a == 1), tb = (trans_b == 1);
+    const bool ta = trans_a != 0;
-    // Physical leading dimension = column count of each contiguous buffer.
+    const bool tb = trans_b != 0;
     const int64_t a_ld = a_c.size(1);
     const int64_t b_ld = b_c.size(1);
-    // Logical GEMM shape derived from the layout flags.
     const int64_t m = ta ? a_c.size(1) : a_c.size(0);
     const int64_t k = ta ? a_c.size(0) : a_c.size(1);
     const int64_t n = tb ? b_c.size(0) : b_c.size(1);
-    const int64_t k2 = tb ? b_c.size(1) : b_c.size(0);
+    TORCH_CHECK(k == (tb ? b_c.size(1) : b_c.size(0)), "inner dim mismatch");
-    TORCH_CHECK(k == k2, "inner dim mismatch");
+    auto output = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16));
-    const bool out_fp8 = (out_dtype == 1);
-    TORCH_CHECK(out_dtype == 0 || out_fp8,
-                "out_dtype must be 0 (bf16) or 1 (fp8 e4m3)");
-    torch::Tensor os;
-    if (out_fp8) {
-        TORCH_CHECK(out_scale.has_value(), "fp8 output requires out_scale");
-        os = out_scale.value();
-        check_scale(os, a, "out_scale");
-    }
-    auto out = torch::empty(
-        {m, n},
-        out_fp8 ? a_c.options().dtype(torch::kFloat8_e4m3fn)
-                : a_c.options().dtype(torch::kBFloat16));
     FP8Params p;
-    pack_gemm_params(p, a_c.data_ptr(), b_c.data_ptr(), out.data_ptr(), sa, sb,
+    pack_gemm(p, a_c.data_ptr(), b_c.data_ptr(), output.data_ptr(), scale, m, n,
-                     out_fp8 ? &os : nullptr, nullptr, nullptr, m, n, k, a_ld,
+              k, a_ld, b_ld);
-                     b_ld);
     if (a.scalar_type() == torch::kFloat8_e4m3fn)
-        dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), out_fp8, ta, tb);
+        dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), ta, tb);
     else
-        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), out_fp8, ta, tb);
+        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), ta, tb);
     C10_CUDA_CHECK(cudaGetLastError());
-    return out;
+    return output;
-}
-
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
-           torch::Tensor>
-linear_forward_fp8(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
-                   torch::Tensor sx, torch::Tensor sw, int64_t fmt,
-                   c10::optional<torch::Tensor> bias_scale,
-                   c10::optional<torch::Tensor> x_ring, int64_t x_ring_idx,
-                   int64_t x_ring_margin, c10::optional<torch::Tensor> w_ring,
-                   int64_t w_ring_idx, int64_t w_ring_margin) {
-    // Pure FP8 forward: quantize x/w (fmt: 0 = E4M3, 1 = E5M2), then the
-    // pre-quantized GEMM; the dequantized BF16 output gets the bias added.
-    // Returns (out, x8, w8, amax_x, amax_w): the quantized operands are
-    // handed back so the policy layer can cache the weight quantization
-    // (torch autocast's cached_cast analog — w8 is reused while the weight
-    // tensor is unchanged, and the backward can share x8/w8 when the fwd/bwd
-    // formats match). amax_x / amax_w come from the quantize kernels
-    // (zero-initialized here; a pre-quantized w reports amax_w = 0 — nothing
-    // to feed a delayed ring). w may itself be pre-quantized fp8 storage
-    // matching fmt (static inference weights): the weight quantize is
-    // skipped, amax_w stays 0, and w8 returns the passed-in w.
-    // When x_ring / w_ring are given (delayed scaling), the quantize kernels
-    // finalize them in-kernel: the returned amax is already folded into the
-    // ring window and the next step's scale is published on device.
-    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required");
-    const auto f8opt = fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn;
-    const bool w_prequant = w.scalar_type() == f8opt;
-    TORCH_CHECK(
-        x.scalar_type() == torch::kBFloat16 &&
-            (w.scalar_type() == torch::kBFloat16 || w_prequant),
-        "x must be bf16; w must be bf16 or pre-quantized fp8 matching fmt");
-    TORCH_CHECK(x.device() == w.device(), "x and w must be on the same device");
-    check_scale(sx, x, "sx");
-    check_scale(sw, x, "sw");
-    check_fp8_device(x);
-    const at::cuda::OptionalCUDAGuard guard(x.device());
-    auto stream = at::cuda::getCurrentCUDAStream();
-
-    auto x_c = x.reshape({-1, w.size(1)}).contiguous();   // [M, K]
-    auto w_c = w.contiguous();                            // [N, K]
-    int64_t m = x_c.size(0), k = x_c.size(1), n = w_c.size(0);
-    TORCH_CHECK(w_c.dim() == 2 && w_c.size(1) == k, "inner dim mismatch");
-    const bool has_bias = bias.defined() && bias.numel() > 0;
-    const bool b_prequant = has_bias && bias.scalar_type() == f8opt;
-    if (has_bias) {
-        TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() &&
-                        bias.numel() == n &&
-                        (bias.scalar_type() == torch::kBFloat16 || b_prequant),
-                    "bias must be CUDA bf16 or pre-quantized fp8 matching fmt, "
-                    "with shape [N]");
-        TORCH_CHECK(b_prequant == bias_scale.has_value(),
-                    "fp8 bias requires bias_scale (and bf16 bias takes none)");
-        if (b_prequant) check_scale(*bias_scale, x, "bias_scale");
-    }
-    auto x8 = torch::empty({m, k}, x_c.options().dtype(f8opt));
-    // Each amax slot feeds a cross-block atomic_max, so it must start at 0.
-    // torch::zeros would route through a fill_ dispatcher (~50us CPU per call
-    // in the profile); a caching-allocator empty + cudaMemsetAsync is ~2us.
-    // Zero both up front: the pre-quantized-w path never quantizes w, so its
-    // amax_w is never atomically written and must not carry stale bytes. The
-    // returned values are the freshly measured (or 0) amax either way.
-    auto amax_x = torch::empty({1}, x.options().dtype(torch::kFloat32));
-    auto amax_w = torch::empty({1}, x.options().dtype(torch::kFloat32));
-    cudaMemsetAsync(amax_x.data_ptr(), 0, sizeof(float), stream.stream());
-    cudaMemsetAsync(amax_w.data_ptr(), 0, sizeof(float), stream.stream());
-    auto out = torch::empty({m, n}, x_c.options());
-
-    auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst,
-                        const torch::Tensor& scale, torch::Tensor* amax,
-                        const c10::optional<torch::Tensor>& ring,
-                        int64_t ring_idx, int64_t ring_margin) {
-        FP8QuantizeParams qp;
-        pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax,
-                             ring ? &*ring : nullptr, ring_idx, ring_margin,
-                             src.numel());
-        if (fmt) {
-            launch_fp8_quantize<FP8Format::E5M2>(qp, stream.stream());
-        } else {
-            launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
-        }
-    };
-    quantize(x_c, x8, sx, &amax_x, x_ring, x_ring_idx, x_ring_margin);
-    // Static inference weights arrive pre-quantized (w8 storage + its scale);
-    // only freshly-loaded bf16 weights quantize here.
-    torch::Tensor w8 = w_prequant
-                           ? w_c
-                           : torch::empty({n, k}, x_c.options().dtype(f8opt));
-    if (!w_prequant)
-        quantize(w_c, w8, sw, &amax_w, w_ring, w_ring_idx, w_ring_margin);
-
-    FP8Params p;
-    // Forward is the NT layout: A = x8 [M,K] (a_ld = k), B = w8 [N,K]
-    // (b_ld = k), out = x @ w^T. The bias is fused into the epilogue (bf16
-    // raw, or fp8 + bias_scale on the static path).
-    auto bias_c = has_bias ? bias.contiguous() : bias;
-    pack_gemm_params(p, x8.data_ptr(), w8.data_ptr(), out.data_ptr(), sx, sw,
-                     nullptr, has_bias ? bias_c.data_ptr() : nullptr,
-                     b_prequant ? &*bias_scale : nullptr, m, n, k, k, k);
-    if (fmt) {
-        launch_fp8_gemm<FP8Format::E5M2, false, RowMajor, ColMajor>(
-            p, stream.stream());
-    } else {
-        launch_fp8_gemm<FP8Format::E4M3, false, RowMajor, ColMajor>(
-            p, stream.stream());
-    }
-    C10_CUDA_CHECK(cudaGetLastError());
-
-    std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
-    shape.push_back(n);
-    return {out.reshape(shape), x8, w8, amax_x, amax_w};
-}
-
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
-linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
-                    std::vector<int64_t> masks, torch::Tensor sg,
-                    torch::Tensor sw, torch::Tensor sx, int64_t fmt,
-                    c10::optional<torch::Tensor> g_ring, int64_t g_ring_idx,
-                    int64_t g_ring_margin) {
-    // Pre-quantized FP8 backward: grad is quantized once (E4M3 or E5M2 per
-    // `fmt`), then dX / dW run as FP8 tensor-core GEMMs sharing g8.
-    // Returns (grad_input, grad_weight, grad_bias, amax_g). With g_ring
-    // (delayed scaling), the g quantize kernel finalizes the ring in-kernel
-    // (amax folded into the window, next step's scale published on device);
-    // 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");
-    // Each operand may be bf16 (quantized here) or already fp8 matching fmt
-    // (reused from the forward — the fp8 cached_cast analog, symmetric with
-    // the forward's pre-quantized w path). Pre-quantized operands skip their
-    // 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");
-    check_fp8_device(g);
-    const at::cuda::OptionalCUDAGuard guard(g.device());
-    auto stream = at::cuda::getCurrentCUDAStream();
-
-    auto g_c = g.reshape({-1, w.size(0)}).contiguous();   // [M, N]
-    auto x_c = x.reshape({-1, x.size(-1)}).contiguous();  // [M, K]
-    auto w_c = w.contiguous();                            // [N, K]
-    int64_t m = g_c.size(0), n = w_c.size(0), k = w_c.size(1);
-    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, 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).
-    // Only needed when a g quantize runs (mask[0]||mask[1]); the bias-only
-    // fallback below overwrites it via .copy_, so a wasted memset elsewhere
-    // is harmless.
-    auto amax_g = torch::empty({1}, g.options().dtype(torch::kFloat32));
-    cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), stream.stream());
-
-    auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst,
-                        const torch::Tensor& scale, torch::Tensor* amax,
-                        const c10::optional<torch::Tensor>& ring,
-                        int64_t ring_idx, int64_t ring_margin) {
-        FP8QuantizeParams qp;
-        pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax,
-                             ring ? &*ring : nullptr, ring_idx, ring_margin,
-                             src.numel());
-        if (fmt) {
-            launch_fp8_quantize<FP8Format::E5M2>(qp, stream.stream());
-        } else {
-            launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
-        }
-    };
-    // Four-layout backward: the gradient and activation tensors keep their
-    // natural row-major layout, and the kernel reads them transposed where the
-    // GEMM needs it (the ColMajor layout tags pick the crosswise stage-load).
-    // No torch-level `.transpose().contiguous()`
-    // copies are required — dX uses g8 [M,N] as A with w8 [N,K] read transposed
-    // as B; dW uses g8 transposed as A with x8 transposed as B.
-    // g is quantized once (amax_g measured here); both GEMMs share g8.
-    auto run_bwd_gemm = [&](const FP8Params& gp, bool trans_a, bool trans_b) {
-        if (fmt)
-            dispatch_gemm<FP8Format::E5M2>(gp, stream.stream(), false, trans_a,
-                                           trans_b);
-        else
-            dispatch_gemm<FP8Format::E4M3>(gp, stream.stream(), false, trans_a,
-                                           trans_b);
-    };
-
-    torch::Tensor g8;
-    if (masks[0] || masks[1]) {
-        if (g_prequant) {
-            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]) {
-        torch::Tensor w8;
-        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(),
-                         grad_input_2d.data_ptr(), sg, sw, nullptr, nullptr,
-                         nullptr, m, k, n, n, k);
-        run_bwd_gemm(gp, false, false);
-    }
-    // dW = g^T @ x: A = g8 [M,N] read transposed (a[p*a_ld + m] = g[p,m]), B =
-    // 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]) {
-        torch::Tensor x8;
-        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] && !g_prequant) {
-        amax_g.copy_(g_c.abs().amax().to(torch::kFloat32));
-    }
-    C10_CUDA_CHECK(cudaGetLastError());
-    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};
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("quantize_bf16", &quantize_bf16, py::arg("x"), py::arg("scale"),
+    m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
-          py::arg("fmt"),
+          py::arg("fmt"));
-          "BF16 to FP8 (E4M3/E5M2) quantize with fused amax; returns (x8, amax)");
+    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
-    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("sa"),
+          py::arg("trans_a") = 0, py::arg("trans_b") = 0);
-          py::arg("sb"), py::arg("out_dtype") = 0,
-          py::arg("out_scale") = py::none(), py::arg("trans_a") = 0,
-          py::arg("trans_b") = 0,
-          "Pre-quantized FP8 GEMM: op(a) @ op(b)^T * (sa * sb); out_dtype "
-          "0=bf16, 1=fp8 e4m3 (requires out_scale); trans_a/trans_b select "
-          "the operand layout (default 0/0 = a@b)");
-    m.def("linear_forward_fp8", &linear_forward_fp8, py::arg("x"),
-          py::arg("w"), py::arg("bias"), py::arg("sx"), py::arg("sw"),
-          py::arg("fmt") = 0, py::arg("bias_scale") = py::none(),
-          py::arg("x_ring") = py::none(), py::arg("x_ring_idx") = 0,
-          py::arg("x_ring_margin") = 0, py::arg("w_ring") = py::none(),
-          py::arg("w_ring_idx") = 0, py::arg("w_ring_margin") = 0,
-          "Pure FP8 linear forward: quantize x/w, pre-quantized GEMM with the "
-          "bias fused into the epilogue; w and bias may be pre-quantized fp8 "
-          "matching fmt (static inference path; fp8 bias requires bias_scale);"
-          " x_ring/w_ring optionally finalize a delayed-scaling ring "
-          "([hist | scale | counter] float32 buffer) in-kernel; returns "
-          "(out, x8, w8, amax_x, amax_w) — x8 is [M,K], w8 is [N,K] (the "
-          "passed-in w on the pre-quantized path)");
-    m.def("linear_backward_fp8", &linear_backward_fp8, py::arg("g"),
-          py::arg("x"), py::arg("w"), py::arg("masks"), py::arg("sg"),
-          py::arg("sw"), py::arg("sx"), py::arg("fmt"),
-          py::arg("g_ring") = py::none(), py::arg("g_ring_idx") = 0,
-          py::arg("g_ring_margin") = 0,
-          "FP8 linear backward; g_ring optionally finalizes the gradient's "
-          "delayed-scaling ring in-kernel; returns (grad_input, grad_weight, "
-          "grad_bias, amax_g)");
 }

csrc/kernels/fp8/quantize.cuh

@@ -0,0 +1,173 @@
+#pragma once
+// FP8 quantize device code — pure CUDA, no torch. Any float input element
+// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over
+// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the
+// FP8QuantizeParams POD, formats and input types ride on template parameters,
+// and the launcher is a plain function usable from both the torch binding and
+// pure C tests.
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+#include "common.h"
+#include "../common/reduce.cuh"
+
+namespace astrai {
+namespace fp8 {
+
+// Input element type traits: one element -> float, and the vectorized
+// unpack of one 16-byte load into kVecElems floats.
+template <typename InT>
+struct quant_in_traits;
+
+template <>
+struct quant_in_traits<__nv_bfloat16> {
+    static constexpr int kVecElems = 8;
+    static __device__ __forceinline__ float to_float(__nv_bfloat16 v) {
+        return __bfloat162float(v);
+    }
+    static __device__ __forceinline__ void load_vec(const uint4& raw,
+                                                    float* f) {
+        const unsigned w[4] = {raw.x, raw.y, raw.z, raw.w};
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            f[2 * j] =
+                __bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu));
+            f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16));
+        }
+    }
+};
+
+template <>
+struct quant_in_traits<__half> {
+    static constexpr int kVecElems = 8;
+    static __device__ __forceinline__ float to_float(__half v) {
+        return __half2float(v);
+    }
+    static __device__ __forceinline__ void load_vec(const uint4& raw,
+                                                    float* f) {
+        const __half2* h2 = reinterpret_cast<const __half2*>(&raw);
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const float2 p = __half22float2(h2[j]);
+            f[2 * j] = p.x;
+            f[2 * j + 1] = p.y;
+        }
+    }
+};
+
+template <>
+struct quant_in_traits<float> {
+    static constexpr int kVecElems = 4;
+    static __device__ __forceinline__ float to_float(float v) { return v; }
+    static __device__ __forceinline__ void load_vec(const uint4& raw,
+                                                    float* f) {
+        f[0] = __uint_as_float(raw.x);
+        f[1] = __uint_as_float(raw.y);
+        f[2] = __uint_as_float(raw.z);
+        f[3] = __uint_as_float(raw.w);
+    }
+};
+
+// Convert one float pair to one packed fp8 pair. The stored bytes see
+// value * mult (round-nearest-even + satfinite).
+template <FP8Format Fmt>
+__device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
+    constexpr __nv_fp8_interpretation_t kFmt =
+        Fmt == FP8Format::E5M2 ? __NV_E5M2 : __NV_E4M3;
+    return static_cast<unsigned>(__nv_cvt_float2_to_fp8x2(
+        make_float2(a, b), __NV_SATFINITE, kFmt));
+}
+
+// Quantize kernel: float input -> FP8 (E4M3 or E5M2), fused amax over raw
+// values.
+template <FP8Format Fmt, typename InT>
+__global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
+    const float mult = *p.scale;
+    const auto* x = static_cast<const InT*>(p.input_ptr);
+    void* x8 = p.output_ptr;
+    float* amax = p.amax;
+    float local_amax = 0.0f;
+    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
+
+    // Vectorized body: one 16B load -> kVecElems fp8 bytes per step (8
+    // elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B
+    // aligned and the binding passes freshly allocated contiguous buffers,
+    // so element 0 keeps the uint4 access natural; a misaligned base
+    // (contiguous view with an odd storage offset) falls back to the scalar
+    // loop below via total_vec = 0.
+    constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
+    const bool aligned =
+        ((reinterpret_cast<uintptr_t>(x) |
+          reinterpret_cast<uintptr_t>(x8)) &
+         15) == 0;
+    const int64_t total_vec = aligned ? p.total / kVecElems : 0;
+    const uint4* xv = reinterpret_cast<const uint4*>(x);
+    for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total_vec;
+         i += stride) {
+        float f[kVecElems];
+        quant_in_traits<InT>::load_vec(xv[i], f);
+        // One 32-bit word packs two fp8x2 pairs (4 elements).
+        unsigned packed[kVecElems / 4];
+#pragma unroll
+        for (int j = 0; j < kVecElems / 4; ++j) {
+            local_amax = fmaxf(
+                local_amax,
+                fmaxf(fmaxf(fabsf(f[4 * j]), fabsf(f[4 * j + 1])),
+                      fmaxf(fabsf(f[4 * j + 2]), fabsf(f[4 * j + 3]))));
+            const unsigned lo =
+                cvt_fp8x2<Fmt>(f[4 * j] * mult, f[4 * j + 1] * mult);
+            const unsigned hi =
+                cvt_fp8x2<Fmt>(f[4 * j + 2] * mult, f[4 * j + 3] * mult);
+            packed[j] = (lo & 0xffffu) | (hi << 16);
+        }
+        if constexpr (kVecElems == 8)
+            reinterpret_cast<uint2*>(x8)[i] =
+                make_uint2(packed[0], packed[1]);
+        else
+            reinterpret_cast<unsigned*>(x8)[i] = packed[0];
+    }
+    // Scalar tail (and full fallback for misaligned bases).
+    for (int64_t i = total_vec * kVecElems + blockIdx.x * blockDim.x +
+                      threadIdx.x;
+         i < p.total; i += stride) {
+        const float v = quant_in_traits<InT>::to_float(x[i]);
+        local_amax = fmaxf(local_amax, fabsf(v));
+        if constexpr (Fmt == FP8Format::E5M2) {
+            reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] =
+                __nv_fp8_e5m2(v * mult);
+        } else {
+            reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] =
+                __nv_fp8_e4m3(v * mult);
+        }
+    }
+    if (amax) {
+        local_amax = warp_reduce_max(local_amax);
+        __shared__ float slots[32];
+        if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax;
+        __syncthreads();
+        if (threadIdx.x == 0) {
+            float v = 0.0f;
+            for (int w = 0; w < (blockDim.x >> 5); ++w)
+                v = fmaxf(v, slots[w]);
+            atomic_max_float(amax, v);
+        }
+    }
+}
+
+template <FP8Format Fmt, typename InT>
+void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
+    constexpr int kThreads = 256;
+    // One block per 256 vectors; at least one block so the scalar tail of a
+    // tiny / misaligned tensor is still covered.
+    constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
+    int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
+    if (blocks < 1) blocks = 1;
+    fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
+}
+
+}  // namespace fp8
+}  // namespace astrai

csrc/tests/fp8_test.cu

@@ -175,15 +175,13 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
                           int k, int a_ld, int b_ld) {
     __nv_fp8_e4m3 *da, *db;
     __nv_bfloat16* dout;
-    float *dsa, *dsb;
+    float* dscale;
     cudaMalloc(&da, (size_t)m * k);
     cudaMalloc(&db, (size_t)n * k);
     cudaMalloc(&dout, (size_t)m * n * 2);
-    cudaMalloc(&dsa, 4);
+    cudaMalloc(&dscale, 4);
-    cudaMalloc(&dsb, 4);
     float one = 1.0f;
-    cudaMemcpy(dsa, &one, 4, cudaMemcpyHostToDevice);
+    cudaMemcpy(dscale, &one, 4, cudaMemcpyHostToDevice);
-    cudaMemcpy(dsb, &one, 4, cudaMemcpyHostToDevice);
     // quantize inputs to e4m3 on host and upload byte-by-byte
     std::vector<unsigned char> qa(m * k), qb(n * k);
     for (int i = 0; i < m * k; ++i) {
@@ -201,14 +199,13 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     p.a_ptr = da;
     p.b_ptr = db;
     p.out_ptr = dout;
-    p.scale_a = dsa;
+    p.scale = dscale;
-    p.scale_b = dsb;
     p.m = m;
     p.n = n;
     p.k = k;
     p.a_ld = a_ld;
     p.b_ld = b_ld;
-    launch_fp8_gemm<FP8Format::E4M3, false, LA, LB, kK, Stages>(p, 0);
+    launch_fp8_gemm<FP8Format::E4M3, LA, LB, kK, Stages>(p, 0);
     cudaError_t e = cudaDeviceSynchronize();
     if (e != cudaSuccess) {
         printf("  CUDA err: %s\n", cudaGetErrorString(e));
@@ -247,8 +244,7 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     cudaFree(da);
     cudaFree(db);
     cudaFree(dout);
-    cudaFree(dsa);
+    cudaFree(dscale);
-    cudaFree(dsb);
     return ok;
 }
 

docs/developer/cuda_kernels.md

@@ -47,16 +47,17 @@ style as attention, but split into **three** files:
 | File | Role |
 |------|------|
 | `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `FP8Params` POD — no torch |
-| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_quantize_kernel` (BF16→FP8 + amax), `fp8_gemm_kernel` (pre-quantized GEMM, 128×64 CTA / 64×16 warp / 3-stage cp.async) — no torch |
+| `fp8/quantize.cuh` | pure-CUDA device code: `fp8_quantize_kernel<Fmt, InT>` (bf16/fp16/fp32 → FP8 + amax, `quant_in_traits<InT>` vectorized unpack) — no torch |
+| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_gemm_kernel` (pre-quantized GEMM, 128×128 CTA / 64×32 warp / multi-stage cp.async, transposed-operand layouts) — no torch |
 | `fp8/ops.cu` | binding only: `check_fp8_device` (sm_89+), param packing, launch dispatch, pybind → module `fp8_ops` |
 
-Scale semantics follow `torch._scaled_mm` (quantization step size: divide by
+Scale semantics: `quantize` takes the quantization *multiplier*, `mm_fp8`
-`scale`; the kernel computes the reciprocal internally — the interface never
+takes the combined dequant scale (`sa * sb`); the strategy layer passes
-takes `*_inv`). `amax` is always returned in the original bf16 domain.
+`scale.reciprocal()` / `sa * sb` respectively. `amax` is always returned in
+the original input domain.
 
 Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless
-primitives (`quantize_bf16` / `mm_fp8` / `linear_forward_fp8` /
+primitives (`quantize` / `mm_fp8`) via `torch.library.custom_op`, and
-`linear_backward_fp8`) via `torch.library.custom_op`, and
 `astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed /
 dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear`
 on CUDA). See the FP8 section in `AGENTS.md` for full detail.

tests/extension/test_fp8_mma.py

@@ -1,9 +1,9 @@
 """FP8 primitives: kernel-level (CUDA) and policy-level (CPU-verifiable) tests.
 
-The kernel-level tests exercise the pure FP8 path (quantize_bf16 + mm_fp8 for
+The kernel-level tests exercise the two stateless primitives (``quantize`` for
-the forward GEMM, quantize + pre-quantized GEMMs for the backward); the
+bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed
-policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks
+operands); the policy-level tests (recipes, autocast context, per-tensor meta,
-of the custom ops) run without a GPU.
+CPU fallbacks of the custom ops) run without a GPU.
 """
 
 import threading
@@ -24,12 +24,7 @@ from astrai.extension.fp8 import (
     fp8_linear_enabled,
     fp8_state,
 )
-from astrai.extension.ops.fp8 import (
+from astrai.extension.ops.fp8 import mm_fp8, quantize
-    linear_backward_fp8,
-    linear_forward_fp8,
-    mm_fp8,
-    quantize_bf16,
-)
 from tests.conftest import skip_no_fp8
 
 
@@ -37,8 +32,11 @@ def _scale(tensor):
     return (tensor.abs().amax().float() / 448.0).clamp_min(1e-12)
 
 
-def _quantize(tensor, scale):
+def _quantize(tensor, scale, fmt="e4m3"):
-    return (tensor.float() / scale).to(torch.float8_e4m3fn).float()
+    """Reference quantize: multiply by the reciprocal (the kernel's exact
+    arithmetic — a plain divide flips fp8 boundary cases by one ulp)."""
+    dtype = torch.float8_e5m2 if fmt == "e5m2" else torch.float8_e4m3fn
+    return (tensor.float() * scale.reciprocal()).to(dtype).float()
 
 
 # --------------------------------------------------------------------------
@@ -57,9 +55,9 @@ def test_fp8_mm_matches_explicit_quantization(m, n, k):
     b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
     scale_a = _scale(a)
     scale_b = _scale(b)
-    a8, _ = quantize_bf16(a, scale_a, "e4m3")
+    a8, _ = quantize(a, scale_a.reciprocal(), "e4m3")
-    b8, _ = quantize_bf16(b, scale_b, "e4m3")
+    b8, _ = quantize(b, scale_b.reciprocal(), "e4m3")
-    out = mm_fp8(a8, b8, scale_a, scale_b)
+    out = mm_fp8(a8, b8, scale_a * scale_b)
     expected = (_quantize(a, scale_a) @ _quantize(b, scale_b) * scale_a * scale_b).to(
         torch.bfloat16
     )
@@ -70,79 +68,59 @@ def test_fp8_mm_matches_explicit_quantization(m, n, k):
 
 
 @skip_no_fp8
-def test_quantize_bf16_returns_amax():
+@pytest.mark.parametrize("in_dtype", [torch.bfloat16, torch.float16, torch.float32])
-    """quantize_bf16 returns (x8, amax); amax tracks the *raw* values and the
+@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"])
-    caller never clears it (zero-initialized inside the kernel entry)."""
+def test_quantize_input_dtypes(in_dtype, fmt):
+    """quantize accepts bf16/fp16/fp32 inputs; bytes and amax match the
+    explicit (value * multiplier) reference."""
     torch.manual_seed(3)
-    x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16)
+    x = torch.randn(64, 128, device="cuda", dtype=torch.float32) * 0.5
+    x = x.to(in_dtype)
     scale = torch.tensor([0.5], device="cuda")
-    x8, amax = quantize_bf16(x, scale, "e4m3")
+    x8, amax = quantize(x, scale, fmt)
-    assert x8.dtype == torch.float8_e4m3fn
+    out_dtype = torch.float8_e5m2 if fmt == "e5m2" else torch.float8_e4m3fn
+    assert x8.dtype == out_dtype
     assert x8.shape == x.shape
     assert amax.shape == (1,)
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
-    ref = (x.float() / 0.5).to(torch.float8_e4m3fn)
+    ref = (x.float() * 0.5).to(out_dtype)
     assert torch.equal(x8, ref)
 
 
 @skip_no_fp8
-def test_quantize_bf16_e5m2_format():
+def test_quantize_e5m2_format():
     x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
-    x8, amax = quantize_bf16(x, torch.tensor([0.1], device="cuda"), "e5m2")
+    x8, amax = quantize(x, torch.tensor([10.0], device="cuda"), "e5m2")
     assert x8.dtype == torch.float8_e5m2
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
 
 
 @skip_no_fp8
-def test_quantize_ring_in_kernel_finalize():
+@pytest.mark.parametrize("trans_a", [False, True])
-    """The quantize kernel finalizes the delayed-scaling ring in-kernel: the
+@pytest.mark.parametrize("trans_b", [False, True])
-    measured amax lands in hist[idx], the window reduces to the next step's
+def test_mm_fp8_transposed_operands(trans_a, trans_b):
-    scale on device, and the counter re-arms for the next launch."""
+    """mm_fp8 handles all four operand layouts via trans_a/trans_b."""
-    torch.manual_seed(21)
+    torch.manual_seed(17)
-    dev = torch.device("cuda")
+    m, n, k = 19, 13, 37
-    ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0))
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)  # A [M][K]
-    x0 = torch.randn(256, 256, device=dev, dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)  # B^T [N][K]
-    w = torch.randn(256, 256, device=dev, dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
-    sw = torch.tensor([1.0], device=dev)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
-    ring.seed(x0, "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
-    hist0 = ring.hist.clone()
+    a_op = a8.t().contiguous() if trans_a else a8
+    b_op = b8 if trans_b else b8.t().contiguous()
 
-    # Step over three fresh tensors: each launch folds its amax into
+    out = mm_fp8(a_op, b_op, sa * sb, trans_a=trans_a, trans_b=trans_b)
-    # hist[idx] and publishes max(hist)/448 as the next scale.
+    assert out.shape == (m, n)
-    idx = 0
+    expected = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
-    for _ in range(3):
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
-        x = torch.randn(256, 256, device=dev, dtype=torch.bfloat16) * (2.0 + 4.0 * _)
-        _ = linear_forward_fp8(
-            x,
-            w,
-            None,
-            ring.scale,
-            sw,
-            "e4m3",
-            None,
-            ring.state,
-            idx,
-            0,
-        )
-        torch.cuda.synchronize()
-        expected_hist = hist0.clone()
-        expected_hist[idx] = x.abs().amax().float()
-        torch.testing.assert_close(ring.hist, expected_hist)
-        expected_scale = (expected_hist.max() / 448.0).reshape(1)
-        torch.testing.assert_close(ring.scale, expected_scale, rtol=1e-6, atol=1e-12)
-        # counter re-armed to int32 zero
-        assert ring.state[-1].view(torch.int32).item() == 0
-        hist0 = expected_hist.clone()
-        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
+    """The delayed scale for step N is computed from amax(steps < N); the
-    aliases the scale the GEMM must dequantize with. The forward must snapshot
+    forward must snapshot the scale before the ring update, so a changing
-    the delayed scale first, so a changing amax across steps does not leak the
+    amax across steps does not leak the next-step scale into the output."""
-    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()
@@ -153,8 +131,7 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
         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 —
+        # 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)
@@ -177,93 +154,57 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
 
 @skip_no_fp8
 def test_fp8_linear_forward_and_backward():
+    """The composed strategy path: forward quantize+GEMM+bias, backward
+    dX/dW GEMMs on transposed operands (E5M2 in hybrid)."""
     torch.manual_seed(7)
     m, n, k = 19, 13, 37
     x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
     weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16)
     bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
-    scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad)
 
-    out, x8, w8, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w)
+    state = f8mod.fp8_state()
-    grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
+    state.reset()
-        grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3"
+    state.default_recipe = DynamicScaling()
-    )
+    try:
+        out, _, _ = f8mod.fp8_linear_forward(x, weight, bias)
 
-    qx = _quantize(x, scale_x)
+        sx, sw = _scale(x), _scale(weight)
-    qw = _quantize(weight, scale_w)
+        qx = _quantize(x, sx)
-    qg = _quantize(grad, scale_g)
+        qw = _quantize(weight, sw)
-    expected_out = (qx @ qw.t() * scale_x * scale_w + bias).to(torch.bfloat16)
+        expected_out = (qx @ qw.t() * sx * sw + bias).to(torch.bfloat16)
-    expected_grad_x = (qg @ qw * scale_g * scale_w).to(torch.bfloat16)
+        torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01)
-    expected_grad_w = (qg.t() @ qx * scale_g * scale_x).to(torch.bfloat16)
 
-    torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01)
+        # backward through the aten::linear integration (hybrid E5M2). The
-    torch.testing.assert_close(grad_x, expected_grad_x, atol=0.125, rtol=0.01)
+        # incoming gradient is 2*out of the *fp8* forward (bf16-rounded), not
-    torch.testing.assert_close(grad_w, expected_grad_w, atol=0.125, rtol=0.01)
+        # 2*exact — derive the reference from the actual output.
-    torch.testing.assert_close(grad_b, grad.sum(0).to(torch.bfloat16))
+        xr = x.detach().clone().requires_grad_()
-    torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1))
+        wr = weight.detach().clone().requires_grad_()
-    torch.testing.assert_close(amax_w, weight.abs().amax().float().reshape(1))
+        br = bias.detach().clone().requires_grad_()
-    torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
+        with fp8_autocast(enabled=True):
+            loss = F.linear(xr, wr, br).float().pow(2).sum()
+        loss.backward()
 
-
+        g = (2 * out.float()).to(torch.bfloat16).float()  # actual grad wrt out
-@skip_no_fp8
+        # the dynamic path measures current-step amax in the bwd fmt (E5M2);
-def test_linear_backward_e5m2_gradients():
+        # amax must be taken in fp32 — a bf16-rounded scale flips E5M2
-    """Hybrid backward: gradient GEMMs run in E5M2 (larger dynamic range)."""
+        # boundary rounding (2-bit mantissa) and the reference drifts.
-    torch.manual_seed(5)
+        e5 = 57344.0
-    m, n, k = 32, 16, 64
+        sg = (g.abs().amax() / e5).clamp_min(1e-12)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 3.0
+        sw5 = (weight.abs().amax().float() / e5).clamp_min(1e-12)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+        sx5 = (x.abs().amax().float() / e5).clamp_min(1e-12)
-    grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) * 10.0
+        expected_grad_x = (
-    sg = _scale(grad) * 0.5
+            _quantize(g, sg, "e5m2") @ _quantize(weight, sw5, "e5m2") * sg * sw5
-    sw = _scale(weight)
+        ).to(torch.bfloat16)
-    sx = _scale(x)
+        expected_grad_w = (
-
+            _quantize(g, sg, "e5m2").t() @ _quantize(x, sx5, "e5m2") * sg * sx5
-    grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
+        ).to(torch.bfloat16)
-        grad, x, weight, [1, 1, 1], sg, sw, sx, "e5m2"
+        torch.testing.assert_close(xr.grad, expected_grad_x, atol=0.5, rtol=0.05)
-    )
+        torch.testing.assert_close(wr.grad, expected_grad_w, atol=0.5, rtol=0.05)
-
+        torch.testing.assert_close(
-    def q5(t, s):
+            br.grad, g.sum(0).to(torch.bfloat16), atol=0.5, rtol=0.05
-        return (t.float() / s).to(torch.float8_e5m2).float()
+        )
-
+    finally:
-    qg = q5(grad, sg)
+        state.reset()
-    qw = q5(weight, sw)
-    qx = q5(x, sx)
-    expected_grad_x = (qg @ qw * sg * sw).to(torch.bfloat16)
-    expected_grad_w = (qg.t() @ qx * sg * sx).to(torch.bfloat16)
-    torch.testing.assert_close(grad_x, expected_grad_x, atol=0.5, rtol=0.05)
-    torch.testing.assert_close(grad_w, expected_grad_w, atol=0.5, rtol=0.05)
-    torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
-
-
-@skip_no_fp8
-def test_fp8_linear_static_fp8_weight_and_bias():
-    """Static fp8 inference: pre-quantized w8/b8 + their scales take the GEMM
-    directly (no weight quantize, amax_w = 0); the bias is fused in the
-    epilogue (bf16 and fp8 bias share the fused path)."""
-    torch.manual_seed(9)
-    m, n, k = 67, 45, 129
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.5
-    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) * 0.5
-    sx, sw, sb = _scale(x), _scale(weight), _scale(bias)
-
-    w8, _ = quantize_bf16(weight, sw, "e4m3")
-    b8, _ = quantize_bf16(bias, sb, "e4m3")
-    out, x8, w8_back, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb)
-
-    qx = _quantize(x, sx)
-    qw = _quantize(weight, sw)
-    qb = _quantize(bias, sb)
-    expected = (qx @ qw.t() * sx * sw + qb * sb).to(torch.bfloat16)
-    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
-    torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1))
-    assert amax_w.item() == 0.0  # nothing measured on the static path
-    assert w8_back is w8  # pre-quantized w handed straight back
-
-    # bf16 bias stays bf16 on the same fused-epilogue path
-    out_bf16bias, *_ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3")
-    expected_b = (qx @ qw.t() * sx * sw + bias.float()).to(torch.bfloat16)
-    torch.testing.assert_close(out_bf16bias, expected_b, atol=0.125, rtol=0.01)
 
 
 @skip_no_fp8
@@ -279,24 +220,24 @@ def test_fp8_linear_backward_outside_autocast():
     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}
+    calls = {"fwd": 0}
-    orig = f8mod.linear_backward_fp8
+    orig = f8mod.fp8_linear_forward
 
     def spy(*args, **kwargs):
-        calls["bwd"] += 1
+        calls["fwd"] += 1
         return orig(*args, **kwargs)
 
-    f8mod.linear_backward_fp8 = spy
+    f8mod.fp8_linear_forward = 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_linear_forward = orig
         f8mod.fp8_state().reset()
 
-    assert calls["bwd"] == 1  # fp8 kernels, not the bf16 fallback
+    assert calls["fwd"] == 1  # fp8 kernels, not the bf16 fallback
     ref = F.linear(xr, wr, br)
     ref.float().pow(2).sum().backward()
 
@@ -319,15 +260,15 @@ def test_mm_fp8_matches_scaled_mm():
     m, n, k = 512, 4096, 4096
     a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
     b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
-    sa = torch.tensor([2.5], device="cuda")
+    sa = _scale(a)
-    sb = torch.tensor([1.5], device="cuda")
+    sb = _scale(b)
-    a8, _ = quantize_bf16(a, sa, "e4m3")
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
-    b8, _ = quantize_bf16(b, sb, "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
-    out = mm_fp8(a8, b8, sa, sb)
+    out = mm_fp8(a8, b8, sa * sb)
     assert out.dtype == torch.bfloat16
     assert out.shape == (m, n)
 
-    ref = (a8.float().double() @ b8.float().double() * 2.5 * 1.5).to(torch.bfloat16)
+    ref = (a8.float().double() @ b8.float().double() * sa * sb).to(torch.bfloat16)
     torch.testing.assert_close(out, ref, atol=6.0, rtol=0.05)
 
     try:
@@ -342,30 +283,6 @@ def test_mm_fp8_matches_scaled_mm():
     )
 
 
-@skip_no_fp8
-def test_mm_fp8_fp8_output():
-    """mm_fp8 with out_dtype='e4m3' produces an FP8 output (layer-to-layer)."""
-    torch.manual_seed(12)
-    m, n, k = 256, 128, 64
-    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
-    sa = torch.tensor([2.0], device="cuda")
-    sb = torch.tensor([1.0], device="cuda")
-    os_ = torch.tensor([0.5], device="cuda")
-    a8, _ = quantize_bf16(a, sa, "e4m3")
-    b8, _ = quantize_bf16(b, sb, "e4m3")
-    out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_)
-    assert out8.dtype == torch.float8_e4m3fn
-    assert out8.shape == (m, n)
-
-    ref = (a8.float().double() @ b8.float().double() * 2.0 * 1.0 * 0.5).to(
-        torch.bfloat16
-    )
-    torch.testing.assert_close(
-        out8.float().to(torch.bfloat16), ref, atol=6.0, rtol=0.05
-    )
-
-
 # --------------------------------------------------------------------------
 # Policy-level (CPU-verifiable)
 # --------------------------------------------------------------------------
@@ -423,23 +340,26 @@ def test_fp8_tensor_meta_delayed_update():
     meta.w.seed(w, "e4m3")
     assert meta.w.initialized
     torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1))
-    # [hist | scale | counter] packing: views alias the single state buffer.
+    # [hist | scale] packing: views alias the single state buffer.
     assert meta.w.state.numel() == 4 + 2
     assert meta.w.hist.data_ptr() == meta.w.state.data_ptr()
     assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr()
-    # counter slot stays int32-zero (float bits) between launches
-    assert meta.w.state[-1].view(torch.int32).item() == 0
     meta.w.advance()
     assert meta.w.idx == 1
 
+    # update folds a fresh amax into the window and publishes the next scale
+    amax = torch.tensor([8.0])
+    meta.w.update(amax, "e4m3")
+    torch.testing.assert_close(meta.w.scale, torch.tensor([8.0 / 448.0]))
 
-def test_quantize_bf16_cpu_fallback():
+
+def test_quantize_cpu_fallback():
     """CPU fallback of the quantize primitive (scale semantics + amax)."""
     x = torch.randn(16, 32, dtype=torch.bfloat16)
-    scale = torch.tensor([0.5])
+    scale = torch.tensor([0.5])  # quantize multiplier
-    x8, amax = quantize_bf16(x, scale, "e4m3")
+    x8, amax = quantize(x, scale, "e4m3")
     assert x8.dtype == torch.float8_e4m3fn
-    ref = (x.float() / 0.5).to(torch.float8_e4m3fn)
+    ref = (x.float() * 0.5).to(torch.float8_e4m3fn)
     assert torch.equal(x8, ref)
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
 
@@ -447,26 +367,12 @@ def test_quantize_bf16_cpu_fallback():
 def test_mm_fp8_cpu_fallback():
     a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
     b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
-    sa = torch.tensor([2.0])
+    scale = torch.tensor([1.0])
-    sb = torch.tensor([0.5])
+    out = mm_fp8(a8, b8, scale)
-    out = mm_fp8(a8, b8, sa, sb)
+    ref = (a8.float() @ b8.float() * 1.0).to(torch.bfloat16)
-    ref = (a8.float() @ b8.float() * 2.0 * 0.5).to(torch.bfloat16)
     torch.testing.assert_close(out, ref)
 
 
-def test_mm_fp8_fp8_output_cpu():
-    """CPU fallback with an FP8 output (out_dtype='e4m3' + out_scale)."""
-    a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
-    b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
-    sa = torch.tensor([2.0])
-    sb = torch.tensor([0.5])
-    os_ = torch.tensor([0.25])
-    out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_)
-    assert out8.dtype == torch.float8_e4m3fn
-    ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn)
-    assert torch.equal(out8, ref)
-
-
 # --------------------------------------------------------------------------
 # torch-autocast parity: context semantics (nesting, thread locality, switch)
 # --------------------------------------------------------------------------
@@ -488,7 +394,7 @@ def test_nested_disabled_region_redispatches_bf16():
     resumes when it exits."""
     x, w = _linear()
     with fp8_autocast(enabled=True):
-        out_fp8 = F.linear(x, w)
+        F.linear(x, w)
         with fp8_autocast(enabled=False):
             out_bf16 = F.linear(x, w)
             assert type(out_bf16.grad_fn).__name__ != "_LinearFp8Backward"