perf: fold the delayed-scaling ring update into the quantize kernel

- the kernel's last block folds amax into the history window and publishes the next scale in-kernel (atomicAdd ticket + fences), replacing the host update chain - quantize bindings split into quantize(transposed) / quantize_dual with fixed arities and a QuantLayout enum; the python adapter becomes a thin attention-style wrapper over pybind (Optional ring_state at the boundary, no torch.library custom_ops) - tests: in-kernel fold vs host reference (exact), dual/transposed orientation byte-equality Benchmark: L20 (sm_89), 1.2B model, full train step. Per-linear fixed overhead 28.8us -> 8.8us; fp8 vs bf16: M=512 77.5ms, M=2048 144.5ms (1.15x), M=8192 527.4ms (1.28x); losses bit-identical.
2026-08-31 14:24:51 +08:00
parent 1cf7d6c76b
commit 962c10c52b
6 changed files with 324 additions and 320 deletions
@@ -36,7 +36,7 @@ from typing import Dict, List, Optional
 import torch
 from torch.library import Library
-from astrai.extension.ops.fp8 import mm_fp8, quantize
+from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
 # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
 FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -92,10 +92,12 @@ class DynamicScaling(FP8Recipe):
 class _ScaleRing:
    """One operand's delayed-scaling state: a float32 buffer
-    ``[hist[n] | scale | counter]`` (views). ``update`` folds the amax
+    ``[hist[n] | scale | legacy | amax | done]`` (views). The quantize
-    returned by the quantize primitive into ``hist[idx]`` and publishes the
+    kernel folds its fused amax into ``hist[idx]`` and publishes the next
-    next scale from the window; ``idx`` advances host-side each step. The
+    scale from the window in its own last block (``fold_args`` passes the
-    trailing slot is a legacy counter kept for state-buffer compatibility.
+    buffer + recipe constants); ``idx`` advances host-side each use. The
    ``amax``/``done`` tail slots are kernel scratch (self-cleaning across
    launches); the legacy slot keeps state-buffer compatibility.
    """
    __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -103,7 +105,7 @@ class _ScaleRing:
    def __init__(self, device: torch.device, recipe: FP8Recipe):
        self.recipe = recipe
        n = recipe.history_len
-        self.state = torch.zeros(n + 2, device=device, dtype=torch.float32)
+        self.state = torch.zeros(n + 4, device=device, dtype=torch.float32)
        self.hist = self.state[:n]
        self.scale = self.state[n : n + 1]
        self.idx = 0
@@ -119,9 +121,14 @@ 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:
+    def fold_args(self, fmt: str) -> dict:
-        self.hist[self.idx].copy_(amax.reshape(()))
+        """Keyword arguments for quantize()'s in-kernel history fold."""
-        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
+        return {
            "ring_state": self.state,
            "hist_idx": self.idx,
            "fp8_max": FP8_MAX[fmt],
            "pow2_margin": float(2**self.recipe.margin),
        }
 class FP8TensorMeta:
@@ -322,11 +329,11 @@ def fp8_linear_forward(
    Composed from the two stateless primitives: quantize x/w with the active
    scales, run the pre-quantized GEMM with the bias fused into its epilogue.
-    Delayed scaling folds
+    Delayed scaling lets the quantize kernel fold the fused amax into the
-    the returned amax into the history ring and publishes the next scale;
+    history ring and publish the next scale in its own last block; dynamic
-    dynamic scaling measures the current amax itself. Training quantizes the
+    scaling measures the current amax itself. Training quantizes the weight
-    weight every step (the optimizer bumps its version, so there is no cast
+    every step (the optimizer bumps its version, so there is no cast cache,
-    cache, matching ``cached_cast``-less behavior).
+    matching ``cached_cast``-less behavior).
    """
    state = fp8_state()
    if cfg is None:
@@ -351,19 +358,19 @@ def fp8_linear_forward(
    if not meta.x.initialized:
        meta.x.seed(x, fmt)
    sx, sw = meta.x.scale.clone(), meta.w.scale.clone()
-    x8, amax_x = quantize(x, sx.reciprocal(), fmt)
+    # The clones feed this call's kernels (stream-ordered before the in-kernel
    # fold overwrites the ring scale slots); the fp8 quantize kernel folds the
    # amax into the history window and publishes the next scale itself.
    x8, _ = quantize(x, sx.reciprocal(), fmt, **meta.x.fold_args(fmt))
    if _is_fp8(w.dtype):
-        w8, amax_w = w, None
+        w8 = w
    else:
-        w8, amax_w = quantize(w, sw.reciprocal(), fmt)
+        w8, _ = quantize(w, sw.reciprocal(), fmt, **meta.w.fold_args(fmt))
    out = mm_fp8(
        x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
    ).reshape(*x.shape[:-1], w.size(0))
    meta.x.update(amax_x, fmt)
    if amax_w is not None:
        meta.w.update(amax_w, fmt)
    meta.x.advance()
-    if amax_w is not None:
+    if not _is_fp8(w.dtype):
        meta.w.advance()
    return out, sx, sw
@@ -411,22 +418,26 @@ class _LinearFp8(torch.autograd.Function):
        # quantize outputs: g8 [m,n] with w8T [k,n] (trans_b=True) gives
        # grad_x, g8T [n,m] with x8T [k,m] gives grad_w — no NN-swap or TT
        # crosswise kernel in the training path. g is consumed in both
-        # orientations, so one dual-layout pass feeds both.
+        # orientations, so quantize_dual's single pass feeds both.
-        g8, g8T, amax_g = quantize(g2, sg.reciprocal(), fmt, layout=2)
+        # The g quantize folds the gradient amax into its ring in-kernel;
-        x8T, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, layout=1)
+        # the x8T/w8T orientation copies discard amax (those rings were
        # folded at forward time).
        g8, g8T, _ = quantize_dual(g2, sg.reciprocal(), fmt, **meta.g.fold_args(fmt))
        x8T, _ = quantize(
            x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, transposed=True
        )
        if _is_fp8(w.dtype):
            # Pre-quantized weight has no transposed copy: keep the swap
            # path for grad_x (grad_w is unaffected).
            grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape)
        else:
-            w8T, _ = quantize(w, sw.reciprocal(), fmt, layout=1)
+            w8T, _ = quantize(w, sw.reciprocal(), fmt, transposed=True)
            grad_x = mm_fp8(g8, w8T, sg * sw, trans_b=True).reshape(x.shape)
        grad_w = mm_fp8(g8T, x8T, sg * sx, trans_b=True)  # g8.T @ x8
        # bias-free linears must not pay the column-sum
        # reduce: g2.sum(0) is another full read of the gradient.
        grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None
        if not ctx.is_dynamic:
            meta.g.update(amax_g, fmt)
            meta.g.advance()
        return grad_x, grad_w, grad_b
@@ -1,14 +1,20 @@
 """FP8 CUDA kernel interface adapter (the only module touching the pybind).
-Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
+Attention-style thin wrappers: one Python entry per binding, called directly
 — no torch.library dispatch layer. Optional arguments (``ring_state``,
 ``bias``) keep native Optional semantics at the pybind boundary, and
 in-place buffer updates (the delayed-scaling ring fold, like attention's
 KV-cache appends) happen on-stream without mutation declarations. CUDA-only:
 non-CUDA or unsupported inputs raise from the binding's TORCH_CHECKs.
- ``quantize(x, scale, fmt) -> (x8, amax)`` — BF16/FP16/FP32 → FP8 with fused amax
+- ``quantize(x, scale, fmt, transposed=False) -> (x8|x8T, amax)`` — BF16/FP16/FP32
  → FP8 with fused amax (``transposed`` picks the orientation; arity is fixed)
 - ``quantize_dual(x, scale, fmt) -> (x8, x8T, amax)`` — both orientations, one read
 - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
-Scale semantics: scales are *quantization steps* — the value divided out when
+``scale`` is the quantization multiplier (device scalar); ``fmt`` is
-quantizing (``x8 = x / scale``). Every primitive computes its own inverse
+``"e4m3"`` or ``"e5m2"``. ``amax`` values are *returned*, never passed as
-internally; callers never pass ``scale_inv``. ``amax`` values are *returned*,
+output arguments.
 never passed as output arguments. ``fmt`` is ``"e4m3"`` or ``"e5m2"``.
 Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
 this module is stateless.
@@ -17,7 +23,6 @@ this module is stateless.
 from typing import Optional, Tuple
 import torch
 from torch.library import custom_op
 from astrai.extension.loader import get_module
@@ -32,192 +37,63 @@ 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
@custom_op("custom::fp8_quantize", mutates_args=())
 def fp8_quantize(
    x: torch.Tensor, scale: torch.Tensor, fmt: int
 ) -> Tuple[torch.Tensor, torch.Tensor]:
    """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 == 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)
@custom_op("custom::fp8_quantize_t", mutates_args=())
 def fp8_quantize_t(
    x: torch.Tensor, scale: torch.Tensor, fmt: int
 ) -> Tuple[torch.Tensor, torch.Tensor]:
    """Transposed-output variant of fp8_quantize: returns ``(x8T, amax)``
    where ``x8T`` is the [cols][rows] row-major transpose of the quantized
    input (the K-contiguous operand orientation for NT GEMMs)."""
@fp8_quantize_t.register_fake
 def _fp8_quantize_t_fake(x, scale, fmt):
    dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
    rows, cols = x.shape[-2], x.shape[-1]
    return (
        torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
        torch.empty(1, device=x.device, dtype=torch.float32),
    )
@fp8_quantize_t.register_kernel("cuda")
 def _fp8_quantize_t_cuda(x, scale, fmt):
    if x.dtype not in _QUANT_INPUT_DTYPES:
        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
    return get_module("fp8_ops").quantize(x, scale, int(fmt), 1)
@fp8_quantize_t.register_kernel("cpu")
 def _fp8_quantize_t_cpu(x, scale, fmt):
    x8, amax = _fp8_quantize_cpu(x, scale, fmt)
    return x8.transpose(-2, -1).contiguous(), amax
@custom_op("custom::fp8_quantize_dual", mutates_args=())
 def fp8_quantize_dual(
    x: torch.Tensor, scale: torch.Tensor, fmt: int
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Dual-orientation quantize: one read of ``x`` produces both the
    row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
    consumed by GEMMs on both orientations (backward ``g``)."""
@fp8_quantize_dual.register_fake
 def _fp8_quantize_dual_fake(x, scale, fmt):
    dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
    rows, cols = x.shape[-2], x.shape[-1]
    return (
        torch.empty(x.shape, device=x.device, dtype=dtype),
        torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
        torch.empty(1, device=x.device, dtype=torch.float32),
    )
@fp8_quantize_dual.register_kernel("cuda")
 def _fp8_quantize_dual_cuda(x, scale, fmt):
    if x.dtype not in _QUANT_INPUT_DTYPES:
        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
    return get_module("fp8_ops").quantize(x, scale, int(fmt), 2)
@fp8_quantize_dual.register_kernel("cpu")
 def _fp8_quantize_dual_cpu(x, scale, fmt):
    x8, amax = _fp8_quantize_cpu(x, scale, fmt)
    return x8, x8.transpose(-2, -1).contiguous(), amax
@fp8_quantize.register_kernel("cuda")
 def _fp8_quantize_cuda(x, scale, fmt):
    if x.dtype not in _QUANT_INPUT_DTYPES:
        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
    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(_fmt_name(fmt)))
    amax = x.abs().amax().float().reshape(1).clamp_min(1e-12)
    return x8, amax
@custom_op("custom::fp8_gemm", mutates_args=())
 def fp8_gemm(
    a: torch.Tensor,
    b: torch.Tensor,
    scale: torch.Tensor,
    trans_a: int = 0,
    trans_b: int = 0,
    bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
    """FP8 GEMM: ``a @ b * scale (+ bias)`` with FP32 accumulation.
    2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules).
    ``bias`` (bf16, length n) fuses into the epilogue in fp32 before the
    single bf16 rounding. The result is always BF16; FP8 output is a
    separate quantize operation.
    """
@fp8_gemm.register_fake
 def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0, bias=None):
    dtype = torch.bfloat16
    rows = a.size(2) if trans_a else a.size(1)
    cols = b.size(1) if trans_b else b.size(2)
    batches = [t.size(0) for t in (a, b) if t.dim() == 3]
    shape = (max(batches), rows, cols) if batches else (rows, cols)
    return torch.empty(shape, device=a.device, dtype=dtype)
@fp8_gemm.register_kernel("cuda")
 def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0, bias=None):
    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, scale, trans_a, trans_b, bias)
@fp8_gemm.register_kernel("cpu")
 def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
    aa = a.float().transpose(-2, -1) if trans_a else a.float()
    bb = b.float().transpose(-2, -1) if trans_b else b.float()
    acc = aa @ bb * scale
    if bias is not None and bias.numel() > 0:
        acc = acc + bias.float()
    return acc.to(torch.bfloat16)
 def quantize(
-    x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3", layout: int = 0
+    x: torch.Tensor,
-) -> tuple:
+    scale: torch.Tensor,
    fmt: str = "e4m3",
    transposed: bool = False,
    ring_state: Optional[torch.Tensor] = None,
    hist_idx: int = 0,
    fp8_max: float = 448.0,
    pow2_margin: float = 1.0,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax.
    ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
-    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. ``layout``
+    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
-    picks the output orientation: 0 = row-major ``(x8, amax)``; 1 =
+    ``transposed=True`` swaps ``x8`` for ``x8T``, the ``[cols][rows]``
-    transposed ``[cols][rows]`` ``(x8T, amax)`` — the K-contiguous operand
+    row-major transpose of the quantized input — the K-contiguous operand
-    orientation NT GEMMs want; 2 = both from one read ``(x8, x8T, amax)``
+    orientation NT GEMMs want — at the same 2-tuple arity.
-    (for tensors consumed in both orientations, e.g. backward ``g``).
+
    ``ring_state`` (a 1D float32 CUDA buffer laid out
    ``[hist n | scale | legacy | amax | done]``) switches on the in-kernel
    delayed-scaling fold: the kernel's last block folds the amax into
    ``hist[hist_idx]`` and publishes the next scale as
    ``max(hist) / fp8_max / pow2_margin`` — the returned ``amax`` is then the
    self-cleaned persistent slot (reads zero). None keeps the classic
    fresh-amax return.
    """
-    # Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a
+    return get_module("fp8_ops").quantize(
-    # 512-wide GEMM): real CUDA tensors of a supported dtype go straight to
+        x,
-    # the extension. Fake/subclass tensors and non-CUDA inputs keep the
+        scale,
-    # custom_op route so torch.compile / meta / fake-tensor tracing and the
+        _fmt_int(fmt),
-    # CPU fallback behave exactly as before.
+        transposed,
-    if (
+        ring_state,
-        type(x) is torch.Tensor
+        hist_idx,
-        and x.is_cuda
+        fp8_max,
-        and x.dtype in _QUANT_INPUT_DTYPES
+        pow2_margin,
-        and fmt in _FMT_TO_INT
+    )
-    ):
+
-        return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt], layout)
+
-    if layout == 0:
+def quantize_dual(
-        return fp8_quantize(x, scale, _fmt_int(fmt))
+    x: torch.Tensor,
-    if layout == 1:
+    scale: torch.Tensor,
-        return fp8_quantize_t(x, scale, _fmt_int(fmt))
+    fmt: str = "e4m3",
-    return fp8_quantize_dual(x, scale, _fmt_int(fmt))
+    ring_state: Optional[torch.Tensor] = None,
    hist_idx: int = 0,
    fp8_max: float = 448.0,
    pow2_margin: float = 1.0,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Dual-orientation quantize: one read of ``x`` produces both the
    row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
    consumed by GEMMs in both orientations (backward ``g``).
    ``ring_state`` switches on the in-kernel delayed-scaling fold exactly as
    in :func:`quantize`.
    """
    return get_module("fp8_ops").quantize_dual(
        x, scale, _fmt_int(fmt), ring_state, hist_idx, fp8_max, pow2_margin
    )
 def mm_fp8(
@@ -237,15 +113,4 @@ def mm_fp8(
    kernel epilogue in fp32 — no separate elementwise pass. The result is
    BF16; FP8 output is a separate quantize operation.
    """
-    # Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep
+    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
    # validation identical on the direct route (bias may be None — the
    # binding resolves it to the no-bias path).
    if (
        type(a) is torch.Tensor
        and a.is_cuda
        and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
    ):
        return get_module("fp8_ops").mm_fp8(
            a, b, scale, int(trans_a), int(trans_b), bias
        )
    return fp8_gemm(a, b, scale, trans_a, trans_b, bias)
@@ -54,19 +54,40 @@ struct Fp8GemmTraits {
                  "warp tile must be a multiple of the m16n8 MMA shape");
 };
 // Quantize output orientation: RowMajor = x8 only; Transposed = the
 // [cols][rows] x8T only; Dual = both from a single read. Transposed/Dual
 // produce K-contiguous operands so crosswise consumers (backward
 // grad_x / grad_w) route through the NT fast path.
 enum class QuantLayout : int {
    RowMajor = 0,
    Transposed = 1,
    Dual = 2,
 };
 // Quantize-kernel parameter POD: float input -> FP8 with fused amax.
 struct FP8QuantizeParams {
    const void* __restrict__ input_ptr = nullptr;
    void* __restrict__ output_ptr = nullptr;
    void* __restrict__ output_transposed_ptr = nullptr;  // [cols][rows]
-    // Output layout: 0 = row-major only, 1 = transposed only, 2 = both from
+    QuantLayout out_layout = QuantLayout::RowMajor;
    // a single read. Modes 1/2 produce K-contiguous operands so crosswise
    // consumers (backward grad_x / grad_w) route through the NT fast path.
    int out_layout = 0;
    const float* __restrict__ scale = nullptr;  // device multiplier
    float* __restrict__ amax = nullptr;         // raw-domain max out
    // Optional delayed-scaling ring fold: when fold_ring is set, the kernel's
    // last-finishing block folds the final amax into hist[hist_idx], reduces
    // the window and publishes the next scale — replacing the host-side
    // update chain. amax then points at a persistent self-cleaning slot
    // (zeroed by the same last block) inside the caller's ring state.
    bool fold_ring = false;
    float* __restrict__ hist = nullptr;  // [hist_len] amax history window
    float* __restrict__ scale_out = nullptr;
    unsigned int* __restrict__ done = nullptr;  // block-completion counter
    int hist_len = 0;
    int hist_idx = 0;
    float fp8_max = 448.0f;   // scale = max(hist) / fp8_max / pow2_margin
    float pow2_margin = 1.0f;
    // Element count (elementwise kernel); the tiled kernel views the same
    // buffer as [rows][cols] row-major.
    int total = 0;
@@ -1,4 +1,4 @@
-// CUDA bindings for the two stateless FP8 primitives.
+// CUDA bindings for the stateless FP8 quantize/GEMM primitives.
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -94,13 +94,17 @@ void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p,
        launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
 }
-}  // namespace
+// Shared binding body for the two quantize entry points: RowMajor /
-
+// Transposed (single output) serve quantize(), Dual (both orientations from
-// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return),
+// one read) serves quantize_dual(). A ring tensor switches
-// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a
+// on the in-kernel delayed-scaling fold: state layout
-// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path.
+// [hist n | scale | legacy | amax | done-as-int], and the returned amax is
-py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
+// the (self-cleaned) persistent slot. Without it, amax is reduced into a
-                    int64_t layout) {
+// fresh buffer armed by a driver memset — cheaper than the zeros() fill
 // kernel.
 py::object quantize_impl(torch::Tensor x, torch::Tensor scale, int64_t fmt,
                         QuantLayout layout, py::object ring, int64_t hist_idx,
                         double fp8_max, double pow2_margin) {
    TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == torch::kBFloat16 ||
                    x.scalar_type() == torch::kHalf ||
@@ -109,9 +113,7 @@ py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
    TORCH_CHECK(fmt == static_cast<int64_t>(FP8Format::E4M3) ||
                    fmt == static_cast<int64_t>(FP8Format::E5M2),
                "unsupported quantization type: expected E4M3 (0) or E5M2 (1)");
-    TORCH_CHECK(layout >= 0 && layout <= 2,
+    TORCH_CHECK(layout == QuantLayout::RowMajor || x.dim() >= 2,
                "layout must be 0 (row-major), 1 (transposed) or 2 (both)");
    TORCH_CHECK(layout == 0 || x.dim() >= 2,
                "transposed quantize layouts need a 2D+ tensor");
    check_scale(scale, x);
    check_fp8_device(x);
@@ -120,41 +122,94 @@ py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
    auto input = x.contiguous();
    auto out_opts = input.options().dtype(
        fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn);
-    // amax is reduced via atomicMax of non-negative values; a driver memset
+    torch::Tensor amax;
-    // arms it cheaper than the zeros() fill kernel (one fewer tensor-op
+    float *ring_hist = nullptr, *ring_scale_out = nullptr;
-    // dispatch + kernel launch on every quantize call).
+    unsigned int* ring_done = nullptr;
-    auto amax = torch::empty({1}, input.options().dtype(torch::kFloat32));
+    int ring_len = 0;
    if (!ring.is_none()) {
        auto st = ring.cast<torch::Tensor>();
        TORCH_CHECK(st.is_cuda() && st.dim() == 1 &&
                        st.scalar_type() == torch::kFloat32,
                    "ring state must be a 1D float32 CUDA tensor");
        const int64_t n = st.numel() - 4;
        TORCH_CHECK(n > 0 && hist_idx >= 0 && hist_idx < n,
                    "ring state too small or hist_idx out of range");
        float* base = st.data_ptr<float>();
        amax = st.narrow(0, n + 2, 1);
        ring_hist = base;
        ring_scale_out = base + n;
        ring_done = reinterpret_cast<unsigned int*>(base + n + 3);
        ring_len = static_cast<int>(n);
    } else {
        amax = torch::empty({1}, input.options().dtype(torch::kFloat32));
        cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream());
    }
    FP8QuantizeParams p;
    p.input_ptr = input.data_ptr();
    p.scale = scale.data_ptr<float>();
    p.amax = amax.data_ptr<float>();
    if (ring_hist) {
        p.fold_ring = true;
        p.hist = ring_hist;
        p.scale_out = ring_scale_out;
        p.done = ring_done;
        p.hist_len = ring_len;
        p.hist_idx = static_cast<int>(hist_idx);
        p.fp8_max = static_cast<float>(fp8_max);
        p.pow2_margin = static_cast<float>(pow2_margin);
    }
    p.total = static_cast<int>(input.numel());
-    p.out_layout = static_cast<int>(layout);
+    p.out_layout = layout;
    p.rows = static_cast<int>(input.size(-2));
    p.cols = static_cast<int>(input.size(-1));
    torch::Tensor output, output_t;
-    if (layout == 0 || layout == 2) {
+    if (layout != QuantLayout::Transposed) {
        output = torch::empty_like(input, out_opts);
        p.output_ptr = output.data_ptr();
    }
-    if (layout >= 1) {
+    if (layout != QuantLayout::RowMajor) {
        output_t = torch::empty({input.size(-1), input.size(-2)}, out_opts);
        p.output_transposed_ptr = output_t.data_ptr();
    }
    const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
-    if (layout != 0)
+    if (layout == QuantLayout::RowMajor)
        launch_quantize_for<true>(input, p, e5m2, stream.stream());
    else
        launch_quantize_for<false>(input, p, e5m2, stream.stream());
    else
        launch_quantize_for<true>(input, p, e5m2, stream.stream());
    C10_CUDA_CHECK(cudaGetLastError());
-    if (layout == 2) return py::make_tuple(output, output_t, amax);
+    if (layout == QuantLayout::Dual)
-    return py::make_tuple(layout == 1 ? output_t : output, amax);
+        return py::make_tuple(output, output_t, amax);
    return py::make_tuple(
        layout == QuantLayout::Transposed ? output_t : output, amax);
 }
 }  // namespace
 // Single-orientation quantize binding: row-major x8, or its [cols][rows]
 // transpose when transposed is set — the K-contiguous operand orientation
 // NT GEMMs want. Returns (x8|x8T, amax).
 py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
                    bool transposed, py::object ring, int64_t hist_idx,
                    double fp8_max, double pow2_margin) {
    const QuantLayout layout =
        transposed ? QuantLayout::Transposed : QuantLayout::RowMajor;
    return quantize_impl(x, scale, fmt, layout, ring, hist_idx, fp8_max,
                         pow2_margin);
 }
 // Dual-orientation quantize binding: one read of x produces both the
 // row-major x8 and its transpose (plus amax), for tensors consumed by GEMMs
 // in both orientations (backward g). Returns (x8, x8T, amax).
 py::object quantize_dual(torch::Tensor x, torch::Tensor scale, int64_t fmt,
                         py::object ring, int64_t hist_idx, double fp8_max,
                         double pow2_margin) {
    return quantize_impl(x, scale, fmt, QuantLayout::Dual, ring, hist_idx,
                         fp8_max, pow2_margin);
 }
 torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
-                     int64_t trans_a, int64_t trans_b, torch::Tensor bias) {
+                     bool trans_a, bool trans_b, py::object bias) {
    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,
@@ -164,6 +219,18 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
                    (b.dim() == 2 || b.dim() == 3),
                "a and b must be 2D or 3D (batched)");
    TORCH_CHECK(a.device() == b.device(), "a and b must share device");
    // Python None and an omitted argument both mean "no bias" — an undefined
    // tensor below. (py::isinstance<torch::Tensor> is false for real tensors
    // here — torch's caster registers no pybind type info — so validate by
    // attempting the cast itself.)
    torch::Tensor bias_t;
    if (!bias.is_none()) {
        try {
            bias_t = bias.cast<torch::Tensor>();
        } catch (const py::cast_error&) {
            TORCH_CHECK(false, "bias must be a torch.Tensor or None");
        }
    }
    check_scale(scale, a);
    check_fp8_device(a);
    const at::cuda::OptionalCUDAGuard guard(a.device());
@@ -180,10 +247,8 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
    torch::Tensor a_st, b_st;
    int64_t a_ld, b_ld, a_bstride, b_bstride;
-    const bool tag_a =
+    const bool tag_a = resolve_operand(a, trans_a, a_ld, a_bstride, a_st);
-        resolve_operand(a, trans_a != 0, a_ld, a_bstride, a_st);
+    const bool tag_b = resolve_operand(b, trans_b, b_ld, b_bstride, b_st);
    const bool tag_b =
        resolve_operand(b, trans_b != 0, b_ld, b_bstride, b_st);
    // GEMM dims from the user flags; storage layout never swaps them.
    const int64_t m = trans_a ? a.size(-1) : a.size(-2);
    const int64_t k = trans_a ? a.size(-2) : a.size(-1);
@@ -207,13 +272,13 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
    p.b_ld = static_cast<int>(b_ld);
    // Fused epilogue bias (bf16, broadcast over rows and batches). An
    // undefined or 0-element tensor keeps the plain scaled output.
-    if (bias.defined() && bias.numel() > 0) {
+    if (bias_t.defined() && bias_t.numel() > 0) {
-        TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kBFloat16,
+        TORCH_CHECK(bias_t.is_cuda() && bias_t.scalar_type() == torch::kBFloat16,
                    "fp8 gemm bias must be a CUDA bf16 tensor");
-        TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n,
+        TORCH_CHECK(bias_t.dim() == 1 && bias_t.size(0) == n,
                    "fp8 gemm bias must be 1D of length n=", n);
-        TORCH_CHECK(bias.is_contiguous(), "fp8 gemm bias must be contiguous");
+        TORCH_CHECK(bias_t.is_contiguous(), "fp8 gemm bias must be contiguous");
-        p.bias_ptr = bias.data_ptr();
+        p.bias_ptr = bias_t.data_ptr();
    }
    p.batch = static_cast<int>(batch);
    p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride;
@@ -227,28 +292,16 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
    return output;
 }
 // mm_fp8 binding: Python None and an omitted argument both mean "no bias",
 // so every Python layer can pass its bias argument through untouched.
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
-          py::arg("fmt"), py::arg("layout") = 0);
+          py::arg("fmt"), py::arg("transposed") = false,
-    m.def(
+          py::arg("ring") = py::none(), py::arg("hist_idx") = 0,
-        "mm_fp8",
+          py::arg("fp8_max") = 448.0, py::arg("pow2_margin") = 1.0);
-        [](torch::Tensor a, torch::Tensor b, torch::Tensor scale,
+    m.def("quantize_dual", &quantize_dual, py::arg("x"), py::arg("scale"),
-           int64_t trans_a, int64_t trans_b, py::object bias) {
+          py::arg("fmt"), py::arg("ring") = py::none(),
-            torch::Tensor t;
+          py::arg("hist_idx") = 0, py::arg("fp8_max") = 448.0,
-            if (!bias.is_none()) {
+          py::arg("pow2_margin") = 1.0);
-                // (py::isinstance<torch::Tensor> is false for real tensors
+    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
-                // here — torch's caster registers no pybind type info — so
+          py::arg("trans_a") = false, py::arg("trans_b") = false,
-                // validate by attempting the cast itself.)
+          py::arg("bias") = py::none());
                try {
                    t = bias.cast<torch::Tensor>();
                } catch (const py::cast_error&) {
                    TORCH_CHECK(false, "bias must be a torch.Tensor or None");
                }
            }
            return mm_fp8(a, b, scale, trans_a, trans_b, t);
        },
        py::arg("a"), py::arg("b"), py::arg("scale"), py::arg("trans_a") = 0,
        py::arg("trans_b") = 0, py::arg("bias") = py::none());
 }
@@ -108,8 +108,13 @@ __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
 // Block-wide amax reduce -> one atomic per block: warp-reduce, park one
 // value per warp, thread 0 folds. kWarps must cover the block's warp count.
 // With p.fold_ring, the last-finishing block additionally folds the final
 // amax into the history window and publishes the next scale (atomicAdd
 // ticket + fences), re-zeroing the amax slot and the counter for the next
 // launch — the host-side delayed-scaling update chain disappears.
 template <int kWarps>
-__device__ __forceinline__ void publish_amax(float* amax, float v) {
+__device__ __forceinline__ void publish_amax(const FP8QuantizeParams& p,
                                             float v) {
    v = warp_reduce_max(v);
    __shared__ float slots[kWarps];
    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
@@ -118,12 +123,23 @@ __device__ __forceinline__ void publish_amax(float* amax, float v) {
    if (tid == 0) {
 #pragma unroll
        for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
-        atomic_max_float(amax, v);
+        atomic_max_float(p.amax, v);
        if (!p.fold_ring) return;
        __threadfence();
        const unsigned int ticket = atomicAdd(p.done, 1u);
        __threadfence();
        if (ticket != gridDim.x - 1u) return;
        p.hist[p.hist_idx] = *p.amax;
        float peak = p.hist[0];
        for (int i = 1; i < p.hist_len; ++i) peak = fmaxf(peak, p.hist[i]);
        *p.scale_out = fmaxf(peak / p.fp8_max / p.pow2_margin, 1e-12f);
        *p.amax = 0.0f;
        *p.done = 0u;
    }
 }
-// Elementwise quantize kernel (out_layout 0): vectorized 16B loads -> fp8
+// Elementwise quantize kernel (QuantLayout::RowMajor): vectorized 16B loads
-// stores, fused amax over raw values.
+// -> fp8 stores, fused amax over raw values.
 template <FP8Format Fmt, typename InT>
 __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
    const float mult = *p.scale;
@@ -174,10 +190,11 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
        local_amax = fmaxf(local_amax, fabsf(v));
        x8[i] = cvt_fp8<Fmt>(v * mult);
    }
-    if (p.amax) publish_amax<8>(p.amax, local_amax);
+    if (p.amax) publish_amax<8>(p, local_amax);
 }
-// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] input
+// Tiled transpose quantize (QuantLayout::Transposed/Dual): reads the
 // [rows][cols] input
 // once and writes the fp8 bytes transposed ([cols][rows], so the contract
 // dim lands K-contiguous for NT GEMM operands) and, in mode 2, the row-major
 // copy too. 64x32 tiles, one native pair load per row (a full 128B warp
@@ -233,7 +250,7 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
            }
        }
    }
-    if (p.out_layout == 2) {
+    if (p.out_layout == QuantLayout::Dual) {
        uint8_t* out = static_cast<uint8_t*>(p.output_ptr);
 #pragma unroll
        for (int j = 0; j < 4; ++j)
@@ -266,11 +283,12 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
            out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
                tile[threadIdx.y * 8 + i][threadIdx.x];
    }
-    if (p.amax) publish_amax<8>(p.amax, local_amax);
+    if (p.amax) publish_amax<8>(p, local_amax);
 }
-// Unified quantize launcher: Tiled selects the transpose kernel (out_layout
+// Unified quantize launcher: Tiled selects the transpose kernel
-// 1/2) over the vectorized elementwise one. The transpose kernel vectorizes
+// (QuantLayout::Transposed/Dual) over the vectorized elementwise one. The
 // transpose kernel vectorizes
 // pair loads in-kernel and falls back to scalar loads at unaligned/ragged
 // rows, so the host side picks only the grid.
 template <FP8Format Fmt, typename InT, bool Tiled = false>
@@ -2,8 +2,9 @@
 The kernel-level tests exercise the two stateless primitives (``quantize`` for
 bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed
-operands); the policy-level tests (recipes, autocast context, per-tensor meta,
+operands); the policy-level tests (recipes, autocast context, per-tensor
-CPU fallbacks of the custom ops) run without a GPU.
+meta) run without a GPU. The primitives themselves are CUDA-only
 (attention-style direct wrappers — no torch.library dispatch layer).
 """
 import threading
@@ -24,7 +25,7 @@ from astrai.extension.fp8 import (
    fp8_linear_enabled,
    fp8_state,
 )
-from astrai.extension.ops.fp8 import mm_fp8, quantize
+from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
 from tests.conftest import skip_no_fp8
@@ -448,37 +449,72 @@ 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] packing: views alias the single state buffer.
+    # [hist | scale | legacy | amax | done] packing: views alias one buffer.
-    assert meta.w.state.numel() == 4 + 2
+    assert meta.w.state.numel() == 4 + 4
    assert meta.w.hist.data_ptr() == meta.w.state.data_ptr()
    assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr()
    meta.w.advance()
    assert meta.w.idx == 1
-    # update folds a fresh amax into the window and publishes the next scale
+    # fold_args hands the kernel the buffer, the slot and the recipe constants
-    amax = torch.tensor([8.0])
+    args = meta.w.fold_args("e4m3")
-    meta.w.update(amax, "e4m3")
+    assert args["ring_state"] is meta.w.state and args["hist_idx"] == 1
-    torch.testing.assert_close(meta.w.scale, torch.tensor([8.0 / 448.0]))
+    assert args["fp8_max"] == 448.0 and args["pow2_margin"] == 1.0
-def test_quantize_cpu_fallback():
+@skip_no_fp8
-    """CPU fallback of the quantize primitive (scale semantics + amax)."""
+@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"])
-    x = torch.randn(16, 32, dtype=torch.bfloat16)
+def test_quantize_dual_and_transposed_orientations(fmt):
-    scale = torch.tensor([0.5])  # quantize multiplier
+    """quantize_dual yields both orientations from one read; quantize's
-    x8, amax = quantize(x, scale, "e4m3")
+    transposed switch keeps the 2-tuple arity with the [cols][rows] layout."""
-    assert x8.dtype == torch.float8_e4m3fn
+    torch.manual_seed(11)
-    ref = (x.float() * 0.5).to(torch.float8_e4m3fn)
+    x = torch.randn(37, 67, device="cuda", dtype=torch.bfloat16) * 3
-    assert torch.equal(x8, ref)
+    mult = _scale(x).reciprocal()
    x8, amax = quantize(x, mult, fmt)
    x8T, _ = quantize(x, mult, fmt, transposed=True)
    d8, d8T, _ = quantize_dual(x, mult, fmt)
    assert x8T.shape == (67, 37)
    assert torch.equal(x8.view(torch.uint8), d8.view(torch.uint8))
    assert torch.equal(x8T.view(torch.uint8), d8T.view(torch.uint8))
    assert torch.equal(x8T.t().contiguous().view(torch.uint8), x8.view(torch.uint8))
    torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
-def test_mm_fp8_cpu_fallback():
+@skip_no_fp8
-    a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
+@pytest.mark.parametrize("fmt,fmax", [("e4m3", 448.0), ("e5m2", 57344.0)])
-    b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
+@pytest.mark.parametrize("margin", [0, 1])
-    scale = torch.tensor([1.0])
+def test_quantize_ring_fold_matches_host_update(fmt, fmax, margin):
-    out = mm_fp8(a8, b8, scale)
+    """The in-kernel delayed-scaling fold matches a host-side reference."""
-    ref = (a8.float() @ b8.float() * 1.0).to(torch.bfloat16)
+    dev = torch.device("cuda")
-    torch.testing.assert_close(out, ref)
+    n, idx = 4, 2
    torch.manual_seed(3)
    x = torch.randn(128, 96, dtype=torch.bfloat16, device=dev) * 3
    mult = torch.tensor([0.01], device=dev)
    pow2m = float(2**margin)
    # Reference: legacy quantize + the host fold it used to return amax for.
    x8_ref, amax = quantize(x, mult, fmt)
    hist = torch.full((n,), 1.0, device=dev)
    hist[idx] = amax.to(torch.float32)
    scale = (hist.max() / fmax / pow2m).clamp_min(1e-12).reshape(1)
    # Fused: same window, fold inside the quantize kernel's last block.
    ring = torch.zeros(n + 4, device=dev)
    ring[:n].fill_(1.0)
    x8, _ = quantize(
        x,
        mult,
        fmt,
        ring_state=ring,
        hist_idx=idx,
        fp8_max=fmax,
        pow2_margin=pow2m,
    )
    assert torch.equal(x8.view(torch.uint8), x8_ref.view(torch.uint8))
    torch.testing.assert_close(ring[:n], hist, rtol=0, atol=0)
    torch.testing.assert_close(ring[n : n + 1], scale, rtol=0, atol=0)
    assert float(ring[n + 2]) == 0.0  # amax slot self-cleaned
    assert int(ring[n + 3].view(torch.int32)) == 0  # done counter reset
 # --------------------------------------------------------------------------