refactor: collapse fp8 recipe hierarchy and state property layers

- Merge DelayedScaling/DynamicScaling and the abstract FP8Recipe base into one FP8Recipe dataclass with a dynamic flag; dispatch now reads cfg.recipe.dynamic instead of isinstance checks
- Drop the _ActiveOrDefault descriptor and the FP8State property views; the persistent defaults are plain default_* attributes and get_weight_meta takes the active recipe explicitly
- Convert FP8TensorMeta to a NamedTuple of the three per-operand rings
- Update tests to the new API; the autocast context test now asserts _active_config push/restore directly

Dockerfile

@@ -57,10 +57,14 @@ COPY docs/ ./docs/
 COPY pyproject.toml .
 COPY README.md .
 
-# Create non-root user matching the host uid/gid (passed via build args)
+# Create non-root user matching the host uid/gid (passed via build args).
+# ubuntu:24.04 ships a default 'ubuntu' user/group at uid/gid 1000, so remove
+# it first to free those ids before creating astrai.
 ARG USER_UID=1000
 ARG USER_GID=1000
-RUN groupadd -g "${USER_GID}" astrai \
+RUN userdel -r ubuntu 2>/dev/null || true \
+    && groupdel ubuntu 2>/dev/null || true \
+    && groupadd -g "${USER_GID}" astrai \
     && useradd -m -u "${USER_UID}" -g astrai astrai \
     && chown -R astrai:astrai /app
 ENV HOME=/home/astrai

astrai/extension/backend/attention.py

@@ -694,12 +694,13 @@ class CudaBackend(AttentionBackend):
 class FlashAttnBackend(AttentionBackend):
     """FlashAttention backend via the optional ``flash-attn`` package.
 
-    Decode (q_len=1, contiguous cache): uses ``flash_attn_with_kvcache``,
+    Decode (q_len=1, contiguous cache): writes K/V to the pool, gathers
-    which reads K/V directly from the flat pool via cache_batch_idx +
+    flat K/V via the ``req_to_token`` page table, and calls
-    cache_seqlens — no materialized KV gather.
+    ``flash_attn_varlen_func`` over the ragged batch
+    (``qo_indptr``/``kv_indptr``).
 
-    Prefill / non-contiguous decode: falls back to KV gather +
+    Prefill: packed 3-D calls share the ``flash_attn_varlen_func`` path;
-    ``flash_attn_func``.
+    dense 4-D calls go through ``flash_attn_func`` (mask-free only).
     """
 
     @classmethod

astrai/extension/fp8.py

@@ -31,15 +31,12 @@ import functools
 from contextvars import ContextVar, Token
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Dict, List, NamedTuple, Optional
 
 import torch
 from torch.library import Library
 
-from astrai.extension.ops.fp8 import (
+from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
-    linear_backward_fp8,
-    linear_forward_fp8,
-)
 
 # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
 FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -59,47 +56,35 @@ class FP8Format(str, Enum):
         return "e5m2" if self is FP8Format.HYBRID else self.value
 
 
+@dataclass
 class FP8Recipe:
     """Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
 
-    ``scale_from_history`` receives the operand's amax tensor (a ring window for
+    ``dynamic=False`` (default) is TE-style delayed scaling: max over the
-    delayed scaling, the current amax for dynamic scaling) and returns the
+    amax history window (amax from *previous* steps; the window trades
-    quantization step. Subclasses set ``history_len`` / ``margin``.
+    responsiveness against stability). ``dynamic=True`` is current-amax
+    scaling (torchao DYNAMIC): measure, then quantize — no history, at an
+    extra pass. ``scale_from_history`` receives the operand's amax tensor
+    (a ring window / the current amax) and returns the quantization step.
     """
 
     history_len: int = 16
     margin: int = 0
+    dynamic: bool = False
 
     def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
         peak = amax.max()
         return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
 
 
-@dataclass
-class DelayedScaling(FP8Recipe):
-    """TE-style delayed scaling: max over the amax history window (amax from
-    *previous* steps; the window trades responsiveness against stability)."""
-
-    history_len: int = 16
-    margin: int = 0
-
-
-@dataclass
-class DynamicScaling(FP8Recipe):
-    """Current-amax scaling (torchao DYNAMIC): measure, then quantize. No
-    history — the scale is derived from the same-step amax, at an extra pass."""
-
-    history_len: int = 1
-    margin: int = 0
-
-
 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 | legacy | amax | done]`` (views). The quantize
-    block records the measured amax into ``hist[idx]``, reduces the window and
+    kernel folds its fused amax into ``hist[idx]`` and publishes the next
-    publishes the next scale entirely on device — the Python-side write/max/write
+    scale from the window in its own last block (``fold_args`` passes the
-    chain is gone. The counter slot stays int32-zero (float bits) between
+    buffer + recipe constants); ``idx`` advances host-side each use. The
-    launches; ``idx`` advances host-side each step.
+    ``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")
@@ -107,14 +92,14 @@ 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
         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,20 +108,25 @@ class _ScaleRing:
         self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
         self.initialized = True
 
+    def fold_args(self, fmt: str) -> dict:
+        """Keyword arguments for quantize()'s in-kernel history fold."""
+        return {
+            "ring_state": self.state,
+            "hist_idx": self.idx,
+            "fp8_max": FP8_MAX[fmt],
+            "pow2_margin": float(2**self.recipe.margin),
+        }
 
-class FP8TensorMeta:
+
-    """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/
+class FP8TensorMeta(NamedTuple):
-    ``g``). Fused kernels record amax while quantizing, so the scale used at step
+    """Per-weight delayed-scaling rings for ``w``, ``x`` and ``g``.
-    N reflects amax from steps < N. DynamicScaling never allocates a meta — it
+
-    measures the current amax inline.
+    Dynamic scaling never allocates a meta; it measures the current amax inline.
     """
 
-    __slots__ = ("w", "x", "g")
+    w: _ScaleRing
-
+    x: _ScaleRing
-    def __init__(self, device: torch.device, recipe: FP8Recipe):
+    g: _ScaleRing
-        self.w = _ScaleRing(device, recipe)
-        self.x = _ScaleRing(device, recipe)
-        self.g = _ScaleRing(device, recipe)
 
 
 @dataclass(frozen=True)
@@ -160,60 +150,38 @@ _active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
 class FP8State:
     """Global fp8 training state: per-tensor metas + out-of-region defaults.
 
-    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set
+    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar``
-    by ``fp8_autocast``. The properties below read that active config when a
+    set by ``fp8_autocast`` (see ``_active``/``_current_config``); these plain
-    region is open and the global defaults otherwise; the setters (and
+    attributes are the persistent defaults applied outside any region —
-    ``fp8_linear_enable``) write the global defaults — the persistent switch
+    ``fp8_linear_enable`` writes ``default_enabled``. The metas registry is
-    applying outside any region. The metas registry is shared across threads
+    shared across threads (GIL-protected); fp8 backward runs on autograd
-    (GIL-protected); fp8 backward runs on autograd engine threads and only
+    engine threads and only touches metas captured on ``ctx`` at forward time.
-    touches metas captured on ``ctx`` at forward time.
     """
 
     def __init__(self):
         self.default_enabled = False
-        self.default_recipe: FP8Recipe = DelayedScaling()
+        self.default_recipe: FP8Recipe = FP8Recipe()
         self.default_format: FP8Format = FP8Format.HYBRID
         self._metas: Dict[tuple, FP8TensorMeta] = {}
 
-    # Active-config views (region config if open, else the defaults).
+    def get_weight_meta(self, w: torch.Tensor, recipe: FP8Recipe) -> FP8TensorMeta:
-    @property
-    def enabled(self) -> bool:
-        cfg = _active_config.get()
-        return cfg.enabled if cfg is not None else self.default_enabled
-
-    @property
-    def recipe(self) -> FP8Recipe:
-        cfg = _active_config.get()
-        return cfg.recipe if cfg is not None else self.default_recipe
-
-    @property
-    def fp8_format(self) -> FP8Format:
-        cfg = _active_config.get()
-        return cfg.fp8_format if cfg is not None else self.default_format
-
-    # Persistent (out-of-region) defaults.
-    @enabled.setter
-    def enabled(self, value: bool) -> None:
-        self.default_enabled = bool(value)
-
-    @recipe.setter
-    def recipe(self, value: FP8Recipe) -> None:
-        self.default_recipe = value
-
-    @fp8_format.setter
-    def fp8_format(self, value: FP8Format) -> None:
-        self.default_format = FP8Format(value)
-
-    def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
         key = (w.data_ptr(), w.shape, w.dtype)
         meta = self._metas.get(key)
         if meta is None:
-            meta = FP8TensorMeta(w.device, self.recipe)
+            meta = FP8TensorMeta(
+                _ScaleRing(w.device, recipe),
+                _ScaleRing(w.device, recipe),
+                _ScaleRing(w.device, recipe),
+            )
             self._metas[key] = meta
         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 = FP8Recipe()
+        self.default_format = FP8Format.HYBRID
         self._metas.clear()
 
 
@@ -275,7 +243,7 @@ class fp8_autocast:
         margin: int = 0,
     ):
         if recipe is None:
-            recipe = DelayedScaling(history_len=update_interval, margin=margin)
+            recipe = FP8Recipe(history_len=update_interval, margin=margin)
         self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
         self._tokens: List[Token] = []
 
@@ -307,18 +275,9 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor
     return recipe.scale_from_history(amax, fmt)
 
 
-_zero_bias: Dict[Optional[int], torch.Tensor] = {}
+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)
-def _empty_bias(x: torch.Tensor) -> torch.Tensor:
-    """Per-device cached 0-element bf16 bias (the binding only checks numel —
-    never mutated), saving a CUDA allocation per bias-less linear."""
-    key = x.device.index
-    t = _zero_bias.get(key)
-    if t is None:
-        t = torch.empty(0, device=x.device, dtype=torch.bfloat16)
-        _zero_bias[key] = t
-    return t
 
 
 def fp8_linear_forward(
@@ -326,60 +285,50 @@ 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 with the bias fused into its epilogue.
-    kernels (amax folded into the window, next scale published on device);
+    Delayed scaling lets the quantize kernel fold the fused amax into the
-    dynamic scaling measures the current amax itself. Training quantizes the
+    history ring and publish the next scale in its own last block; dynamic
-    weight every step (the optimizer bumps its version, so there is no cast
+    scaling measures the current amax itself. Training quantizes the weight
-    cache, matching ``cached_cast``-less behavior).
+    every step (the optimizer bumps its version, so there is no cast cache,
+    matching ``cached_cast``-less behavior).
     """
     state = fp8_state()
     if cfg is None:
         cfg = _current_config()
     fmt = cfg.fp8_format.fwd()
-    margin = cfg.recipe.margin
+    if cfg.recipe.dynamic:
-    if bias is None:
-        bias = _empty_bias(x)
-    if isinstance(cfg.recipe, DynamicScaling):  # measure-then-quantize, no state
         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)
+        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
+        # Bias fuses into the GEMM epilogue (fp32 add before the single bf16
+        # rounding — one rounding fewer than the separate out + bias pass);
+        # None passes through to the kernel's no-bias path.
+        out = mm_fp8(
+            x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
+        ).reshape(*x.shape[:-1], w.size(0))
         return out, sx, sw
 
-    meta = state.get_weight_meta(w)
+    meta = state.get_weight_meta(w, cfg.recipe)
     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).
+    # The clones feed this call's kernels (stream-ordered before the in-kernel
-    # Snapshot the scales first so the GEMM dequantizes with the SAME scale the
+    # fold overwrites the ring scale slots); the fp8 quantize kernel folds the
-    # operands were quantized with, and so the backward can reuse this step's
+    # amax into the history window and publishes the next scale itself.
-    # scale (gradient consistency with the forward). The ring finalize still
+    x8, _ = quantize(x, sx.reciprocal(), fmt, **meta.x.fold_args(fmt))
-    # publishes the next step's scale into the original slot.
+    if _is_fp8(w.dtype):
-    if w.dtype is not torch.bfloat16:  # static pre-quantized weight
+        w8 = w
-        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, _ = quantize(w, sw.reciprocal(), fmt, **meta.w.fold_args(fmt))
-    sx = meta.x.scale.clone()
+    out = mm_fp8(
-    sw = sw_arg.clone()
+        x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
-    out, _x8, _w8, _ax, _aw = linear_forward_fp8(
+    ).reshape(*x.shape[:-1], w.size(0))
-        x,
-        w_arg,
-        bias,
-        sx,
-        sw,
-        fmt,
-        None,
-        meta.x.state,
-        meta.x.idx,
-        margin,
-        w_ring,
-        meta.w.idx,
-        margin,
-    )
     meta.x.advance()
-    if w_ring is not None:
+    if not _is_fp8(w.dtype):
         meta.w.advance()
     return out, sx, sw
 
@@ -401,8 +350,8 @@ class _LinearFp8(torch.autograd.Function):
         ctx.save_for_backward(x, w, sx, sw)
         ctx.fmt_bwd = cfg.fp8_format.bwd()
         ctx.recipe = cfg.recipe
-        ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
+        ctx.is_dynamic = cfg.recipe.dynamic
-        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w)
+        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w, cfg.recipe)
         return out
 
     @staticmethod
@@ -410,38 +359,45 @@ 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.
+        # Flatten leading dims (the forward GEMMs ran on [-1, N] / [-1, K]
+        # views; the kernels only accept 2D operands).
+        g2 = g.reshape(-1, g.size(-1))
         if ctx.is_dynamic:
-            sg = _dynamic_scale(g, ctx.recipe, fmt)
+            sg = _dynamic_scale(g2, 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)
+                meta.g.seed(g2, 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(
+        # Backward GEMMs route through the NT fast path via transposed
-            g,
+        # quantize outputs: g8 [m,n] with w8T [k,n] (trans_b=True) gives
-            x,
+        # grad_x, g8T [n,m] with x8T [k,m] gives grad_w — no NN-swap or TT
-            w,
+        # crosswise kernel in the training path. g is consumed in both
-            list(ctx.needs_input_grad),
+        # orientations, so quantize_dual's single pass feeds both.
-            sg,
+        # The g quantize folds the gradient amax into its ring in-kernel;
-            sw,
+        # the x8T/w8T orientation copies discard amax (those rings were
-            sx,
+        # folded at forward time).
-            fmt,
+        g8, g8T, _ = quantize_dual(g2, sg.reciprocal(), fmt, **meta.g.fold_args(fmt))
-            ring,
+        x8T, _ = quantize(
-            idx,
+            x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, transposed=True
-            ctx.recipe.margin,
         )
+        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, 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.advance()  # the g quantize kernel finalized the ring in-kernel
+            meta.g.advance()
-        return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
+        return grad_x, grad_w, grad_b
 
 
 # ---------------------------------------------------------------------------

astrai/extension/ops/fp8.py

@@ -1,25 +1,28 @@
 """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_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → 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)
-- ``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
+``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.
 """
 
-from typing import List, Optional, Tuple
+from typing import Optional, Tuple
 
 import torch
-from torch.library import custom_op
 
 from astrai.extension.loader import get_module
 
@@ -34,212 +37,80 @@ def _fmt_int(fmt: str) -> int:
         raise ValueError(f"unsupported fp8 format {fmt!r} (expected 'e4m3' or 'e5m2')")
 
 
-def _fmt_dtype(fmt: str) -> torch.dtype:
+def quantize(
-    return torch.float8_e5m2 if _fmt_int(fmt) else torch.float8_e4m3fn
+    x: torch.Tensor,
-
+    scale: torch.Tensor,
-
+    fmt: str = "e4m3",
-@custom_op("custom::fp8_quantize", mutates_args=())
+    transposed: bool = False,
-def fp8_quantize(
+    ring_state: Optional[torch.Tensor] = None,
-    x: torch.Tensor, scale: torch.Tensor, fmt: int
+    hist_idx: int = 0,
+    fp8_max: float = 448.0,
+    pow2_margin: float = 1.0,
 ) -> 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 the quantization multiplier (device scalar); ``fmt`` selects
+    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
+    ``transposed=True`` swaps ``x8`` for ``x8T``, the ``[cols][rows]``
+    row-major transpose of the quantized input — the K-contiguous operand
+    orientation NT GEMMs want — at the same 2-tuple arity.
 
-@fp8_quantize.register_fake
+    ``ring_state`` (a 1D float32 CUDA buffer laid out
-def _fp8_quantize_fake(x, scale, fmt):
+    ``[hist n | scale | legacy | amax | done]``) switches on the in-kernel
-    dtype = torch.float8_e5m2 if fmt else torch.float8_e4m3fn
+    delayed-scaling fold: the kernel's last block folds the amax into
-    return (
+    ``hist[hist_idx]`` and publishes the next scale as
-        torch.empty(x.shape, device=x.device, dtype=dtype),
+    ``max(hist) / fp8_max / pow2_margin`` — the returned ``amax`` is then the
-        torch.empty(1, device=x.device, dtype=torch.float32),
+    self-cleaned persistent slot (reads zero). None keeps the classic
+    fresh-amax return.
+    """
+    return get_module("fp8_ops").quantize(
+        x,
+        scale,
+        _fmt_int(fmt),
+        transposed,
+        ring_state,
+        hist_idx,
+        fp8_max,
+        pow2_margin,
     )
 
 
-@fp8_quantize.register_kernel("cuda")
+def quantize_dual(
-def _fp8_quantize_cuda(x, scale, fmt):
+    x: torch.Tensor,
-    if x.dtype != torch.bfloat16:
+    scale: torch.Tensor,
-        raise TypeError(f"fp8 quantize requires bf16 input, got {x.dtype}")
+    fmt: str = "e4m3",
-    return get_module("fp8_ops").quantize_bf16(x, scale, 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
-@fp8_quantize.register_kernel("cpu")
+    in :func:`quantize`.
-def _fp8_quantize_cpu(x, scale, fmt):
-    x8 = (x.float() / scale).to(_fmt_dtype("e5m2" if fmt else "e4m3"))
-    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,
-    sa: torch.Tensor,
-    sb: torch.Tensor,
-    out_dtype: int = 0,
-    out_scale: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation.
-
-    ``out_dtype``: 0 = BF16 (default), 1 = FP8 E4M3 (requires ``out_scale``,
-    the quantization step for the output — mirrors ``torch._scaled_mm``).
     """
-
+    return get_module("fp8_ops").quantize_dual(
-
+        x, scale, _fmt_int(fmt), ring_state, hist_idx, fp8_max, pow2_margin
-@fp8_gemm.register_fake
+    )
-def _fp8_gemm_fake(a, b, sa, sb, out_dtype=0, out_scale=None):
-    dtype = torch.float8_e4m3fn if out_dtype else torch.bfloat16
-    return torch.empty((a.size(0), 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):
-    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)
-
-
-@fp8_gemm.register_kernel("cpu")
-def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None):
-    acc = a.float() @ b.float() * sa * sb
-    if out_dtype:
-        os_ = 1.0 if out_scale is None else out_scale
-        return (acc * os_).to(torch.float8_e4m3fn)
-    return acc.to(torch.bfloat16)
-
-
-def quantize_bf16(
-    x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``.
-
-    ``scale`` is the quantization step (device scalar); ``fmt`` selects
-    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor — the caller
-    never clears it.
-    """
-    return fp8_quantize(x, scale, _fmt_int(fmt))
 
 
 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,
+    bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``.
+    """Pre-quantized FP8 GEMM: ``a @ b * scale (+ bias)``.
 
-    ``a``/``b`` must be FP8 tensors of the same format (E4M3 or E5M2);
+    ``a``/``b`` must be FP8 tensors of the same format, 2D or 3D (batched,
-    ``sa``/``sb`` are their quantization steps. ``out_dtype`` is ``"bf16"``
+    matmul-style broadcast on the batch dim). Inner-transposed views (e.g.
-    (default) or ``"e4m3"`` — FP8 output for layer-to-layer pipelines, which
+    ``x.t()``) fold into the layout at zero copy. ``scale`` is their combined
-    requires ``out_scale`` (the output quantization step).
+    dequantization scale. ``bias`` (CUDA bf16 1D of length n) adds inside the
+    kernel epilogue in fp32 — no separate elementwise pass. The result is
+    BF16; FP8 output is a separate quantize operation.
     """
-    if out_dtype not in ("bf16", "e4m3"):
+    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
-        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,
-    )

astrai/inference/runtime/executor.py

@@ -67,11 +67,15 @@ class DecodeSteadyState:
 
     When the same ordered task set decodes one token per step, sampling
     params and task signature are reused; only positions advance by 1.
+    ``last_tokens`` keeps that step's sampled ids on-device so the next
+    step with an unchanged signature can fill ``input_ids`` via a
+    device-to-device copy.
     """
 
     task_sig: tuple
     positions: list[int]
     sampling_info: SamplingBatchInfo
+    last_tokens: Optional[Tensor] = None
 
 
 def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
@@ -250,6 +254,13 @@ class Executor:
         return_logprobs: bool = False,
         info: Optional[SamplingBatchInfo] = None,
     ):
+        """Sample from ``logits`` and return ``(host_payload, tokens)``.
+
+        ``host_payload`` is the scheduler-facing list (token ids, or
+        ``(token_id, logprob)`` tuples with ``return_logprobs``);
+        ``tokens`` is the ``[B]`` device tensor that produced it, kept
+        for the steady-state decode fast path.
+        """
         info = info or _build_sampling_batch_info(tasks, self.device)
         if info.has_freq:
             history_lists = [
@@ -284,14 +295,14 @@ class Executor:
             return_logprobs=return_logprobs,
         )
         if not return_logprobs:
-            return result.tolist()
+            return result.tolist(), result
 
         tokens, logprobs = result
         tokens_list = tokens.tolist()
         logprobs_list = logprobs.tolist()
         for task, logprob in zip(tasks, logprobs_list):
             task.output_logprobs.append(float(logprob))
-        return list(zip(tokens_list, logprobs_list))
+        return list(zip(tokens_list, logprobs_list)), tokens
 
     def execute_prefill(
         self,
@@ -336,7 +347,8 @@ class Executor:
                 torch.arange(1, batch_sz + 1, device=self.device) * q_len - 1
             ]
 
-        return tasks, self._sample_logits(logits, tasks, return_logprobs)
+        step_out, _ = self._sample_logits(logits, tasks, return_logprobs)
+        return tasks, step_out
 
     def execute_decode(
         self, tasks: List[Task], return_logprobs: bool = False
@@ -360,24 +372,30 @@ class Executor:
 
         b = len(tasks)
         ws = self._workspace
+        task_ids = [t.task_id for t in tasks]
+        cur_positions = [t.next_pos for t in tasks]
+        task_sig = tuple(task_ids)
 
         # ---- pre-replay: update input buffers in-place ----
 
-        input_ids = ws.fill_input_ids(
+        # When the previous decode step ran this same ordered task set, its
-            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
+        # sampled tokens are still on-device and map 1:1 onto the current
-        )
+        # slots — fill input ids device-to-device.  inference_mode guards
-
+        # the read because the source was produced under sampling's
-        task_ids = [t.task_id for t in tasks]
+        # inference-mode context.
-        cur_positions = [t.next_pos for t in tasks]
+        cached = self._decode_cache
+        sig_match = cached is not None and cached.task_sig == task_sig
+        if sig_match and cached.last_tokens is not None:
+            with torch.inference_mode():
+                input_ids = ws.fill_input_ids_from_device(cached.last_tokens)
+        else:
+            input_ids = ws.fill_input_ids(
+                [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
+            )
 
         kv_cache = self.task_cache.bind(task_ids, ws)
 
-        task_sig = tuple(task_ids)
+        reuse_decode_state = self.task_cache.bind_was_steady and sig_match
-        reuse_decode_state = (
-            self.task_cache.bind_was_steady
-            and self._decode_cache is not None
-            and self._decode_cache.task_sig == task_sig
-        )
         if reuse_decode_state:
             info = self._decode_cache.sampling_info
             ws.position_ids[:b] += 1
@@ -418,4 +436,8 @@ class Executor:
                 )
             logits = outputs["logits"]
 
-        return self._sample_logits(logits, tasks, return_logprobs, info=info)
+        step_out, tokens_dev = self._sample_logits(
+            logits, tasks, return_logprobs, info=info
+        )
+        self._decode_cache.last_tokens = tokens_dev
+        return step_out

astrai/inference/workspace.py

@@ -139,6 +139,18 @@ class InferenceWorkspace:
         self.input_ids[:b].copy_(pin[:b])
         return self.input_ids[:b]
 
+    def fill_input_ids_from_device(self, tokens: Tensor) -> Tensor:
+        """Copy device-resident ``[B]`` token ids into the device buffer.
+
+        Steady-state decode fast path: when the executor's cached task
+        signature still matches, the previous step's sampled tokens map
+        1:1 onto the current slots, so the ids transfer device-to-device
+        instead of round-tripping through the host staging buffers.
+        """
+        b = tokens.size(0)
+        self.input_ids[:b].copy_(tokens)
+        return self.input_ids[:b]
+
     def decode_mask(self, position_ids: Tensor, total_len: int) -> Tensor:
         """Return the ``[B, 1, total_len]`` validity mask for this step.
 

astrai/logging.py

@@ -1,11 +1,13 @@
 import logging
 import os
 
+from astrai.parallel.setup import get_rank, get_world_size
+
 
 class _DistributedContextFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
-        record.rank = os.environ.get("RANK", "0")
+        record.rank = str(get_rank())
-        record.world_size = os.environ.get("WORLD_SIZE", "1")
+        record.world_size = str(get_world_size())
         return True
 
 

astrai/parallel/setup.py

@@ -30,15 +30,13 @@ def get_current_device():
 def get_world_size() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size()
-    else:
+    return int(os.environ.get("WORLD_SIZE", "1"))
-        return 1
 
 
 def get_rank() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_rank()
-    else:
+    return int(os.environ.get("RANK", "0"))
-        return 0
 
 
 @contextmanager

astrai/trainer/train_callback.py

@@ -116,6 +116,8 @@ class GradientCheckpointingCallback(TrainCallback):
             del module._original_forward
 
     def on_train_begin(self, context: TrainContext):
+        if not self.modules:
+            return
         context.model.apply(self._enable)
         logger.info("Gradient checkpointing enabled")
 

csrc/CMakeLists.txt

@@ -52,13 +52,16 @@ set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
 # globally unique across families) and their per-family source paths under
 # kernels/. `loader.py` auto-discovers the .so files in astrai/extension/lib/,
 # so this CMake registry is the single place to register a new kernel.
+#
+# FP8 MMA instructions require sm_89+. Keep the target out of the build on
+# older architectures instead of instantiating templates that cannot compile.
+# The remaining kernels are still useful on sm_80+ (including sm_86).
 set(KERNEL_NAMES
     attn_decode
     attn_prefill
     attn_paged_decode
     attn_paged_prefill
     rotary_emb
-    fp8_ops
 )
 set(KERNEL_SRCS
     attention/decode.cu
@@ -66,9 +69,17 @@ set(KERNEL_SRCS
     attention/paged_decode.cu
     attention/paged_prefill.cu
     rotary/rotary_emb.cu
-    fp8/ops.cu
 )
 
+if(ASTRAI_CUDA_ARCH GREATER_EQUAL 89)
+    list(APPEND KERNEL_NAMES fp8_ops)
+    list(APPEND KERNEL_SRCS fp8/ops.cu)
+else()
+    message(WARNING
+        "FP8 operator disabled: ASTRAI_CUDA_ARCH=${ASTRAI_CUDA_ARCH} "
+        "requires compute capability 89 or newer")
+endif()
+
 list(LENGTH KERNEL_NAMES _kernel_count)
 math(EXPR _kernel_last "${_kernel_count} - 1")
 foreach(i RANGE ${_kernel_last})

csrc/kernels/attention/common.h

@@ -16,6 +16,11 @@ enum TensorLayout : int {
 // Split-KV workspace cap: max decode splits per (batch, q_head).
 constexpr int MAX_SPLITS = 32;
 
+// Paged-prefill host Q-tile granularity in q rows: one q_tile_to_index unit
+// covers this many query rows of one request.  Must match Q_TILE_ROWS in
+// astrai/inference/workspace.py, which builds the device-side tile maps.
+constexpr int HOST_Q_TILE_ROWS = 64;
+
 
 // Unified attention params covering BOTH addressing modes:
 //   - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v).

csrc/kernels/attention/dispatchers.cuh

@@ -88,8 +88,17 @@ struct PrefillLauncherMMA {
     static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
         using Config = PrefillConfigMap<HEAD_DIM, IsCausal>;
         using Traits = KernelTraits<HEAD_DIM, Config::BC, Config::WARPS, Config::STAGES>;
-        constexpr int ROWS = Traits::BR * Config::WARPS;
+        // GQA head packing: HB = min(G, WARPS) q-heads of one kv-head group
-        dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
+        // share each block's K/V stream (~HB× less global K/V traffic).
+        // Each head gets WPH = WARPS/HB 16-row chunks per block, so per-head
+        // rows drop from 64 to BR*WPH while total mma work per K/V byte is
+        // unchanged.  G=1 (MHA) reproduces the historical grid exactly.
+        const int G = p.q_head / p.kv_head;
+        const int HB = std::min(G, Config::WARPS);
+        const int WPH = Config::WARPS / HB;
+        constexpr int BR = Traits::BR;
+        dim3 grid(QSchedule::packed_grid_x(p, BR * WPH),
+                  p.kv_head * ((G + HB - 1) / HB),
                   QSchedule::host_grid_batch(p));
         dim3 block(Traits::NUM_THREADS);
         attn_prefill_split_q_mma_kernel<Traits, QSchedule, KV, IsCausal, HasMask>

csrc/kernels/attention/layout_policies.cuh

@@ -56,6 +56,20 @@ struct DenseQSchedule {
         q_tile = blockIdx.x;
     }
 
+    // GQA-packed prefill mapping: HB q-heads of one kv-head group share a
+    // block's K/V stream, each head owning `rows` = BR*WPH consecutive q rows
+    // per block.  Dense tensors tile q_len directly, one block per range.
+    HOST_FORCEINLINE int packed_grid_x(
+        const AttentionParams<bf16>& p, int rows) {
+        return (p.q_len + rows - 1) / rows;
+    }
+
+    DEVICE_FORCEINLINE void map_packed_block(
+        const AttentionParams<bf16>&, int rows, int& batch, int& row_base) {
+        batch = blockIdx.z;
+        row_base = blockIdx.x * rows;
+    }
+
     DEVICE_FORCEINLINE int q_len(
         const AttentionParams<bf16>& p, int) {
         return p.q_len;
@@ -84,6 +98,23 @@ struct PackedQSchedule {
         q_tile = p.q_tile_to_index[blockIdx.x];
     }
 
+    // GQA-packed prefill mapping: the host tile maps are built in
+    // HOST_Q_TILE_ROWS granularity, so each host tile splits into
+    // HOST_Q_TILE_ROWS / rows packed blocks along blockIdx.x.
+    HOST_FORCEINLINE int packed_grid_x(
+        const AttentionParams<bf16>& p, int rows) {
+        return p.num_q_tiles * (HOST_Q_TILE_ROWS / rows);
+    }
+
+    DEVICE_FORCEINLINE void map_packed_block(
+        const AttentionParams<bf16>& p, int rows, int& batch, int& row_base) {
+        const int hb = HOST_Q_TILE_ROWS / rows;
+        const int host_tile = blockIdx.x / hb;
+        batch = p.q_tile_to_batch[host_tile];
+        row_base = p.q_tile_to_index[host_tile] * HOST_Q_TILE_ROWS
+            + (blockIdx.x - host_tile * hb) * rows;
+    }
+
     DEVICE_FORCEINLINE int q_len(
         const AttentionParams<bf16>& p, int batch) {
         return p.qo_indptr[batch + 1] - p.qo_indptr[batch];

csrc/kernels/attention/prefill_split_q_mma.cuh

@@ -13,6 +13,13 @@ namespace attention {
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
 // cores via mma.sync.m16n8k16 (f32 accumulate).
 //
+// GQA head packing (FA2/FA3-style): HB = min(G, WARPS) query heads of one
+// kv-head group share a block's K/V tiles, so each K/V element is read from
+// global memory once per block instead of once per q head (~HB× less K/V
+// traffic).  WARPS = WPH × HB: warp w handles head slot w/WPH, chunk w%WPH;
+// all warps of a block cover the same token range, keeping the causal sweep
+// end block-uniform.  G=1 (MHA) degenerates to the unpadded layout.
+//
 // KV = ContigKV (dense [batch, kv_head, kv_len, head_dim]) or PagedKV
 //      (flat pool + req_to_token, ragged batches via qo_indptr/kv_indptr).
 // IsCausal and HasMask are compile-time bools — the compiler eliminates all
@@ -26,11 +33,23 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int gid = lane >> 2;   // 0..7
     const int tid4 = lane & 3;   // 0..3
 
-    const int q_head = blockIdx.y;
+    const int G = p.q_head / p.kv_head;
-    int batch, q_tile;
+    const int HB = min(G, Traits::WARPS);   // q heads packed per block
-    QSchedule::map_block(p, batch, q_tile);
+    const int WPH = Traits::WARPS / HB;     // 16-row chunks per head
-    const int kv_head = q_head / (p.q_head / p.kv_head);
+    const int BPG = (G + HB - 1) / HB;      // blocks per GQA group
-    const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;
+    const int chunk = warp % WPH;
+
+    int batch, row_base;
+    QSchedule::map_packed_block(p, Traits::BR * WPH, batch, row_base);
+    const int kv_head = blockIdx.y / BPG;
+    const int slot = blockIdx.y - kv_head * BPG;
+    const int head_idx = slot * HB + warp / WPH;
+    // G % HB tail blocks have idle head slots: clamp to the last head so all
+    // warps do valid work (cp.async + __syncthreads stay block-uniform) and
+    // just skip the O store via `active`.
+    const bool active = head_idx < G;
+    const int q_head = kv_head * G + min(head_idx, G - 1);
+    const int qrow0 = row_base + chunk * Traits::BR;
 
     // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
     const int seq_len    = KV::kv_len(p, batch);
@@ -62,11 +81,11 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int qr0 = qrow0 + gid;
     const int qr1 = qrow0 + gid + 8;
 
-    // Causal tile-skip bounds (dead code when IsCausal == false)
+    // Causal tile-skip bounds (dead code when IsCausal == false).
+    // max_kv is per-warp (its own 16 rows); block_max_kv is the last row of
+    // the whole block's range and must be uniform for the shared sweep loop.
     const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
-    const int block_max_kv =
+    const int block_max_kv = row_base + WPH * Traits::BR - 1 + causal_off;
-        q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
-        + causal_off;
 
     int t_end = tiles - 1;
     if constexpr (IsCausal) {
@@ -144,13 +163,13 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     #pragma unroll
     for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
         int d = dn8 * 8 + 2 * tid4;
-        if (qr0 < q_len) {
+        if (active && qr0 < q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
                                                       Oacc[dn8][1] * rl0);
             *reinterpret_cast<__nv_bfloat162*>(
                 &p.o_ptr[o_base + qr0 * p.q_l_stride + d * p.q_d_stride]) = v;
         }
-        if (qr1 < q_len) {
+        if (active && qr1 < q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
                                                      Oacc[dn8][3] * rl1);
             *reinterpret_cast<__nv_bfloat162*>(

csrc/kernels/common/cp_async.cuh

@@ -2,8 +2,13 @@
 //
 // One header for the async-copy pipeline used by both the attention kernels
 // (predicated 16-byte K/V tile staging) and the fp8 GEMM (predicated operand
-// staging + wait_group dispatch). PTX requires wait_group's operand to be an
+// staging + the fixed-depth wait_group). The emitter is split from its
-// immediate, hence the template forms.
+// policies: cp_async_16_raw owns the single PTX site, and each wrapper states
+// one destination contract (generic pointer vs loop-carried shared offset)
+// and one predication contract (unconditional vs zero-fill-when-false), so
+// call sites never pass a dead `true` predicate or re-convert a carried
+// offset. PTX requires wait_group's operand to be an immediate, hence the
+// template form below.
 
 #pragma once
 
@@ -11,15 +16,14 @@
 
 namespace astrai {
 
-// Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill.
+// Raw emitter: read src_size bytes (<= 16) from gmem into the shared
-// src_size=0 means no bytes are read, so an out-of-bounds address is safe.
+// offset. src_size = 0 reads nothing, so a predicated-off call zero-fills
-// BypassL1 defaults to .cg (L2 only); false selects .ca (L1 + L2).
+// its destination without touching the (possibly out-of-range) source.
-// `T` is the smem element type; only the destination pointer's type matters.
+// BypassL1 selects .cg (L2 only, default) vs .ca (L1 + L2).
-template <typename T, bool BypassL1 = true>
+template <bool BypassL1 = true>
-__device__ __forceinline__ void cp_async_16(T* smem_ptr, const void* gmem_ptr,
+__device__ __forceinline__ void cp_async_16_raw(unsigned smem_addr,
-                                            bool pred) {
+                                                const void* gmem_ptr,
-    const unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
+                                                int src_size) {
-    const int src_size = pred ? 16 : 0;
     if constexpr (BypassL1) {
         asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
                      :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
@@ -29,6 +33,32 @@ __device__ __forceinline__ void cp_async_16(T* smem_ptr, const void* gmem_ptr,
     }
 }
 
+// Unconditional 16-byte copy to a generic shared pointer.
+// `T` is the smem element type; only the destination pointer's type matters.
+template <typename T, bool BypassL1 = true>
+__device__ __forceinline__ void cp_async_16(T* smem_ptr,
+                                            const void* gmem_ptr) {
+    cp_async_16_raw<BypassL1>(__cvta_generic_to_shared(smem_ptr), gmem_ptr,
+                              16);
+}
+
+// Predicated: full copy when `pred`, zero-fill otherwise.
+template <typename T, bool BypassL1 = true>
+__device__ __forceinline__ void cp_async_16(T* smem_ptr, const void* gmem_ptr,
+                                            bool pred) {
+    cp_async_16_raw<BypassL1>(__cvta_generic_to_shared(smem_ptr), gmem_ptr,
+                              pred ? 16 : 0);
+}
+
+// Predicated raw-offset form: the destination is an already-converted
+// shared-memory offset (e.g. a loop-carried swizzled stage address), so
+// steady-state prefetch sites issue one LDGSTS straight from the register.
+template <bool BypassL1 = true>
+__device__ __forceinline__ void cp_async_16(unsigned smem_addr,
+                                            const void* gmem_ptr, bool pred) {
+    cp_async_16_raw<BypassL1>(smem_addr, gmem_ptr, pred ? 16 : 0);
+}
+
 // Commit all outstanding cp.async ops of this thread as one group.
 __device__ __forceinline__ void cp_async_commit_group() {
     asm volatile("cp.async.commit_group;");
@@ -49,21 +79,4 @@ __device__ __forceinline__ void cp_async_wait_group() {
     asm volatile("cp.async.wait_group %0;" :: "n"(KeepGroups));
 }
 
-// Runtime dispatch over cp_async_wait_group<N>: unrolls into a compare
-// ladder over [0, MaxKeepGroups] so the immediate-only PTX constraint is
-// hidden behind a runtime `keep_groups` (used by the fp8 GEMM pipeline,
-// whose remaining-tile count is dynamic).
-template <int MaxKeepGroups>
-__device__ __forceinline__ void cp_async_wait_group_dispatch(int keep_groups) {
-    static_assert(MaxKeepGroups >= 0 && MaxKeepGroups <= 7,
-                  "cp.async.wait_group supports immediates in [0, 7]");
-    if (keep_groups == MaxKeepGroups) {
-        cp_async_wait_group<MaxKeepGroups>();
-    } else if constexpr (MaxKeepGroups > 0) {
-        cp_async_wait_group_dispatch<MaxKeepGroups - 1>(keep_groups);
-    } else {
-        cp_async_wait_group<0>();
-    }
-}
-
 }  // namespace astrai

csrc/kernels/fp8/common.h

@@ -11,125 +11,118 @@
 namespace astrai {
 namespace fp8 {
 
-// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
+// Compile-time FP8 format: E4M3 (forward, max 448) or E5M2 (gradients,
-// E5M2 (gradient / large dynamic range, max 57344).
+// max 57344).
 enum class FP8Format : int {
     E4M3 = 0,
     E5M2 = 1,
 };
 
-// Operand memory layouts as types (CUTLASS-style tags). The tag names the
+// Operand storage tags (CUTLASS-style) relative to the canonical matrices
-// storage order of the raw buffer relative to the operand's canonical GEMM
+// A [M][K] / B [K][N]: A RowMajor = [M][K] (default), A ColMajor = [K][M],
-// matrix — A is [M][K], B is [K][N]:
+// B RowMajor = [K][N], B ColMajor = [N][K] (the nn.Linear weight). Selection
-//   A RowMajor = [M][K] storage (K-contiguous rows; the default)
+// is by type at compile time (see gemm.cuh's stage loads).
-//   A ColMajor = [K][M] storage (M-contiguous; A^T)
-//   B RowMajor = [K][N] storage (N-contiguous; the plain a @ b operand)
-//   B ColMajor = [N][K] storage (K-contiguous; the nn.Linear weight layout)
-// Empty tags: selection happens by type at compile time (see load_operand_tile).
 struct RowMajor {};
 struct ColMajor {};
 
-// Transpose of a layout tag: the same buffer with the rows and contract dims
+// Compile-time tile configuration, mirroring KernelTraits in the attention
-// swapped. B's tag is relative to the canonical [K][N] GEMM matrix, so the
+// kernels: CTA tile, warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128
-// stage-load (which views any operand as [rows][contract]) sees the transposed
+// CTA, 32x32 on the 64x64 small CTA) and cp.async pipeline depth.
-// tag — this trait makes that inversion explicit.
+template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
-template <typename Layout>
+          int WarpM = 64, int WarpN = 32>
-struct transpose_layout;
-template <>
-struct transpose_layout<RowMajor> {
-    using type = ColMajor;
-};
-template <>
-struct transpose_layout<ColMajor> {
-    using type = RowMajor;
-};
-template <typename Layout>
-using transpose_layout_t = typename transpose_layout<Layout>::type;
-
-// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
-// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
-// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile and
-// the cp.async pipeline depth.
-template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages>
 struct Fp8GemmTraits {
     static constexpr FP8Format kFormat = Fmt;
     static constexpr int kBlockM = BlockM;
     static constexpr int kBlockN = BlockN;
     static constexpr int kK = K;
     static constexpr int kStages = Stages;
+    static constexpr int kWarpM = WarpM;
+    static constexpr int kWarpN = WarpN;
     static constexpr bool kIsE5M2 = (Fmt == FP8Format::E5M2);
     static constexpr __nv_fp8_interpretation_t kNvFormat =
         kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
     static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
 
-    // Derived launch geometry: 64x32 warp tiles give the CTA thread count.
+    // Derived geometry: warp tiles tile the CTA. The smem budget is
-    // The shared-memory budget is layout-aware (crosswise operands add K-
+    // layout-aware, so it lives in Fp8GemmSmem (gemm.cuh).
-    // major staging + a canonical buffer), so it lives in Fp8GemmSmem in
+    static constexpr int kWarpsM = BlockM / WarpM;
-    // gemm.cuh together with the resident-CTA hint for __launch_bounds__.
+    static constexpr int kWarpsN = BlockN / WarpN;
-    static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
+    static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
+    static_assert(kWarpsM * WarpM == BlockM && kWarpsN * WarpN == BlockN,
+                  "warp tiles must exactly tile the CTA");
+    static_assert(WarpM % 16 == 0 && WarpN % 8 == 0,
+                  "warp tile must be a multiple of the m16n8 MMA shape");
 };
 
-// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional
+// Quantize output orientation: RowMajor = x8 only; Transposed = the
-// delayed-scaling ring finalization. Separate from FP8Params so each
+// [cols][rows] x8T only; Dual = both from a single read. Transposed/Dual
-// operator owns exactly the fields it touches (the GEMM never reads amax /
+// produce K-contiguous operands so crosswise consumers (backward
-// ring state). Same NSDMI rationale: amax / ring_state gate optional paths
+// grad_x / grad_w) route through the NT fast path.
-// via null checks. Still an aggregate, still trivially copyable.
+enum class QuantLayout : int {
+    RowMajor = 0,
+    Transposed = 1,
+    Dual = 2,
+};
+
+// Quantize-kernel parameter POD: float input -> FP8 with fused amax.
 struct FP8QuantizeParams {
-    // BF16 input and FP8 output buffers; scale_a is the quantization step
+    const void* __restrict__ input_ptr = nullptr;
-    // (device scalar). amax_a (may be null) is zero-initialized by the
+    void* __restrict__ output_ptr = nullptr;
-    // binding and receives the raw-domain absolute maximum.
+    void* __restrict__ output_transposed_ptr = nullptr;  // [cols][rows]
-    const void* __restrict__ a_ptr = nullptr;
+    QuantLayout out_layout = QuantLayout::RowMajor;
-    void* __restrict__ out_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;  // device multiplier
-    // [hist[ring_len] | scale | counter] with ring_len = numel - 2. When
+    float* __restrict__ amax = nullptr;         // raw-domain max out
-    // 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).
+    // 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;
+    int rows = 0;
+    int cols = 0;
 };
 
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
-// through the pre-quantized GEMM kernels. Each kernel touches only the
+// through the kernels; each kernel touches only the fields it needs.
-// fields it needs; buffers are raw pointers packed by the torch binding.
-// Pointer members default to null (same NSDMI rationale as AttentionParams:
-// 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
+    // FP8 operands + output; scales are quantization steps (device
-    // quantization steps (device scalars).
+    // scalars). Optional bf16 bias fuses into the epilogue (fp32 add before
+    // the single bf16 rounding); null disables.
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
-    const void* __restrict__ bias = nullptr;
+    const void* __restrict__ bias_ptr = 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;
 
-    // Shapes. `int` covers every realistic LLM shape; the kernels promote
+    const float* __restrict__ scale = nullptr;
-    // to int64 for all pointer arithmetic.
+    // NN-swap mode (canonicalize_gemm): the kernel computes the transposed
-    int m, n, k;
+    // problem and the epilogue scatters D[row][col] to out[col * p.m + row]
+    // in the caller's [M][N] buffer. Zero in the plain orientation.
+    int out_transposed = 0;
+    int m, n, k;  // int covers LLM shapes; kernels promote to int64
 
-    // Physical leading dimensions (column count, i.e. row stride) of A and
+    // Batched (bmm) geometry: grid.z steps these element strides (0
-    // B. For a non-transposed operand the stride equals the contract dim;
+    // broadcasts the operand across batches).
-    // for a transposed operand it is the operand's own column count. The
+    int batch = 1;
-    // binding packs these so the kernel reads both buffers either naturally
+    int64_t a_batch_stride = 0;
-    // or transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
+    int64_t b_batch_stride = 0;
+    int64_t out_batch_stride = 0;
+
+    // Physical leading dims (row strides) of A and B; the binding packs
+    // them so the kernel reads each buffer naturally or transposed per the
+    // LayoutA/LayoutB tags.
     int a_ld, b_ld;
 };
 

csrc/kernels/fp8/gemm/epilogue.cuh

@@ -0,0 +1,180 @@
+#pragma once
+// Collective epilogue: fused bias, the bf16 scatter of the fp32 accumulators
+// through the reclaimed operand shared memory, and the coalesced copy-out.
+
+#include "../common.h"
+#include "policy.cuh"
+
+namespace astrai {
+namespace fp8 {
+
+template <typename Policy>
+struct Fp8CollectiveEpilogue {
+    using Traits = typename Policy::Traits;
+    static constexpr bool kStreamOut = Policy::kStreamOut;
+    static constexpr int kBlockM = Traits::kBlockM;
+    static constexpr int kBlockN = Traits::kBlockN;
+    static constexpr int kMt = Traits::kWarpM / 16;
+    static constexpr int kNt = Traits::kWarpN / 8;
+
+    __nv_bfloat16* const tile_out;
+    const float output_scale;
+    const __nv_bfloat16* const bias;
+    const int64_t m, n;
+    const bool t_out;
+    const int row_elems, row_chunks;
+    const int warp_m, warp_n, group, thread_in_group;
+    const int64_t block_m, block_n;
+
+    __device__ Fp8CollectiveEpilogue(char* smem, const FP8Params& p,
+                                     int64_t block_m, int64_t block_n, int tid)
+        : tile_out(reinterpret_cast<__nv_bfloat16*>(smem)),
+          output_scale(*p.scale),
+          bias(reinterpret_cast<const __nv_bfloat16*>(p.bias_ptr)),
+          m(p.m), n(p.n), t_out(p.out_transposed != 0),
+          row_elems(t_out ? kBlockM : kBlockN),
+          row_chunks(row_elems / 8),
+          warp_m((tid >> 5) / Traits::kWarpsN),
+          warp_n((tid >> 5) % Traits::kWarpsN),
+          group((tid & 31) >> 2),
+          thread_in_group(tid & 3),
+          block_m(block_m), block_n(block_n) {}
+
+    // Swizzled address of one 16B chunk (row r, chunk c) of the staged
+    // tile. Plain orientation: kBlockM rows of kBlockN elems; out-
+    // transposed (swap dispatch): rows and row length trade places. Both
+    // row-chunk counts are powers of two, keeping the XOR swizzle
+    // well-defined.
+    __device__ __forceinline__ __nv_bfloat16* out_chunk(int r, int c) const {
+        return tile_out + (size_t)r * row_elems +
+               ((c ^ (r & (row_chunks - 1))) * 8);
+    }
+    __device__ __forceinline__ __nv_bfloat16* out_elem(int r, int v) const {
+        return out_chunk(r, v >> 3) + (v & 7);
+    }
+
+    // Scatter the accumulators into the staging tile: the operand rings are
+    // dead once the mainloop ends, so their space stages the bf16 output
+    // tile. Threads scatter (STS.32 of bf16x2 pairs), a barrier makes the
+    // tile coherent, then the whole CTA copies it out in fully-coalesced
+    // 16B chunks. The 16B-chunk XOR swizzle keeps both the scatter and the
+    // gather conflict-free.
+    __device__ __forceinline__ void stage(float acc[kNt][kMt][4]) const {
+        // Fused bias: added to the fp32 accumulator before the single bf16
+        // rounding. The per-lane loads are L1 broadcasts; rows past the
+        // edge skip the load (their smem slots never copy out). Under
+        // out_transposed the bias indexes D-cols = the kernel's rows.
+        const int local_col0 = warp_n * Traits::kWarpN + thread_in_group * 2;
+        const int64_t bias_col0 = block_n * kBlockN;
+        const int64_t bias_row0 = block_m * kBlockM;
+        if (!t_out) {
+#pragma unroll
+            for (int nt = 0; nt < kNt; ++nt) {
+                const int col = local_col0 + nt * 8;
+                const int64_t gcol = bias_col0 + col;
+                const float b0 =
+                    bias && gcol < n ? __bfloat162float(bias[gcol]) : 0.0f;
+                const float b1 =
+                    bias && gcol + 1 < n ? __bfloat162float(bias[gcol + 1])
+                                         : 0.0f;
+#pragma unroll
+                for (int mt = 0; mt < kMt; ++mt) {
+                    const int r0 = warp_m * Traits::kWarpM + group + mt * 16;
+                    const float* tile_acc = acc[nt][mt];
+                    // Two bf16x2 stores per accumulator tile: rows g and
+                    // g+8 of the m16n8 output, columns tig*2/tig*2+1 inside
+                    // one 16B chunk.
+                    const int off = col & 7;  // element offset in the chunk
+                    *reinterpret_cast<__nv_bfloat162*>(
+                        out_chunk(r0, col >> 3) + off) =
+                        __floats2bfloat162_rn(tile_acc[0] * output_scale + b0,
+                                              tile_acc[1] * output_scale + b1);
+                    *reinterpret_cast<__nv_bfloat162*>(
+                        out_chunk(r0 + 8, col >> 3) + off) =
+                        __floats2bfloat162_rn(tile_acc[2] * output_scale + b0,
+                                              tile_acc[3] * output_scale + b1);
+                }
+            }
+        } else {
+            // Transposed scatter: accumulator (kernel row r0, col) is
+            // D[col0_global + col][row0_global + r0], staged at T[col][r0].
+            // The acc pair spans two staged rows, so these are scalar
+            // stores (the swap path is the rare NN layout). OOB elements
+            // store dead lanes of the tile, never copied out.
+#pragma unroll
+            for (int nt = 0; nt < kNt; ++nt) {
+                const int col = local_col0 + nt * 8;
+#pragma unroll
+                for (int mt = 0; mt < kMt; ++mt) {
+                    const int r0 = warp_m * Traits::kWarpM + group + mt * 16;
+                    const int64_t grow = bias_row0 + r0;
+                    const float b =
+                        bias && grow < m ? __bfloat162float(bias[grow]) : 0.0f;
+                    const float* tile_acc = acc[nt][mt];
+                    *out_elem(col, r0) =
+                        __float2bfloat16(tile_acc[0] * output_scale + b);
+                    *out_elem(col + 1, r0) =
+                        __float2bfloat16(tile_acc[1] * output_scale + b);
+                    *out_elem(col, r0 + 8) =
+                        __float2bfloat16(tile_acc[2] * output_scale + b);
+                    *out_elem(col + 1, r0 + 8) =
+                        __float2bfloat16(tile_acc[3] * output_scale + b);
+                }
+            }
+        }
+    }
+
+    // Coalesced copy-out: thread -> one 16B chunk; consecutive threads walk
+    // a row so each global transaction covers a full 128B line. Under the
+    // swap the staged rows are D-rows counted from block_n's stripe while
+    // the row length is kernel m', so row/stride flip to the swapped dims.
+    __device__ __forceinline__ void store(__nv_bfloat16* out_bf16) const {
+        constexpr int kTotalChunks =
+            kBlockM * (kBlockN / 8);  // == kBlockN * (kBlockM/8)
+        const int64_t row0_global = block_m * kBlockM;
+        const int64_t col0_global = block_n * kBlockN;
+        for (int idx = threadIdx.x; idx < kTotalChunks; idx += kCtaThreads) {
+            const int r = idx / row_chunks;
+            const int c = idx % row_chunks;
+            const uint4 v = *reinterpret_cast<const uint4*>(out_chunk(r, c));
+            const int64_t row = t_out ? (int64_t)block_n * kBlockN + r
+                                      : row0_global + r;
+            const int64_t col = t_out ? row0_global + (int64_t)c * 8
+                                      : col0_global + (int64_t)c * 8;
+            const int64_t rows_total = t_out ? n : m;
+            const int64_t row_stride = t_out ? m : n;
+            if (row >= rows_total) break;  // rows are consecutive: nothing left
+            auto* dst = out_bf16 + row * row_stride + col;
+            if (col + 8 <= row_stride &&
+                (reinterpret_cast<uintptr_t>(dst) & 15) == 0) {
+                if constexpr (kStreamOut) {
+                    // Evict-first streaming store knob: neutral on L20
+                    // squares, -3..4% on rects; kept for other SKUs.
+                    __stcs(reinterpret_cast<uint4*>(dst), v);
+                } else {
+                    *reinterpret_cast<uint4*>(dst) = v;
+                }
+            } else {
+                // Row-edge chunk or an odd-stride row base: spill the
+                // elements that survive the row edge.
+                const __nv_bfloat16* elems =
+                    reinterpret_cast<const __nv_bfloat16*>(&v);
+                for (int e = 0; e < 8 && col + e < row_stride; ++e)
+                    dst[e] = elems[e];
+            }
+        }
+    }
+
+    __device__ __forceinline__ void run(float acc[kNt][kMt][4],
+                                        __nv_bfloat16* out_bf16) {
+        stage(acc);
+        __syncthreads();
+        store(out_bf16);
+    }
+
+  private:
+    static constexpr int kCtaThreads = Traits::kCtaThreads;
+};
+
+}  // namespace fp8
+}  // namespace astrai

csrc/kernels/fp8/gemm/load.cuh

@@ -0,0 +1,224 @@
+#pragma once
+// Operand loaders: swizzled shared-memory staging for congruous operands
+// (cp.async, predicated and interior variants, plus the loop-carried
+// prefetch state) and the direct LDG+PRMT path for crosswise operands.
+// The staging invariants and the swizzle derivation live in
+// docs/developer/cuda_kernels.md.
+
+#include "../../common/cp_async.cuh"
+#include "../common.h"
+#include "policy.cuh"
+
+namespace astrai {
+namespace fp8 {
+
+// log2 of a compile-time power of two (for the swizzle shifts).
+template <int N, int Acc = 0>
+struct log2_const : log2_const<(N >> 1), Acc + 1> {};
+template <int Acc>
+struct log2_const<1, Acc> {
+    static constexpr int value = Acc;
+};
+
+// Swizzled address inside a flat [rows * K] staging tile: the 16B chunk
+// index is XORed with the row bits at [3, 3+log2(kChunks)) so a warp's
+// ldmatrix fragment load (8 consecutive rows x 16B) hits all 32 banks
+// exactly once; chunks stay contiguous, so cp.async staging is unaffected.
+template <int K, typename T8>
+__device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) {
+    constexpr int kChunks = K / 16;  // 16B chunks per row
+    static_assert(kChunks >= 1 && (kChunks & (kChunks - 1)) == 0,
+                  "swizzle needs a power-of-two 16B-chunk count");
+    constexpr int kShift = 3 - log2_const<kChunks>::value;
+    return tile + row * K +
+           ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + (col & 15));
+}
+
+// Stage-load a CONGRUOUS operand (contract-contiguous storage — the only
+// cp.async-able shape) into the flat [rows * K] swizzled tile. kInterior
+// drops all predication: valid only for a fully interior CTA (whole rows,
+// 16B-aligned base|ld, k_base + K <= contract); the address math then folds
+// to one immediate XOR per chunk (see the design notes). Crosswise operands
+// go through load_crosswise_direct instead.
+template <typename T8, int K, int RowsTile, int kThreads,
+          bool kInterior = false>
+__device__ __forceinline__ void
+load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
+                  int64_t contract, int64_t ld, int tid, int64_t k_base,
+                  int64_t block_row) {
+    constexpr int kChunks = K / 16;
+    static_assert(RowsTile * kChunks % kThreads == 0,
+                  "tile chunks must divide evenly across threads");
+    constexpr int kCpt = RowsTile * kChunks / kThreads;  // chunks per thread
+    constexpr int kCpr = kChunks / kCpt;  // chunks per row slice
+    const int r = tid / kCpr;
+    const int c0 = (tid % kCpr) * kCpt * 16;
+    if constexpr (kInterior) {
+        const char* src = reinterpret_cast<const char*>(
+            operand + (block_row + r) * ld + k_base + c0);
+        const uintptr_t dst =
+            reinterpret_cast<uintptr_t>(tile_at<K>(tile, r, c0));
+#pragma unroll
+        for (int j = 0; j < kCpt; ++j)
+            astrai::cp_async_16(reinterpret_cast<T8*>(dst ^ (j << 4)),
+                                src + j * 16);
+    } else {
+        const int64_t row = block_row + r;
+        const bool row_ok = row < rows;
+        // k_base and every c are multiples of 16, so all chunks share the
+        // row base's alignment verdict.
+        const auto* src = operand + row * ld + k_base;
+        const bool chunk_aligned = (reinterpret_cast<uintptr_t>(src) & 15) == 0;
+#pragma unroll
+        for (int j = 0; j < kCpt; ++j) {
+            const int c = c0 + j * 16;
+            T8* dst = tile_at<K>(tile, r, c);
+            if (row_ok && chunk_aligned && k_base + c + 15 < contract) {
+                astrai::cp_async_16(dst, src + c);
+            } else {
+                // Tail chunk / misaligned base / OOB row: scalar fill.
+#pragma unroll
+                for (int i = 0; i < 16; ++i)
+                    dst[i] =
+                        row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f);
+            }
+        }
+    }
+}
+
+// Loop-carried prefetch state for one congruous operand ring: per-thread
+// (r, c0) mapping with the swizzled stage destination and global source
+// pointer carried across k-tiles, so each prefetch chunk is one LDGSTS
+// issued straight from registers. The guard is a property of the operand's
+// layout, so it lives in the type: the false specialization (crosswise
+// operand) is an empty no-op.
+template <bool kAsync, typename T8, int kK, int kRowsTile, int kThreads>
+struct PrefetchCarry;
+
+template <typename T8, int kK, int kRowsTile, int kThreads>
+struct PrefetchCarry<true, T8, kK, kRowsTile, kThreads> {
+    static constexpr int kCpt = kRowsTile * (kK / 16) / kThreads;
+    static constexpr int kCpr = (kK / 16) / kCpt;
+    unsigned wr = 0;    // current stage's swizzled destination offset
+    unsigned wr0 = 0;   // slot-0 wrap base
+    unsigned wrEnd = 0; // one-past-the-ring sentinel
+    const char* src = nullptr;  // current tile's global source bytes
+
+    __device__ __forceinline__ PrefetchCarry(
+        const T8* ring, int ringSlots, int stageElems, const T8* operand,
+        int64_t ld, int64_t blockRow, int tid, int firstTile) {
+        const int r = tid / kCpr;
+        const int c0 = (tid % kCpr) * kCpt * 16;
+        const T8* slot0 = ring + (firstTile % ringSlots) * stageElems;
+        const unsigned laneOff = static_cast<unsigned>(
+            (const char*)tile_at<kK>(slot0, r, c0) - (const char*)slot0);
+        const unsigned base = __cvta_generic_to_shared(ring) + laneOff;
+        wr = base + (unsigned)((firstTile % ringSlots) * stageElems);
+        wr0 = base;
+        wrEnd = base + (unsigned)(ringSlots * stageElems);
+        src = reinterpret_cast<const char*>(
+                  operand + (blockRow + r) * ld + c0) +
+              (int64_t)firstTile * kK;
+    }
+
+    // Emit this thread's chunks for the current tile; pf false (loop tail)
+    // zero-fills into the slot compute(i-1) already released.
+    __device__ __forceinline__ void emit(bool pf) const {
+#pragma unroll
+        for (int j = 0; j < kCpt; ++j)
+            astrai::cp_async_16(wr ^ (unsigned)(j << 4), src + j * 16, pf);
+    }
+
+    __device__ __forceinline__ void advance(int stageElems) {
+        wr += (unsigned)stageElems;
+        if (wr == wrEnd) wr = wr0;
+        src += kK;
+    }
+};
+
+template <typename T8, int kK, int kRowsTile, int kThreads>
+struct PrefetchCarry<false, T8, kK, kRowsTile, kThreads> {
+    __device__ __forceinline__ PrefetchCarry(
+        const T8*, int, int, const T8*, int64_t, int64_t, int, int) {}
+    __device__ __forceinline__ void emit(bool) const {}
+    __device__ __forceinline__ void advance(int) {}
+};
+
+// Direct (synchronous) crosswise load into a canonical rotating stage:
+// LDG.128 x4 (4 consecutive contract bytes x 16 rows) + in-register PRMT
+// transpose + 16 STS.32. Crosswise operands cannot cp.async into the
+// canonical tile (a 16B global run holds one contract byte for each of 16
+// rows), so they take this path; a staged smem->smem variant measured
+// 15-20% slower and was removed (see git history).
+template <typename T8, int K, int RowsTile, int kThreads>
+__device__ __forceinline__ void
+load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows,
+                      int64_t contract, int64_t ld, int tid, int64_t k_base,
+                      int64_t block_row) {
+    constexpr int kQuads = K / 4;    // 4-byte contract quads per tile
+    constexpr int kGroups = RowsTile / 16;
+    constexpr int kTChunks = kQuads * kGroups;  // 64B chunks per tile
+    // r0 is a multiple of 16 and p*ld preserves alignment whenever ld has
+    // it, so every run of a chunk shares one alignment verdict.
+    const bool run_aligned =
+        ((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
+    for (int chunk = tid; chunk < kTChunks; chunk += kThreads) {
+        const int quad = chunk / kGroups;
+        const int rg = chunk % kGroups;
+        const int64_t r0 = block_row + rg * 16;
+        const bool rows_full = r0 + 15 < rows;
+        if (rows_full && run_aligned) {
+            const int64_t p0 = k_base + quad * 4;
+            uint4 v[4];
+#pragma unroll
+            for (int s = 0; s < 4; ++s) {
+                // Contract tail: a run past k carries zero bytes; they flow
+                // through the PRMT transpose like any other value.
+                if (p0 + s < contract)
+                    v[s] = *reinterpret_cast<const uint4*>(
+                        operand + (p0 + s) * ld + r0);
+                else
+                    v[s] = make_uint4(0u, 0u, 0u, 0u);
+            }
+            const unsigned* bytes = reinterpret_cast<const unsigned*>(v);
+#pragma unroll
+            for (int i = 0; i < 16; ++i) {
+                // word i = row r0+i's quad: byte i of each of the four runs
+                // [v0.b(i), v1.b(i), v2.b(i), v3.b(i)].
+                const unsigned nib = i & 3;
+                const unsigned sel = nib | ((nib + 4) << 4);
+                const unsigned w01 =
+                    __byte_perm(bytes[0 + (i >> 2)], bytes[4 + (i >> 2)], sel);
+                const unsigned w23 =
+                    __byte_perm(bytes[8 + (i >> 2)], bytes[12 + (i >> 2)], sel);
+                *reinterpret_cast<unsigned*>(tile_at<K>(tile, rg * 16 + i,
+                                                        quad * 4)) =
+                    __byte_perm(w01, w23, 0x5410u);
+            }
+        } else {
+            // Row-tail or misaligned chunk: byte-granular gather with
+            // per-row predication; contract-tail columns zero-fill.
+#pragma unroll
+            for (int s = 0; s < 4; ++s) {
+                const int col = quad * 4 + s;
+                if (k_base + col >= contract) {
+#pragma unroll
+                    for (int i = 0; i < 16; ++i)
+                        *tile_at<K>(tile, rg * 16 + i, col) = T8(0.0f);
+                    continue;
+                }
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int64_t r_idx = r0 + i;
+                    *tile_at<K>(tile, rg * 16 + i, col) =
+                        r_idx < rows
+                            ? operand[(k_base + col) * ld + r_idx]
+                            : T8(0.0f);
+                }
+            }
+        }
+    }
+}
+
+}  // namespace fp8
+}  // namespace astrai

csrc/kernels/fp8/gemm/mainloop.cuh

@@ -0,0 +1,336 @@
+#pragma once
+// Collective mainloop: shared-memory stage rings, the gmem->smem stage loads
+// (congruous cp.async / crosswise LDG+PRMT), the per-lane ldmatrix fragment
+// addressing and the software-pipelined mma.sync loop. The fragment
+// addressing scheme and the fast-loop peel rationale live in
+// docs/developer/cuda_kernels.md.
+
+#include <type_traits>
+
+#include "../../common/mma.cuh"
+#include "../common.h"
+#include "load.cuh"
+#include "policy.cuh"
+
+namespace astrai {
+namespace fp8 {
+
+template <typename Policy>
+struct Fp8CollectiveMainloop {
+    using Traits = typename Policy::Traits;
+    using LayoutA = typename Policy::LayoutTagA;
+    using LayoutB = typename Policy::LayoutTagB;
+    using Smem = Fp8GemmSmem<Traits, LayoutA, LayoutB>;
+    static constexpr bool kFastLoop = Policy::kFastLoop;
+    using T8 = std::conditional_t<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>;
+    static constexpr int kBlockM = Traits::kBlockM;
+    static constexpr int kBlockN = Traits::kBlockN;
+    static constexpr int kK = Traits::kK;
+    static constexpr int kStages = Traits::kStages;
+    static constexpr int kCtaThreads = Traits::kCtaThreads;
+    static constexpr bool kDirectA = Smem::kDirectA;
+    static constexpr bool kDirectB = Smem::kDirectB;
+    static_assert(kStages >= 1 && kStages <= 8,
+                  "FP8 GEMM stages must be in [1, 8]");
+    // CTA = (BlockM/WarpM) x (BlockN/WarpN) warps, each warp computing
+    // kMt x kNt m16n8k32 MMAs. Rings rotate kStages+1 buffers (see
+    // Fp8GemmSmem) — one __syncthreads per k-tile.
+    static constexpr int kMt = Traits::kWarpM / 16;  // 16-row MMA tiles per warp
+    static constexpr int kNt = Traits::kWarpN / 8;   // 8-col MMA tiles per warp
+    static constexpr int kSegs = kK / kMmaK;  // mma-sized k segments per tile
+    static constexpr int kARing = Smem::kRingDepth;
+    static constexpr int kBRing = Smem::kRingDepth;
+    static constexpr int kAStageBytes = kBlockM * kK;
+    static constexpr int kBStageBytes = kBlockN * kK;
+
+    T8* const a_base;
+    T8* const b_base;
+    const T8* const a;
+    const T8* const b;
+    const int64_t m, n, k, a_ld, b_ld;
+    const int tid;
+    const int64_t block_m, block_n;
+    const int warp_m, warp_n;
+    const int a_row0;  // + mt * 16 in the loop
+    const int b_row0;  // + nt * 8
+    const int64_t tile_count;
+    // Interior-CTA peel (kFastLoop instantiations only): whole-CTA,
+    // 16B-aligned, K without tail — the mainloop then runs a compile-time
+    // specialized copy with no per-chunk predication (measured +4.5..10% on
+    // the issue-bound small CTA; the 128x128 CTA regressed, so only the
+    // small CTA opts in). The verdict is uniform per CTA.
+    const bool fast_cta;
+
+    __device__ Fp8CollectiveMainloop(char* smem, const T8* a, const T8* b,
+                                     int64_t m, int64_t n, int64_t k,
+                                     int64_t a_ld, int64_t b_ld, int tid,
+                                     int2 block)
+        : a_base(reinterpret_cast<T8*>(smem)),
+          b_base(reinterpret_cast<T8*>(smem + kARing * kAStageBytes)),
+          a(a), b(b), m(m), n(n), k(k), a_ld(a_ld), b_ld(b_ld), tid(tid),
+          block_m(block.x), block_n(block.y),
+          warp_m((tid >> 5) / Traits::kWarpsN),
+          warp_n((tid >> 5) % Traits::kWarpsN),
+          a_row0(warp_m * Traits::kWarpM),
+          b_row0(warp_n * Traits::kWarpN),
+          tile_count((k + kK - 1) / kK),
+          fast_cta(kFastLoop && !kDirectA && !kDirectB &&
+                   ((int64_t)block.x * kBlockM + kBlockM <= m) &&
+                   ((int64_t)block.y * kBlockN + kBlockN <= n) &&
+                   ((reinterpret_cast<uintptr_t>(a) | (uint64_t)a_ld) & 15) == 0 &&
+                   ((reinterpret_cast<uintptr_t>(b) | (uint64_t)b_ld) & 15) == 0 &&
+                   (k % kK) == 0) {}
+
+    // Stage-slot helpers: rings rotate one slot per k-tile, so callers
+    // either compute the slot from the tile index (prologue, generic loop)
+    // or carry an advancing pointer (steady-state fast loop).
+    __device__ __forceinline__ T8* a_stage_of(int64_t tile) const {
+        return a_base + (size_t)(tile % kARing) * kAStageBytes;
+    }
+    __device__ __forceinline__ T8* b_stage_of(int64_t tile) const {
+        return b_base + (size_t)(tile % kBRing) * kBStageBytes;
+    }
+    // Asynchronous congruous loads for one k-tile: cp.async into the
+    // canonical rings; kFast selects the predication-free interior copy
+    // (fast_cta admits only congruous operands). Called after the
+    // post-compute barrier, alongside the commit.
+    template <bool kFast = false>
+    __device__ __forceinline__ void load_async(T8* a_stage, T8* b_stage,
+                                               int64_t k_base) const {
+        if constexpr (!kDirectA)
+            load_operand_tile<T8, kK, kBlockM, kCtaThreads, kFast>(
+                a_stage, a, m, k, a_ld, tid, k_base, block_m * kBlockM);
+        if constexpr (!kDirectB)
+            load_operand_tile<T8, kK, kBlockN, kCtaThreads, kFast>(
+                b_stage, b, n, k, b_ld, tid, k_base, block_n * kBlockN);
+    }
+    // Synchronous direct-crosswise loads for one k-tile. In the steady
+    // state this runs right after barrier 1, so the LDG latency and the
+    // PRMT transpose overlap the MMA phase instead of stalling the
+    // inter-barrier window.
+    __device__ __forceinline__ void load_direct(T8* a_stage, T8* b_stage,
+                                                int64_t k_base) const {
+        if constexpr (kDirectA)
+            load_crosswise_direct<T8, kK, kBlockM, kCtaThreads>(
+                a_stage, a, m, k, a_ld, tid, k_base, block_m * kBlockM);
+        if constexpr (kDirectB)
+            load_crosswise_direct<T8, kK, kBlockN, kCtaThreads>(
+                b_stage, b, n, k, b_ld, tid, k_base, block_n * kBlockN);
+    }
+
+    // Prime the pipeline: kStages committed groups, one per stage slot.
+    // The commit is unconditional — when K is shorter than the pipeline the
+    // skipped stages commit empty groups, so the group sequence stays
+    // tile-indexed and the steady-state wait count never needs a runtime
+    // dispatch.
+    __device__ __forceinline__ void prologue() const {
+#pragma unroll
+        for (int stage = 0; stage < kStages; ++stage) {
+            if (stage < tile_count) {
+                if (fast_cta)
+                    load_async<true>(a_stage_of(stage), b_stage_of(stage),
+                                     (int64_t)stage * kK);
+                else
+                    load_async(a_stage_of(stage), b_stage_of(stage),
+                               (int64_t)stage * kK);
+                load_direct(a_stage_of(stage), b_stage_of(stage),
+                            (int64_t)stage * kK);
+            }
+            astrai::cp_async_commit_group();
+        }
+    }
+
+    // Steady-state mainloop, compile-time specialized on kFast: the fast
+    // copy runs predication-free loads with loop-carried read/write
+    // pointers; the generic copy keeps full predication. kFastLoop=false
+    // instantiates only the generic copy.
+    template <bool kFast>
+    __device__ __forceinline__ void run_loop(float acc[kNt][kMt][4]) const {
+        const int lane = tid & 31;
+        // Fast-path write carries: one per congruous operand (crosswise
+        // operands get the empty no-op type), targeting the first
+        // prefetched tile (kStages). Steady-state read carries: the LDSM
+        // base of the current k-tile's stage with the lane offset folded
+        // in, advanced one stage per iteration with an equality wrap —
+        // replaces the per-k-tile (tile % ring) * stage_bytes
+        // recomputation (a UIMAD.WIDE magic-division ladder in SASS).
+        PrefetchCarry<!kDirectA, T8, kK, kBlockM, kCtaThreads> carry_a(
+            a_base, kARing, kAStageBytes, a, a_ld, block_m * kBlockM, tid,
+            kStages);
+        PrefetchCarry<!kDirectB, T8, kK, kBlockN, kCtaThreads> carry_b(
+            b_base, kBRing, kBStageBytes, b, b_ld, block_n * kBlockN, tid,
+            kStages);
+        const unsigned a_rd0 = __cvta_generic_to_shared(a_base) + a_lane_off(lane);
+        const unsigned b_rd0 =
+            __cvta_generic_to_shared(b_base) +
+            (kPairB ? b4_lane_off(lane) : b_lane_off(lane));
+        const unsigned a_rd_end = a_rd0 + (unsigned)(kARing * kAStageBytes);
+        const unsigned b_rd_end = b_rd0 + (unsigned)(kBRing * kBStageBytes);
+        unsigned a_rd = a_rd0, b_rd = b_rd0;
+        for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) {
+        // In the steady state exactly kStages-1 younger groups are in flight
+        // when this fires; the tail's unconditional (possibly empty)
+        // commits keep that invariant true for every iteration.
+        const bool prefetch = tile_index + kStages < tile_count;
+        astrai::cp_async_wait_group<kStages - 1>();
+        // Barrier 1: every thread's cp.async for this stage is complete
+        // before any thread reads tiles written by other threads.
+        __syncthreads();
+
+        // Direct chunks for tile i+kStages: issue LDG+PRMT+STS now so the
+        // global-load latency hides behind the MMA phase below.
+        if (prefetch)
+            load_direct(a_stage_of(tile_index + kStages),
+                        b_stage_of(tile_index + kStages),
+                        (tile_index + kStages) * kK);
+
+        const unsigned a_addr = a_rd;
+        const unsigned b_addr = b_rd;
+        // Per-k_seg base pair (cuBLAS's scheme): seg s lives at the seg-0
+        // base XOR (s<<5) — one LOP3 per extra seg per k-tile, never per
+        // fragment. Every LDSM below addresses [base + immediate].
+        unsigned a_seg[kSegs], b_seg[kSegs];
+#pragma unroll
+        for (int s = 0; s < kSegs; ++s) {
+            a_seg[s] = a_addr ^ (unsigned)(s * kSegXor);
+            b_seg[s] = b_addr ^ (unsigned)(s * kSegXor);
+        }
+
+        // kNt ldmatrix.x2 (B) + kMt ldmatrix.x4 (A) feed kMt*kNt*2 mma.sync
+        // per k_seg — 0.5 load instructions per MMA. B fragments
+        // double-buffer across k_segs; kPairB folds the two adjacent nt
+        // fragments of one pair into a single x4 (see b4_lane_off).
+        unsigned b_frag[2][kNt][2];
+        unsigned b_frag4[2][kNt / 2][4];
+        load_b_frags(b_frag[0][0], b_frag4[0][0], b_seg[0]);
+#pragma unroll
+        for (int k_seg = 0; k_seg < kSegs; ++k_seg) {
+            const int bcur = k_seg & 1, bnext = bcur ^ 1;
+            if (k_seg + 1 < kSegs)
+                load_b_frags(b_frag[bnext][0], b_frag4[bnext][0],
+                             b_seg[k_seg + 1]);
+        // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 is
+        // issued before the MMAs consuming row mt, so the LDS latency hides
+        // behind tensor-pipe work. Costs 4 extra registers.
+        unsigned a_frag[kMt + 1][4];
+        astrai::ldmatrix_x4_lane(a_frag[0], a_seg[k_seg]);
+#pragma unroll
+        for (int mt = 0; mt < kMt; ++mt) {
+            if (mt + 1 < kMt)
+                astrai::ldmatrix_x4_lane(a_frag[mt + 1],
+                                         a_seg[k_seg] + (mt + 1) * kMtStep);
+#pragma unroll
+            for (int nt = 0; nt < kNt; ++nt) {
+                const unsigned* bops =
+                    kPairB ? (b_frag4[bcur][nt >> 1] + (nt & 1) * 2)
+                           : b_frag[bcur][nt];
+                astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt], bops,
+                                     acc[nt][mt]);
+            }
+        }
+        // Next tile's LDGSTS chunks inside the MMA phase: A's after the
+        // first k_seg's MMA batch, B's after the last.
+        if constexpr (kFast) {
+            if (k_seg == 0) carry_a.emit(prefetch);
+            if (k_seg == kSegs - 1) carry_b.emit(prefetch);
+        }
+        }
+        // Generic loop (no interleaved prefetch): the next tile's predicated
+        // loads run after the MMA phase.
+        if constexpr (!kFast) {
+            if (prefetch) {
+                load_async(a_stage_of(tile_index + kStages),
+                           b_stage_of(tile_index + kStages),
+                           (tile_index + kStages) * kK);
+            }
+        }
+        // Unconditional commit: empty in the tail, it pads the group
+        // sequence so the fixed wait above stays correct.
+        astrai::cp_async_commit_group();
+        a_rd += (unsigned)kAStageBytes;
+        if (a_rd == a_rd_end) a_rd = a_rd0;
+        b_rd += (unsigned)kBStageBytes;
+        if (b_rd == b_rd_end) b_rd = b_rd0;
+        if constexpr (kFast) {
+            carry_a.advance(kAStageBytes);
+            carry_b.advance(kBStageBytes);
+        }
+        }
+    }
+
+    __device__ __forceinline__ void accumulate(float acc[kNt][kMt][4]) const {
+        if constexpr (kFastLoop) {
+            if (fast_cta)
+                run_loop<true>(acc);
+            else
+                run_loop<false>(acc);
+        } else {
+            run_loop<false>(acc);
+        }
+    }
+
+  private:
+    // Per-lane ldmatrix fragment addressing (base-pair scheme, mirrored
+    // from the cuBLAS SASS; derivation in the design notes): one base
+    // register per operand per k_seg, every fragment offset an LDSM
+    // immediate — zero address arithmetic inside the MMA phase.
+    __device__ __forceinline__ unsigned a_lane_off(int lane) const {
+        const int r7 = lane & 7;          // row within the 8-row matrix
+        const int rh8 = (lane >> 3) & 1;  // +8 rows (A: lanes 8-15, 24-31)
+        const int rh16 = lane >> 4;       // +1 chunk (A: lanes 16-31)
+        constexpr int kChunks = kK / 16;
+        constexpr int kShift = 3 - log2_const<kChunks>::value;  // tile_at's shift
+        const unsigned lswz =
+            static_cast<unsigned>((r7 >> kShift) & (kChunks - 1));
+        // Stage-relative, loop-invariant per-lane base; A's fragment row
+        // carries the +8-row (rh8) and +1-chunk (rh16) halves.
+        return static_cast<unsigned>((a_row0 + rh8 * 8 + r7) * kK +
+                                     ((rh16 ^ lswz) << 4));
+    }
+    __device__ __forceinline__ unsigned b_lane_off(int lane) const {
+        const int r7 = lane & 7;
+        const int rh8 = (lane >> 3) & 1;  // +8 rows (B uses rh8 as its chunk half)
+        constexpr int kChunks = kK / 16;
+        constexpr int kShift = 3 - log2_const<kChunks>::value;
+        const unsigned lswz =
+            static_cast<unsigned>((r7 >> kShift) & (kChunks - 1));
+        return static_cast<unsigned>((b_row0 + r7) * kK + ((rh8 ^ lswz) << 4));
+    }
+    // x4-paired B loads: one ldmatrix.x4 feeds the two adjacent nt
+    // fragments. Lane contract: lanes 0-7 address rows n0..n7 chunk c,
+    // lanes 8-15 rows n0..n7 chunk c+1, lanes 16-23 rows n8..n15 chunk c,
+    // lanes 24-31 rows n8..n15 chunk c+1. The +8-row step never reaches
+    // the swizzle source bits for kK <= 64; kK=128 swizzles on row[2:0]
+    // where +8 flips bits, so that config keeps the x2 loads.
+    static constexpr unsigned kMtStep = 16 * kK;  // bytes per m-tile row step
+    static constexpr unsigned kNtStep = 8 * kK;   // bytes per n-tile row step
+    static constexpr unsigned kSegXor = 32;       // chunk-index +2 per k_seg
+    static constexpr bool kPairB = kK / 16 <= 4;
+    static_assert(!kPairB || kNt % 2 == 0, "B pairing needs even kNt");
+    static constexpr unsigned kPairStep = 16 * kK;  // bytes per nt-pair row step
+    __device__ __forceinline__ unsigned b4_lane_off(int lane) const {
+        return b_lane_off(lane) + (lane >> 4) * kPairStep / 2;
+    }
+
+    // One k_seg's B-fragment loads, shared by the initial fill and the
+    // double-buffer's next-seg fill. frag2/frag4 are the flat bases of one
+    // b_frag / b_frag4 buffer (the unused one is never touched).
+    __device__ __forceinline__ void
+    load_b_frags(unsigned* frag2, unsigned* frag4, unsigned seg_base) const {
+#pragma unroll
+        for (int p = 0; p < kNt / 2; ++p) {
+            if constexpr (kPairB) {
+                astrai::ldmatrix_x4_lane(frag4 + p * 4,
+                                         seg_base + p * kPairStep);
+            } else {
+                astrai::ldmatrix_x2_lane(frag2 + p * 4,
+                                         seg_base + p * 2 * kNtStep);
+                astrai::ldmatrix_x2_lane(frag2 + p * 4 + 2,
+                                         seg_base + (p * 2 + 1) * kNtStep);
+            }
+        }
+    }
+};
+
+}  // namespace fp8
+}  // namespace astrai

csrc/kernels/fp8/gemm/policy.cuh

@@ -0,0 +1,53 @@
+#pragma once
+// Kernel policy layer: shared-memory budget, occupancy hint and the
+// single Policy type the kernel and collectives take (CUTLASS-style
+// consolidation of traits + layout tags + scheduling knobs).
+
+#include <type_traits>
+
+#include "../common.h"
+
+namespace astrai {
+namespace fp8 {
+
+// m16n8k32 (see astrai::mma_shape<fp8 type>::k in common/mma.cuh)
+constexpr int kMmaK = 32;
+
+// Layout-aware shared-memory budget and occupancy hint. Every operand ring
+// holds kStages+1 buffers: the load for tile i+kStages targets slot
+// (i-1)%(kStages+1) — already consumed — so neither load path needs a
+// post-compute barrier (one __syncthreads per k-tile; see the design notes
+// in docs/developer/cuda_kernels.md). The 48KB static watermark picks the
+// resident-CTA hint for __launch_bounds__.
+template <typename Traits, typename LayoutA, typename LayoutB>
+struct Fp8GemmSmem {
+    // Crosswise (direct-load) operands: A ColMajor storage, B RowMajor
+    // storage (B's tag is relative to the canonical [K][N]).
+    static constexpr bool kDirectA = std::is_same_v<LayoutA, ColMajor>;
+    static constexpr bool kDirectB = std::is_same_v<LayoutB, RowMajor>;
+    static constexpr int kRingDepth = Traits::kStages + 1;
+    static constexpr int kBytes =
+        kRingDepth * (Traits::kBlockM + Traits::kBlockN) * Traits::kK;
+    static constexpr int kMinCtas = kBytes <= 48 * 1024 ? 2 : 1;
+};
+
+template <FP8Format Fmt_, int BlockM_, int BlockN_, typename LayoutA_,
+          typename LayoutB_, int WarpM_, int WarpN_, int kK_, int Stages_,
+          int GroupRaster_, bool StreamOut_ = false, bool FastLoop_ = false>
+struct Fp8GemmPolicy {
+    using Traits =
+        Fp8GemmTraits<Fmt_, BlockM_, BlockN_, kK_, Stages_, WarpM_, WarpN_>;
+    using LayoutTagA = LayoutA_;
+    using LayoutTagB = LayoutB_;
+    static constexpr int kGroupRaster = GroupRaster_;
+    static constexpr bool kStreamOut = StreamOut_;
+    static constexpr bool kFastLoop = FastLoop_;
+    using Smem = Fp8GemmSmem<Traits, LayoutA_, LayoutB_>;
+    // Flattened for __launch_bounds__, which takes no dependent type names.
+    static constexpr int kCtaThreads = Traits::kCtaThreads;
+    static constexpr int kMinCtas = Smem::kMinCtas;
+    static constexpr int kSmemBytes = Smem::kBytes;
+};
+
+}  // namespace fp8
+}  // namespace astrai

csrc/kernels/fp8/gemm/scheduler.cuh

@@ -0,0 +1,28 @@
+#pragma once
+// Tile scheduler: the linear CTA id maps to (block_m, block_n) in grouped
+// (L2-friendly) raster — consecutive CTAs share one B column stripe — or
+// plain N-fastest raster (kRasterGroup=0, the measured best for dX's
+// crosswise-B layouts where grouping was neutral).
+
+namespace astrai {
+namespace fp8 {
+
+template <int kRasterGroup>
+struct Fp8GemmTileScheduler {
+    static __device__ int2 tile(const uint3& block, const dim3& blocks) {
+        if constexpr (kRasterGroup > 0) {
+            constexpr int kGroupM = kRasterGroup;
+            const int bid = int(block.y) * int(blocks.x) + int(block.x);
+            const int group_first_m = (bid / (kGroupM * int(blocks.x))) * kGroupM;
+            const int group_rows =
+                min(int(blocks.y) - group_first_m, kGroupM);  // M-tail group is short
+            return int2{group_first_m + bid % group_rows,
+                        (bid % (kGroupM * int(blocks.x))) / group_rows};
+        } else {
+            return int2{int(block.y), int(block.x)};
+        }
+    }
+};
+
+}  // namespace fp8
+}  // namespace astrai

csrc/kernels/fp8/ops.cu

@@ -1,520 +1,307 @@
-// FP8 GEMM torch binding: tensor validation, FP8Params packing, template
+// CUDA bindings for the stateless FP8 quantize/GEMM 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) ----
+// Inner-layout resolution for one GEMM operand. The user flag names the
+// math (0 = last two dims are [rows][contract], 1 = transposed); the
+// storage may independently be a col-major view (.t() of a contiguous
+// buffer), which folds into the returned dispatch flag at zero copy — the
+// kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the
+// user flag only. Tensors whose inner dims are neither natural layout fall
+// back to .contiguous().
+bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
+                     int64_t& batch_stride, torch::Tensor& storage) {
+    torch::Tensor t = t_in;
+    bool col_major = false;
+    if (t.stride(-1) != 1) {
+        if (t.stride(-2) == 1) {
+            col_major = true;
+        } else {
+            t = t.contiguous();
+        }
+    }
+    storage = t;
+    ld = col_major ? t.stride(-1) : t.stride(-2);
+    batch_stride = t.dim() == 3 ? t.stride(0) : 0;
+    return flag ^ col_major;
+}
 
-void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
+// Dtype dispatch over the unified quantize launcher.
-                      const torch::Tensor& sa, const torch::Tensor& sb,
+template <bool Tiled, FP8Format Fmt>
-                      const torch::Tensor* out_scale, const void* bias,
+void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p,
-                      const torch::Tensor* bias_scale, int64_t m, int64_t n,
+                      cudaStream_t stream) {
-                      int64_t k, int64_t a_ld, int64_t b_ld) {
+    switch (x.scalar_type()) {
-    p.a_ptr = a;
+    case torch::kHalf:
-    p.b_ptr = b;
+        launch_fp8_quantize<Fmt, __half, Tiled>(p, stream);
-    p.out_ptr = out;
+        break;
-    p.scale_a = sa.data_ptr<float>();
+    case torch::kFloat32:
-    p.scale_b = sb.data_ptr<float>();
+        launch_fp8_quantize<Fmt, float, Tiled>(p, stream);
-    p.out_scale = out_scale ? out_scale->data_ptr<float>() : nullptr;
+        break;
-    p.bias = bias;
+    default:
-    p.bias_scale = bias_scale ? bias_scale->data_ptr<float>() : nullptr;
+        launch_fp8_quantize<Fmt, __nv_bfloat16, Tiled>(p, stream);
+    }
+}
+
+template <bool Tiled>
+void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p,
+                         bool e5m2, cudaStream_t stream) {
+    if (e5m2)
+        launch_for_dtype<Tiled, FP8Format::E5M2>(x, p, stream);
+    else
+        launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
+}
+
+// Shared binding body for the two quantize entry points: RowMajor /
+// Transposed (single output) serve quantize(), Dual (both orientations from
+// one read) serves quantize_dual(). A ring tensor switches
+// on the in-kernel delayed-scaling fold: state layout
+// [hist n | scale | legacy | amax | done-as-int], and the returned amax is
+// the (self-cleaned) persistent slot. Without it, amax is reduced into a
+// 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 ||
+                    x.scalar_type() == torch::kFloat32,
+                "x must be bf16, fp16 or fp32");
+    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 == QuantLayout::RowMajor || x.dim() >= 2,
+                "transposed quantize layouts need a 2D+ tensor");
+    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 out_opts = input.options().dtype(
+        fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn);
+    torch::Tensor amax;
+    float *ring_hist = nullptr, *ring_scale_out = nullptr;
+    unsigned int* ring_done = nullptr;
+    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 = layout;
+    p.rows = static_cast<int>(input.size(-2));
+    p.cols = static_cast<int>(input.size(-1));
+    torch::Tensor output, output_t;
+    if (layout != QuantLayout::Transposed) {
+        output = torch::empty_like(input, out_opts);
+        p.output_ptr = output.data_ptr();
+    }
+    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 == QuantLayout::RowMajor)
+        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 == QuantLayout::Dual)
+        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,
+                     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,
+                "a and b must be fp8");
+    TORCH_CHECK(a.scalar_type() == b.scalar_type(), "a and b must share format");
+    TORCH_CHECK((a.dim() == 2 || a.dim() == 3) &&
+                    (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());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    // Batched operands follow matmul broadcast rules: 2D acts as a batch
+    // of 1; a size-1 batch broadcasts across the other side (stride 0).
+    const int64_t batch_a = a.dim() == 3 ? a.size(0) : 1;
+    const int64_t batch_b = b.dim() == 3 ? b.size(0) : 1;
+    TORCH_CHECK(batch_a == batch_b || batch_a == 1 || batch_b == 1,
+                "batch dim mismatch (got ", batch_a, " and ", batch_b, ")");
+    const int64_t batch = std::max(batch_a, batch_b);
+    TORCH_CHECK(batch <= 65535, "batch dim exceeds the grid.z launch limit");
+
+    torch::Tensor a_st, b_st;
+    int64_t a_ld, b_ld, a_bstride, b_bstride;
+    const bool tag_a = resolve_operand(a, trans_a, a_ld, a_bstride, a_st);
+    const bool tag_b = resolve_operand(b, trans_b, 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);
+    const int64_t n = trans_b ? b.size(-2) : b.size(-1);
+    TORCH_CHECK(k == (trans_b ? b.size(-1) : b.size(-2)), "inner dim mismatch");
+
+    const bool batched_out = a.dim() == 3 || b.dim() == 3;
+    torch::Tensor output =
+        batched_out
+            ? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16))
+            : torch::empty({m, n}, a.options().dtype(torch::kBFloat16));
+    FP8Params p;
+    p.a_ptr = a_st.data_ptr();
+    p.b_ptr = b_st.data_ptr();
+    p.out_ptr = output.data_ptr();
+    p.scale = scale.data_ptr<float>();
     p.m = static_cast<int>(m);
     p.n = static_cast<int>(n);
     p.k = static_cast<int>(k);
     p.a_ld = static_cast<int>(a_ld);
     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.
-// Pack the quantize params, optionally wiring the delayed-scaling ring.
+    if (bias_t.defined() && bias_t.numel() > 0) {
-// ring (may be null) packs [hist[len] | scale | counter]; len/margin come
+        TORCH_CHECK(bias_t.is_cuda() && bias_t.scalar_type() == torch::kBFloat16,
-// from the active recipe and idx is the caller's slot for this step.
+                    "fp8 gemm bias must be a CUDA bf16 tensor");
-void pack_quantize_params(FP8QuantizeParams& p, const void* x, void* x8,
+        TORCH_CHECK(bias_t.dim() == 1 && bias_t.size(0) == n,
-                          const torch::Tensor& scale, torch::Tensor* amax,
+                    "fp8 gemm bias must be 1D of length n=", n);
-                          const torch::Tensor* ring, int64_t ring_idx,
+        TORCH_CHECK(bias_t.is_contiguous(), "fp8 gemm bias must be contiguous");
-                          int64_t ring_margin, int64_t total) {
+        p.bias_ptr = bias_t.data_ptr();
-    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);
+    p.batch = static_cast<int>(batch);
-}
+    p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride;
-
+    p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride;
-// ---- GEMM launch dispatch (runtime flags -> compile-time kernel variants) ----
+    p.out_batch_stride = m * n;
-
-template <FP8Format Fmt, int Variant>
-void launch_gemm_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);
-}
-
-template <FP8Format Fmt>
-void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool out_fp8,
-                   bool trans_a, bool trans_b) {
-    // Encode the runtime flags as [output FP8, transpose A, transpose B].
-    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 1: launch_gemm_variant<Fmt, 1>(p, stream); break;
-        case 2: launch_gemm_variant<Fmt, 2>(p, stream); break;
-        case 3: launch_gemm_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
-
-// ---------------------------------------------------------------------------
-// Entry points
-// ---------------------------------------------------------------------------
-
-std::tuple<torch::Tensor, torch::Tensor> quantize_bf16(torch::Tensor x,
-                                                       torch::Tensor scale,
-                                                       int64_t fmt) {
-    // BF16 -> FP8 quantize with fused amax. fmt: 0 = E4M3, 1 = E5M2.
-    // Returns (x8, amax); the caller never clears amax (zero-initialized here).
-    TORCH_CHECK(x.is_cuda() && scale.is_cuda(), "CUDA tensors required");
-    TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
-    check_scale(scale, x, "scale");
-    check_fp8_device(x);
-    const at::cuda::OptionalCUDAGuard guard(x.device());
-    auto stream = at::cuda::getCurrentCUDAStream();
-
-    auto x_c = x.contiguous();
-    auto x8 = torch::empty_like(
-        x_c, x_c.options().dtype(fmt ? torch::kFloat8_e5m2
-                                      : torch::kFloat8_e4m3fn));
-    // 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,
-                         nullptr, 0, 0, x_c.numel());
-    if (fmt) {
-        launch_fp8_quantize<FP8Format::E5M2>(p, stream.stream());
-    } else {
-        launch_fp8_quantize<FP8Format::E4M3>(p, stream.stream());
-    }
-    C10_CUDA_CHECK(cudaGetLastError());
-    return {x8, amax};
-}
-
-torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
-                     torch::Tensor sb, int64_t out_dtype,
-                     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)");
-    TORCH_CHECK(a.scalar_type() == b.scalar_type(),
-                "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");
-    check_scale(sa, a, "sa");
-    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);
-    // Physical leading dimension = column count of each contiguous buffer.
-    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 == k2, "inner dim mismatch");
-    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,
-                     out_fp8 ? &os : nullptr, nullptr, nullptr, m, n, k, a_ld,
-                     b_ld);
     if (a.scalar_type() == torch::kFloat8_e4m3fn)
-        dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), out_fp8, ta, tb);
+        gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
     else
-        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), out_fp8, ta, tb);
+        gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
     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"), py::arg("transposed") = false,
-          "BF16 to FP8 (E4M3/E5M2) quantize with fused amax; returns (x8, amax)");
+          py::arg("ring") = py::none(), py::arg("hist_idx") = 0,
-    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("sa"),
+          py::arg("fp8_max") = 448.0, py::arg("pow2_margin") = 1.0);
-          py::arg("sb"), py::arg("out_dtype") = 0,
+    m.def("quantize_dual", &quantize_dual, py::arg("x"), py::arg("scale"),
-          py::arg("out_scale") = py::none(), py::arg("trans_a") = 0,
+          py::arg("fmt"), py::arg("ring") = py::none(),
-          py::arg("trans_b") = 0,
+          py::arg("hist_idx") = 0, py::arg("fp8_max") = 448.0,
-          "Pre-quantized FP8 GEMM: op(a) @ op(b)^T * (sa * sb); out_dtype "
+          py::arg("pow2_margin") = 1.0);
-          "0=bf16, 1=fp8 e4m3 (requires out_scale); trans_a/trans_b select "
+    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
-          "the operand layout (default 0/0 = a@b)");
+          py::arg("trans_a") = false, py::arg("trans_b") = false,
-    m.def("linear_forward_fp8", &linear_forward_fp8, py::arg("x"),
+          py::arg("bias") = py::none());
-          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,311 @@
+#pragma once
+// FP8 quantize device code — pure CUDA, no torch: kernels take the
+// FP8QuantizeParams POD, format and input type ride on template parameters,
+// and the launcher is shared by the torch binding and the 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, the unpack of one
+// 16-byte load into kVecElems floats, and a native 2-element pair load.
+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 __nv_bfloat162* b2 =
+            reinterpret_cast<const __nv_bfloat162*>(&raw);
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const float2 p = __bfloat1622float2(b2[j]);
+            f[2 * j] = p.x;
+            f[2 * j + 1] = p.y;
+        }
+    }
+    static __device__ __forceinline__ void load_pair(const __nv_bfloat16* p,
+                                                     float* f) {
+        const float2 v = __bfloat1622float2(
+            *reinterpret_cast<const __nv_bfloat162*>(p));
+        f[0] = v.x;
+        f[1] = v.y;
+    }
+};
+
+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;
+        }
+    }
+    static __device__ __forceinline__ void load_pair(const __half* p,
+                                                     float* f) {
+        const float2 v =
+            __half22float2(*reinterpret_cast<const __half2*>(p));
+        f[0] = v.x;
+        f[1] = v.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) {
+        const unsigned* w = reinterpret_cast<const unsigned*>(&raw);
+#pragma unroll
+        for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]);
+    }
+    static __device__ __forceinline__ void load_pair(const float* p,
+                                                     float* f) {
+        f[0] = p[0];
+        f[1] = p[1];
+    }
+};
+
+// One float -> one fp8 byte (round-nearest-even + satfinite).
+template <FP8Format Fmt>
+__device__ __forceinline__ uint8_t cvt_fp8(float v) {
+    if constexpr (Fmt == FP8Format::E5M2)
+        return __nv_fp8_e5m2(v).__x;
+    else
+        return __nv_fp8_e4m3(v).__x;
+}
+
+// One float pair -> one packed fp8x2 word (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));
+}
+
+// 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(const FP8QuantizeParams& p,
+                                             float v) {
+    v = warp_reduce_max(v);
+    __shared__ float slots[kWarps];
+    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
+    if ((tid & 31) == 0) slots[tid >> 5] = v;
+    __syncthreads();
+    if (tid == 0) {
+#pragma unroll
+        for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
+        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 (QuantLayout::RowMajor): vectorized 16B loads
+// -> fp8 stores, 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);
+    uint8_t* x8 = static_cast<uint8_t*>(p.output_ptr);
+    float local_amax = 0.0f;
+    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
+
+    // One 16B load -> kVecElems bytes per step. Torch allocations are >=16B
+    // aligned, so element 0 keeps the uint4 access natural; a misaligned
+    // base (odd storage offset view) falls to the scalar tail 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));
+        x8[i] = cvt_fp8<Fmt>(v * mult);
+    }
+    if (p.amax) publish_amax<8>(p, local_amax);
+}
+
+// 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
+// read); rows whose pair is unaligned or ragged (odd widths, misaligned
+// bases) fall back to element loads in place. Staging goes through a byte
+// tile whose pitch keeps the store stride coprime with the 32 banks.
+// (+25-35% over the former 32x32 scalar kernel on sub-4M tensors; ~5%
+// slower once DRAM-saturated — accepted for the single-kernel shape.)
+template <FP8Format Fmt, typename InT>
+__global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
+    constexpr int kTileC = 64, kTileR = 32;
+    // 34B pitch: staging stride is 17 words (coprime with the 32 banks) so
+    // the pair-byte stores stay conflict-free, and the byte-wise consume
+    // reads still span distinct words.
+    __shared__ uint8_t tile[kTileC][kTileR + 2];
+    const float mult = *p.scale;
+    const auto* x = static_cast<const InT*>(p.input_ptr);
+    const int r0 = blockIdx.y * kTileR;
+    const int c0 = blockIdx.x * kTileC;
+    const int r = r0 + threadIdx.y * 4;
+    const int c = c0 + threadIdx.x * 2;  // cols even => the pair is in-bounds
+
+    uint8_t q[4][2];
+    float local_amax = 0.0f;
+    // Vectorize the pair when both elements are in-bounds and the native
+    // 2-element load is aligned; odd widths, misaligned bases and ragged
+    // edges fall back to element loads row by row.
+    constexpr int kPairAlign = 2 * (int)sizeof(InT);
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        q[j][0] = 0;
+        q[j][1] = 0;
+        if (r + j < p.rows && c < p.cols) {
+            const InT* a = x + (int64_t)(r + j) * p.cols + c;
+            if (c + 1 < p.cols &&
+                (reinterpret_cast<uintptr_t>(a) & (kPairAlign - 1)) == 0) {
+                float f[2];
+                quant_in_traits<InT>::load_pair(a, f);
+#pragma unroll
+                for (int k = 0; k < 2; ++k) {
+                    local_amax = fmaxf(local_amax, fabsf(f[k]));
+                    q[j][k] = cvt_fp8<Fmt>(f[k] * mult);
+                }
+            } else {
+                const float v0 = quant_in_traits<InT>::to_float(a[0]);
+                local_amax = fmaxf(local_amax, fabsf(v0));
+                q[j][0] = cvt_fp8<Fmt>(v0 * mult);
+                if (c + 1 < p.cols) {
+                    const float v1 = quant_in_traits<InT>::to_float(a[1]);
+                    local_amax = fmaxf(local_amax, fabsf(v1));
+                    q[j][1] = cvt_fp8<Fmt>(v1 * mult);
+                }
+            }
+        }
+    }
+    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)
+            if (r + j < p.rows && c < p.cols) {
+                uint8_t* o = out + (int64_t)(r + j) * p.cols + c;
+                const int64_t off = (int64_t)(r + j) * p.cols + c;
+                if (c + 1 < p.cols && (off & 1) == 0)
+                    *reinterpret_cast<unsigned short*>(o) =
+                        (unsigned short)(q[j][0] | (q[j][1] << 8));
+                else {
+                    o[0] = q[j][0];
+                    if (c + 1 < p.cols) o[1] = q[j][1];
+                }
+            }
+    }
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+#pragma unroll
+        for (int k = 0; k < 2; ++k)
+            tile[threadIdx.x * 2 + k][threadIdx.y * 4 + j] = q[j][k];
+    __syncthreads();
+    // Transposed scatter: output element (c, r) lives at c * rows + r;
+    // threadIdx.x tracks r so each warp writes one contiguous run. tile is
+    // [col][row]; warp y walks 8 columns, threads read down one column.
+    uint8_t* out_t = static_cast<uint8_t*>(p.output_transposed_ptr);
+#pragma unroll
+    for (int i = 0; i < 8; ++i) {
+        const int oc = c0 + threadIdx.y * 8 + i;
+        if (oc < p.cols && r0 + threadIdx.x < p.rows)
+            out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
+                tile[threadIdx.y * 8 + i][threadIdx.x];
+    }
+    if (p.amax) publish_amax<8>(p, local_amax);
+}
+
+// Unified quantize launcher: Tiled selects the transpose kernel
+// (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>
+void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
+    if constexpr (Tiled) {
+        const dim3 grid((p.cols + 63) / 64, (p.rows + 31) / 32);
+        if (grid.x == 0 || grid.y == 0) return;
+        fp8_quantize_tiled_kernel<Fmt, InT><<<grid, dim3(32, 8), 0, stream>>>(p);
+    } else {
+        constexpr int kThreads = 256;
+        constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
+        // Grid-stride loops: any grid >= 1 is correct; one block per 256
+        // vectors plus the tail block covers tiny and misaligned tensors.
+        const int64_t blocks = 1 + p.total / (kVecElems * kThreads);
+        fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
+    }
+}
+
+}  // namespace fp8
+}  // namespace astrai

csrc/tests/fp8_test.cu

@@ -170,20 +170,46 @@ static bool test_single_mma() {
 // Part 2: GEMM correctness — layouts x K-tiles vs fp32 CPU reference
 // ---------------------------------------------------------------------------
 
+// Naive fp32 reference on the GPU: same layout interpretation as the CPU
+// loop it replaces (O(m*n) to check instead of O(m*n*k) to compute).
+__global__ static void
+naive_gemm_ref(const __nv_fp8_e4m3* a, const __nv_fp8_e4m3* b, float* out,
+               int m, int n, int k, int a_ld, int b_ld, int a_rm, int b_rm) {
+    const int i = blockIdx.y * 32 + threadIdx.y;
+    const int j = blockIdx.x * 32 + threadIdx.x;
+    if (i >= m || j >= n) return;
+    float acc = 0.f;
+    for (int kk = 0; kk < k; ++kk) {
+        float av = a_rm ? (float)a[i * a_ld + kk] : (float)a[kk * a_ld + i];
+        float bv = b_rm ? (float)b[kk * b_ld + j] : (float)b[j * b_ld + kk];
+        acc += av * bv;
+    }
+    out[i * n + j] = acc;
+}
+
+// Big-CTA policies for the direct-layout cases: kK/Stages vary per case;
+// the fast interior loop follows the dual-congruous rule, grouped raster 8
+// matches the production dispatch.
+template <typename LA, typename LB>
+constexpr bool kCaseFast =
+    !std::is_same_v<LA, ColMajor> && !std::is_same_v<LB, RowMajor>;
+template <typename LA, typename LB, int kK, int Stages>
+using CasePolicy =
+    Fp8GemmPolicy<FP8Format::E4M3, 128, 128, LA, LB, 64, 32, kK, Stages, 8,
+                  false, kCaseFast<LA, LB>>;
+
 template <typename LA, typename LB, int kK, int Stages>
 static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
-                          int k, int a_ld, int b_ld) {
+                          int k, int a_ld, int b_ld, int dispatch = 0) {
     __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 +227,33 @@ 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);
+    float* d_ref;
+    cudaMalloc(&d_ref, (size_t)m * n * 4);
+    naive_gemm_ref<<<dim3((n + 31) / 32, (m + 31) / 32), dim3(32, 32)>>>(
+        da, db, d_ref, m, n, k, a_ld, b_ld,
+        !std::is_same_v<LA, ColMajor>, !std::is_same_v<LB, ColMajor>);
+    std::vector<float> href((size_t)m * n);
+    cudaMemcpy(href.data(), d_ref, href.size() * 4, cudaMemcpyDeviceToHost);
+    cudaFree(d_ref);
+
+    if (dispatch == 1)
+        // Production route, NN: the dual-N-contiguous problem has no
+        // dedicated instantiation — canonicalize_gemm swaps to the
+        // transposed <ColMajor, ColMajor> kernel with its out-transposed
+        // epilogue (see gemm.cuh).
+        gemm<FP8Format::E4M3>(p, 0, false, false);
+    else if (dispatch == 2)
+        // Production route, NT: exercises plan_gemm's small/narrow/big
+        // selection for this shape.
+        gemm<FP8Format::E4M3>(p, 0, false, true);
+    else
+        launch_policy<CasePolicy<LA, LB, kK, Stages>>(p, 0);
     cudaError_t e = cudaDeviceSynchronize();
     if (e != cudaSuccess) {
         printf("  CUDA err: %s\n", cudaGetErrorString(e));
@@ -221,20 +266,7 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     bool ok = true;
     for (int i = 0; i < m && ok; ++i) {
         for (int j = 0; j < n && ok; ++j) {
-            float ref = 0;
+            const float ref = href[(size_t)i * n + j];
-            for (int kk = 0; kk < k; ++kk) {
-                // A reference reads the actual uploaded buffer: LA ColMajor
-                // means the buffer is [K][M] (ha_t), else [M][K].
-                float av = std::is_same_v<LA, ColMajor>
-                               ? (float)__nv_fp8_e4m3(ha[kk * m + i])
-                               : (float)__nv_fp8_e4m3(ha[i * k + kk]);
-                float bv;
-                if (std::is_same_v<LB, ColMajor>)
-                    bv = (float)__nv_fp8_e4m3(hb[j * k + kk]);
-                else
-                    bv = (float)__nv_fp8_e4m3(hb[kk * n + j]);
-                ref += av * bv;
-            }
             float got =
                 __bfloat162float(__ushort_as_bfloat16(hb16[i * n + j]));
             float err = fabsf(got - ref);
@@ -247,8 +279,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;
 }
 
@@ -258,6 +289,7 @@ static bool test_gemm() {
     } cfgs[] = {
         {128, 128, 128}, {256, 128, 256}, {128, 256, 64},
         {100, 130, 96},  {64, 64, 160},   {300, 200, 320},
+        {2048, 256, 512}, {1024, 1024, 512},
     };
     bool all = true;
     for (auto& c : cfgs) {
@@ -278,12 +310,12 @@ static bool test_gemm() {
         printf(" NT K64:");
         all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(ha, hb_colmajor, c.m,
                                                         c.n, c.k, c.k, c.k);
-        printf(" NN K32:");
+        printf(" NN swap:");
-        all &= run_gemm_case<RowMajor, RowMajor, 32, 3>(ha, hb_rowmajor, c.m,
+        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(
-                                                        c.n, c.k, c.k, c.n);
+            ha, hb_rowmajor, c.m, c.n, c.k, c.k, c.n, /*dispatch=*/1);
-        printf(" NN K64:");
+        printf(" NT disp:");
-        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(ha, hb_rowmajor, c.m,
+        all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(
-                                                        c.n, c.k, c.k, c.n);
+            ha, hb_colmajor, c.m, c.n, c.k, c.k, c.k, /*dispatch=*/2);
         printf(" TN K32:");
         all &= run_gemm_case<ColMajor, ColMajor, 32, 3>(ha_t, hb_colmajor, c.m,
                                                         c.n, c.k, c.m, c.k);

docker-compose.yml

@@ -1,5 +1,6 @@
 services:
   server:
+    image: astrai:latest
     build:
       context: .
       dockerfile: Dockerfile
@@ -20,7 +21,7 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              count: 1
+              count: all
               capabilities: [gpu]
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
@@ -31,6 +32,7 @@ services:
     restart: unless-stopped
 
   server-cpu:
+    image: astrai:latest
     profiles: [cpu]
     build:
       context: .
@@ -54,6 +56,7 @@ services:
     restart: unless-stopped
 
   trainer:
+    image: astrai:latest
     profiles: [train]
     build:
       context: .

docs/developer/architecture.md

@@ -4,7 +4,7 @@
 
 - [Class Diagram](#class-diagram) — Full Mermaid class diagram across 10+ namespaces
 - [Module Overview](#module-overview) — Component inventory per module
-- [Design Patterns](#design-patterns) — 15 documented patterns with classes
+- [Design Patterns](#design-patterns) — 16 documented patterns with classes
 - [Core Relationships](#core-relationships) — 11 key inter-component relationships
 
 ## Class Diagram
@@ -816,8 +816,8 @@ classDiagram
 
         class Executor {
             +AutoModel model
-            +AutoTokenizer tokenizer
             +PagePool kv_cache
+            +TaskCacheManager task_cache
             +InferenceWorkspace _workspace
             +Optional[str] device
             +Optional[torch.dtype] dtype
@@ -845,6 +845,7 @@ classDiagram
 
         class InferenceScheduler {
             +PagePool _cache
+            +TaskCacheManager _task_cache
             +Executor _executor
             +TaskManager _task_mgr
             +Event _stop_event
@@ -888,6 +889,24 @@ classDiagram
             +release(pages)
         }
 
+        class AllocationStrategy {
+            <<abstract>>
+            +alloc(state, prompt_ids) bool
+            +free(state)
+            +extend(state, pos) bool
+            +write_indices(state, prompt_ids)
+            +record_hashes(state, prompt_ids, start_logical_page)
+        }
+
+        class ContiguousStrategy {
+            +write_indices(state, prompt_ids)
+        }
+
+        class PagedStrategy {
+            -Allocator _alloc
+            -RadixCache _prefix
+        }
+
         class KVStorage {
             +int size
             +Tensor k_buffer
@@ -926,14 +945,21 @@ classDiagram
             +bool contiguous
             -KVStorage _storage
             -ReqToTokenPool _req_pool
-            -Allocator _alloc
+            -AllocationStrategy _strategy
-            -RadixCache _prefix
+            +strategy AllocationStrategy
+            +req_pool ReqToTokenPool
+            +bind_tasks(req_indices, seq_lens, workspace, device, start_pos, incremental) KVCache
+        }
+
+        class TaskCacheManager {
+            -PagePool _pool
+            -Dict _states
             +task_alloc(task_id, prompt_ids) bool
             +task_free(task_id)
             +task_extend(task_id, pos) bool
             +task_cached(task_id) int
             +task_record_hashes(task_id, prompt_ids, start_logical_page)
-            +bind_tasks(task_ids, workspace, device, start_pos) KVCache
+            +bind(task_ids, workspace) KVCache
         }
 
     class Task {
@@ -1316,17 +1342,22 @@ classDiagram
     PositionIdStrategy <|-- DocResetPositionId
     PositionIdStrategy <|-- ContinuousPositionId
     StoreWriter <|-- BinWriter
+    AllocationStrategy <|-- ContiguousStrategy
+    AllocationStrategy <|-- PagedStrategy
     RawRollout <|-- RolloutResult
     LaunchStrategy <|-- TorchrunStrategy
     LaunchStrategy <|-- LocalStrategy
     %% --- Composition (strong ownership, part destroyed with whole) ---
     PagePool *-- KVStorage
     PagePool *-- ReqToTokenPool
-    PagePool *-- Allocator
+    PagePool *-- AllocationStrategy
-    PagePool *-- RadixCache
+    PagedStrategy *-- Allocator
+    PagedStrategy *-- RadixCache
+    TaskCacheManager o-- PagePool
     RadixCache *-- RadixNode
     InferenceEngine *-- InferenceScheduler
     InferenceScheduler *-- PagePool
+    InferenceScheduler *-- TaskCacheManager
     InferenceScheduler *-- Executor
     Executor *-- InferenceWorkspace
     InferenceScheduler *-- TaskManager
@@ -1419,7 +1450,7 @@ classDiagram
     Task --> TaskStatus
     InferenceEngine --> AutoModel
     Executor --> AutoModel
-    Executor --> AutoTokenizer
+    Executor --> TaskCacheManager
     TaskManager --> AutoTokenizer
 
 ```
@@ -1436,7 +1467,7 @@ classDiagram
 | **astrai.model** | ModelFactory, AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, TaskCacheManager, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, AllocationStrategy, ContiguousStrategy, PagedStrategy, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
 | **astrai.extension** | `backend` policy package, `ops` kernel-wrapper package, `fp8.py` FP8 strategy layer, AttentionBackend, TorchNativeBackend, CudaBackend, FlashAttnBackend, attention, attn_backend, ATTN_BACKEND, apply_rotary_emb, is_available | Stable API over attention/rotary/FP8 execution policy and optional CUDA kernels |
 | **astrai.optim** | OptimizerFactory, MuonAdamW, NoraNadamW, ManoAdamW, composite_step/composite_zero_grad/composite_state_dict, partition_optimizer_parameters | Built-in optimizers (`muon_adamw` / `nora_nadamw` / `mano_adamw`) with shared composite-optimizer helpers |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
@@ -1478,4 +1509,4 @@ classDiagram
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-08-22
+> Document Update Time: 2026-08-29

docs/developer/cuda_kernels.md

@@ -41,25 +41,99 @@ Standalone benchmark vs torch complex-multiply (48 calls = 24 layers × q+k): 6-
 
 The `fp8_ops` family (`csrc/kernels/fp8/`) accelerates bf16 linear layers by
 quantizing to FP8 and running tensor-core GEMMs (**requires sm_89+**; fp8
-`mma.sync.m16n8k32` only exists on Ada/Hopper). It follows the same three-layer
+`mma.sync.m16n8k32` only exists on Ada/Hopper). Same three-layer style as
-style as attention, but split into **three** files:
+attention; the GEMM device code is split humming/CUTLASS-style into one
+layered directory:
 
 | File | Role |
 |------|------|
-| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `FP8Params` POD — no torch |
+| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `FP8Params` / `FP8QuantizeParams` PODs, layout tags — 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: vectorized `fp8_quantize_kernel` + 32×32-tile transpose kernel (out_layout 0/1/2), `quant_in_traits<InT>` unpack — no torch |
+| `fp8/gemm/policy.cuh` | smem budget / occupancy hint (`Fp8GemmSmem`) + `Fp8GemmPolicy` (traits + layouts + knobs — the kernel's single template parameter) |
+| `fp8/gemm/load.cuh` | operand loaders: swizzle (`tile_at`), congruous cp.async (predicated + interior), `PrefetchCarry`, crosswise LDG+PRMT direct load |
+| `fp8/gemm/scheduler.cuh` | CTA id → (block_m, block_n) grouped/plain raster |
+| `fp8/gemm/mainloop.cuh` | `Fp8CollectiveMainloop`: stage rings, stage loads, fragment addressing, pipelined mma.sync loop |
+| `fp8/gemm/epilogue.cuh` | `Fp8CollectiveEpilogue`: fused bias + bf16 smem scatter + coalesced copy-out |
+| `fp8/gemm.cuh` | umbrella: `fp8_gemm_kernel<Policy>` orchestrator + host planning (`plan_gemm` / `launch_plan`; 64×64 / 128×64 / 128×128 CTA) + entry `gemm<Fmt>(params, stream, trans_a, trans_b)` = `canonicalize_gemm` → `plan_gemm` → `launch_plan` |
 | `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*; the
-`scale`; the kernel computes the reciprocal internally — the interface never
+strategy layer passes `scale.reciprocal()` and the kernel multiplies by it.
-takes `*_inv`). `amax` is always returned in the original bf16 domain.
+`mm_fp8` takes the combined dequant scale (`sa * sb`). `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 (`fp8_quantize` / `fp8_gemm`) via `torch.library.custom_op`, with
-`linear_backward_fp8`) via `torch.library.custom_op`, and
+plain `quantize` / `mm_fp8` wrappers, and `astrai/extension/fp8.py` is the
-`astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed /
+strategy layer (`fp8_autocast`, delayed / dynamic scaling recipes,
-dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear`
+`fp8_linear_forward/backward` wiring `aten::linear` on CUDA). See the FP8
-on CUDA). See the FP8 section in `AGENTS.md` for full detail.
+section in `AGENTS.md` for full detail.
+
+#### FP8 GEMM design notes
+
+The load-bearing invariants behind the kernel code (all measurements on
+L20/sm_89 unless noted):
+
+**Swizzle.** Staging tiles are flat `[rows * kK]`; `tile_at` XORs the 16B
+chunk index with row bits at `[3, 3+log2(kChunks))` so a warp's ldmatrix
+fragment load (8 consecutive rows × 16B) hits all 32 banks exactly once
+(the unswizzled row word-stride is `kK/4` words, so rows `r` and
+`r + 8/kChunks` collide mod 32). Chunks stay contiguous, so cp.async
+staging is unaffected.
+
+**Fragment addressing (base-pair scheme).** One base register per operand
+per k_seg, every fragment offset an LDSM immediate. The closure works
+because the XOR swizzle's source bits come only from the lane's
+row-within-matrix `r7`: the 8/16-row fragment steps never reach them, so
+`addr(s, mt) = lane_base + mt*(16*kK) ^ (s<<5)` for A and
+`addr(s, nt) = lane_base + nt*(8*kK) ^ (s<<5)` for B. This replaced
+runtime offset tables that spilled at 131 registers (~55 of 146 hot-loop
+instructions were address math; cuBLAS's inner loop has ~0). Steady-state
+read pointers advance one stage per iteration with an equality wrap,
+replacing the per-k-tile `(tile % ring) * stage_bytes` recomputation
+(UIMAD.WIDE magic-division ladder).
+
+**Pipeline depth and barriers.** Every operand ring holds `kStages+1`
+buffers: the load for tile `i+kStages` targets slot `(i-1)%(kStages+1)`,
+which compute(i-1) finished reading before this iteration's barrier — no
+post-compute barrier, one `__syncthreads` per k-tile. Prologue and tail
+commits are unconditional so the group sequence stays tile-indexed and the
+fixed `wait_group<kStages-1>` is iteration-invariant (a runtime
+wait-count dispatch ladder cost 16 instructions/k-tile). A lean
+`kStages`-deep ring trading the barrier for a 4th resident CTA measured
++5..9% slower at 1280³ and was removed.
+
+**Crosswise loads.** Crosswise operands (A `[K][M]` / B `[N][K]` storage)
+cannot cp.async into the canonical tile; they take the direct LDG.128×4 +
+in-register PRMT transpose + STS.32 path. A staged variant (cp.async into
+K-major staging + per-tile smem→smem transpose) measured 15-20% slower
+across every probed shape including DRAM-streaming B (git history 5745c2f).
+
+**Fast-loop peel.** When both operands are congruous, the whole CTA is
+interior, base|ld is 16B-aligned and K has no tail, the mainloop switches
+to a predication-free copy with loop-carried prefetch state: +4.5..10% on
+the issue-bound 64×64 CTA (256³..1024³), −3% on the 128×128 CTA, so only
+the small CTA opts in.
+
+**Launch planning crossovers** (L20, TFLOPS, big vs alternative):
+crosswise problems keep the 64×64 s3 CTA below ~1.5 waves of 128×128
+tiles (M=256: 129.7 vs 113.1; 1024³: 107.2 vs 94.8; the big CTA wins from
+M=640/1536³ on). Dual-congruous wave band picks narrow vs big by
+`ceil(tiles/sm) * T_tile` with `T_narrow ≈ 0.53 * T_big` (M=384: 134.3 vs
+114.4 narrow wins; M=1024: 202.5 vs 178.8 big wins). Sub-wave: narrow
+wins past ~3/8 of a wave (1024³ 174 vs 131T), the big CTA's operand reuse
+wins past ~5/8 (forcing 64×64 there cost 2048³ 123→171T). Non-128-divisible
+shapes with 64-divisibility take the 64×64 CTA (edge tiles otherwise drag
+the single wave; 1088³: 76 vs 93T). Persistent schedules (static
+round-robin and atomic ticket) both measured worse on L20 (−4..−8%; the
+ticket variant recovers L2 locality but its loop-head barrier costs what
+the CTA-restart overlap saves).
+
+**NN swap.** The dual-N-contiguous problem runs as its transpose
+`E = B^T @ A^T` over swapped operands with an out-transposed epilogue
+scatter (CUTLASS-sm90 `is_swapAB`): one instantiation fewer per tile
+config, at the cost of a scalar-store scatter on a path no LLM-linear
+operand pair hits.
 
 ## Build System
 
@@ -98,7 +172,9 @@ unset, `setup.py` auto-detects the real GPU capability through
 - **sm_80+** (Ampere and later): enables the tensor-core MMA path
   (`mma.sync.m16n8k16.bf16` for bf16 attention, `mma.sync.m16n8k32` for FP8).
 - **sm_89+**: required for the FP8 family (`fp8_ops`) — FP8 tensor-core
-  instructions only exist on Ada/Hopper and newer.
+  instructions only exist on Ada/Hopper and newer. On older architectures,
+  CMake emits a warning and skips the `fp8_ops` target so the remaining CUDA
+  kernels still build successfully.
 - **`-DASTRAI_NO_MMA`** is a manual escape hatch only — the build never defines
   it automatically. To disable the MMA path, add it to `NVCC_FLAGS` yourself;
   all supported build targets are sm_80+.
@@ -112,7 +188,7 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
              --ptxas-options=-O3,-v --extra-device-vectorization --threads=16
 ```
 
-Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all six kernel targets in parallel via `cmake --build -j N`. The target list is the **single source of truth**: `KERNEL_NAMES` and the parallel `KERNEL_SRCS` list in `csrc/CMakeLists.txt`; `astrai/extension/loader.py` auto-discovers the compiled `.so` files.
+Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all registered kernel targets in parallel via `cmake --build -j N` (the five base targets always; `fp8_ops` additionally on sm_89+). The target list is the **single source of truth**: `KERNEL_NAMES` and the parallel `KERNEL_SRCS` list in `csrc/CMakeLists.txt`; `astrai/extension/loader.py` auto-discovers the compiled `.so` files.
 
 ## Python Extension Architecture
 
@@ -229,7 +305,7 @@ cycle belong under `TYPE_CHECKING`.
 
 - **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
 - **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
-- **`FlashAttnBackend`**: Optional flash-attn dispatch with `flash_attn_with_kvcache` fast path.
+- **`FlashAttnBackend`**: Optional flash-attn dispatch via `flash_attn_varlen_func` over gathered flat K/V.
 - **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
 
 Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash``
@@ -298,19 +374,26 @@ q_tile_to_batch = [0, 0, 1, 2, 2, 2]
 q_tile_to_index = [0, 1, 0, 0, 1, 2]
 ```
 
-Paged prefill launches:
+Paged prefill launches (MMA path, GQA head packing):
 
 ```text
-grid.x = num_q_tiles  # 6, exactly the valid ragged work items
+grid.x = num_q_tiles * HB   # HB = min(G, WARPS): q heads packed per block
-grid.y = q_heads
+grid.y = kv_heads * ceil(G / HB)
 grid.z = 1
 ```
 
-Each block resolves its request and request-local tile in O(1):
+The tensor-core prefill kernel packs `HB = min(G, WARPS)` query heads of one
+kv-head group into a block, so K/V tiles stream once per block instead of once
+per q head (~HB× less global K/V traffic). Warp `w` handles head slot `w / WPH`
+and 16-row chunk `w % WPH`, where `WPH = WARPS / HB`; `G = q_heads / kv_heads`
+and `G = 1` (MHA) degenerates to the historical one-head-per-block layout.
+Each host Q tile (64 rows, `Q_TILE_ROWS`) splits into `HB` packed blocks along
+`grid.x`. Each block resolves its request and request-local row range in O(1):
 
 ```cpp
-batch = q_tile_to_batch[blockIdx.x];
+host_tile = blockIdx.x / HB;
-q_tile = q_tile_to_index[blockIdx.x];
+batch = q_tile_to_batch[host_tile];
+row_base = q_tile_to_index[host_tile] * 64 + (blockIdx.x % HB) * (64 / HB);
 ```
 
 The kernel then uses `qo_indptr[batch]` for the packed Q base and adjacent
@@ -332,7 +415,7 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
 Test files:
 - `attn_test.cu` — decode + prefill kernels (correctness tables + benchmarks)
 - `attn_paged_test.cu` — paged decode/prefill kernels
-- `fp8_mma_test.cu` — BF16→FP8→BF16 MMA demo (sm_89)
+- `fp8_test.cu` — single-warp bf16→fp8→mma.sync sanity check + full FP8 GEMM correctness (sm_89)
 
 ## Benchmarks
 
@@ -359,7 +442,9 @@ csrc/
 ├── kernels/
 │   ├── common/                       # cross-family pure-CUDA helpers (no torch)
 │   │   ├── device.cuh                #   sm_at_least(), kMinSmForFp8* constants
-│   │   └── mma.cuh                   #   shared mma_sync<InT> + mma_shape<InT> (bf16 m16n8k16 / fp8 m16n8k32) + ldmatrix_x2/x4<T>
+│   │   ├── mma.cuh                   #   shared mma_sync<InT> + mma_shape<InT> (bf16 m16n8k16 / fp8 m16n8k32) + ldmatrix_x2/x4<T>
+│   │   ├── cp_async.cuh              #   cp.async 16B primitives (predicated copy, commit/wait groups)
+│   │   └── reduce.cuh                #   warp_reduce_max, atomic_max_float
 │   ├── attention/                    # attention family (module names keep the attn_* prefix)
 │   │   ├── common.h                  #   AttentionParams POD, TensorLayout enum (BHLD/BLHD)
 │   │   ├── warp_utils.cuh            #   warp reduction helpers
@@ -370,7 +455,7 @@ csrc/
 │   │   ├── decode_split_kv.cuh       #   decode kernel, scalar (split-KV)
 │   │   ├── decode_split_kv_mma.cuh   #   decode kernel, MMA + split-K
 │   │   ├── prefill_split_q.cuh       #   prefill kernel, scalar (split-Q)
-│   │   ├── prefill_split_q_mma.cuh   #   prefill kernel, MMA (split-Q, packed/ragged Q schedule)
+│   │   ├── prefill_split_q_mma.cuh   #   prefill kernel, MMA (split-Q, GQA head packing, packed/ragged Q schedule)
 │   │   ├── decode.cu                 #   → module attn_decode
 │   │   ├── prefill.cu                #   → module attn_prefill
 │   │   ├── paged_decode.cu           #   → module attn_paged_decode
@@ -378,16 +463,23 @@ csrc/
 │   ├── rotary/
 │   │   └── rotary_emb.cu             # rotary embedding (kernel + binding in one file) → module rotary_emb
 │   └── fp8/                          # FP8 family (module name fp8_ops)
-│       ├── common.h                  #   FP8Format enum, Fp8GemmTraits, FP8Params POD (no torch)
+│       ├── common.h                  #   FP8Format enum, Fp8GemmTraits, FP8Params / FP8QuantizeParams PODs, layout tags (no torch)
-│       ├── gemm.cuh                  #   FP8 device code: quantize + pre-quantized GEMM kernels (no torch)
+│       ├── quantize.cuh              #   quantize kernels: vectorized + 32×32-tile transpose (out_layout 0/1/2) (no torch)
-│       └── mm.cu                     #   binding only: validation, param packing, launch dispatch, pybind
+│       ├── gemm.cuh                  #   GEMM umbrella: kernel orchestrator + host launch planning (no torch)
+│       ├── gemm/                     #   GEMM device layers (humming/CUTLASS-style split)
+│       │   ├── policy.cuh            #     smem budget / occupancy hint + Fp8GemmPolicy
+│       │   ├── load.cuh              #     operand loaders (swizzle, congruous cp.async, crosswise direct)
+│       │   ├── scheduler.cuh         #     grouped/plain raster mapping
+│       │   ├── mainloop.cuh          #     stage rings + pipelined mma.sync mainloop
+│       │   └── epilogue.cuh          #     fused bias + bf16 scatter + copy-out
+│       └── ops.cu                    #   binding only: validation, param packing, launch dispatch, pybind
 └── tests/
     ├── test_utils.cuh                # Shared test utilities (now_ms, f2bf, bf2f, randf)
     ├── attn_test.cu                  # Decode + prefill kernels
     ├── attn_paged_test.cu            # Paged decode/prefill kernels
-    └── fp8_mma_test.cu               # BF16→FP8→BF16 MMA demo
+    └── fp8_test.cu                   # MMA demo + GEMM correctness across layouts/K tiles/ragged shapes
 ```
 
 Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
 
-> Document Update Time: 2026-08-22
+> Document Update Time: 2026-08-29

docs/developer/docker-serving.md

@@ -17,7 +17,7 @@ scripts/serve.sh             preflight, Compose wrapper, lifecycle
         └── server.py --config /run/astrai/serve.yaml
 ```
 
-`scripts/tools/serve_runtime.py` reads `runtime:` plus the two container-side
+`scripts/docker/serve_runtime.py` reads `runtime:` plus the two container-side
 values Compose needs (`server.port` for the port mapping, `server.device` for
 the preflight GPU check). `scripts/tools/server.py --config` reads `server:`.
 Explicit CLI arguments to `server.py` override `server:` YAML values.
@@ -54,8 +54,9 @@ server:
 - `runtime.gpu.enabled: true` (default) selects the `server` service with an
   NVIDIA device reservation; `false` selects `server-cpu` (no GPU passthrough).
   When disabled, `server.device` must be `cpu`.
-- `runtime.gpu.devices` is either `all` or a single-device list such as `[0]`;
+- `runtime.gpu.devices` is `all` (default) or a single-device list such as `[0]`;
-  the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: 1`.
+  the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: all`; the
+  env var performs the only filtering.
 - `environment` values are explicitly passed to the serving container. Keep
   host-specific settings here; they are not universal defaults.
 - `server.device` must agree with `runtime.gpu.enabled`; `preflight` enforces it.
@@ -89,9 +90,9 @@ bash scripts/serve.sh status [CONFIG]
 
 `preflight` validates Docker, the model directory
 (`config.json` + `model.safetensors`), GPU/device consistency, and the
-rendered Compose configuration. `up` starts the container detached and
+rendered Compose configuration. `up` starts the container detached; `run`
-rebuilds the image when the code changed (`--build`); `run` keeps it in the
+keeps it in the foreground. Both reuse the existing image; run
-foreground. The wrapper manages a fixed container name
+`bash scripts/serve.sh build [CONFIG]` after code changes. The wrapper manages a fixed container name
 (`astrai-server` or `astrai-server-<job_name>`); the plain
 `docker compose up -d` / `docker compose --profile cpu up -d` path keeps
 working with defaults (port 8000, `./params`).
@@ -99,7 +100,7 @@ working with defaults (port 8000, `./params`).
 ## Hard Rules
 
 1. Keep Docker settings in `runtime` and server settings in `server`.
-2. Filter GPUs once: the `server` service reserves one device; a `devices`
+2. Filter GPUs once: Compose passes `count: all`; a `devices`
    list becomes `CUDA_VISIBLE_DEVICES`.
 3. `runtime.gpu.enabled: false` requires `server.device: cpu`.
 4. In Docker, `server.port` must match the published container port (default

docs/developer/docker-training.md

@@ -18,7 +18,7 @@ scripts/train.sh             preflight, Compose wrapper, lifecycle, timer
               └── train.py --config /run/astrai/train.yaml
 ```
 
-The two parsers deliberately own different sections. `scripts/tools/train_runtime.py`
+The two parsers deliberately own different sections. `scripts/docker/train_runtime.py`
 reads only `runtime`; `scripts/tools/train.py` reads only
 `model/data/parallel/training/ckpt/log`. Explicit trainer arguments after `--`
 override training YAML values.
@@ -42,10 +42,11 @@ runtime:
     stop_timeout_seconds: 600
     checkpoint_keep_last: 5
     # max_duration_hours: 12
-  # Add host-specific workarounds only when required:
+  # Optional; entries are passed verbatim into the trainer container
+  # (see "Per-Job Environment"):
   # environment:
-  #   NCCL_P2P_DISABLE: "1"
+  #   ASTR_LOG_LEVEL: DEBUG
-  #   NCCL_NET_GDR_LEVEL: "0"
+  #   ASTR_BACKEND: torch_native
 ```
 
 - Relative paths resolve from the YAML file's directory, not the current shell.
@@ -57,11 +58,85 @@ runtime:
   Use `fsdp` explicitly when model sharding is required.
 - To select specific physical GPUs, replace `all` with a list such as
   `devices: [0, 1]`.
-- `environment` values are explicitly passed to the training container. Keep
+- `environment` entries apply only to the job defined by this YAML file, not to
-  host-specific NCCL workarounds here; they are not universal defaults.
+  the host or to other jobs. Keep the section omitted unless this job's GPU
+  selection needs it; see [Per-Job Environment](#per-job-environment).
 - `max_duration_hours` starts a detached host timer that calls the same graceful
   `stop` command. A manual stop cancels the timer.
 
+## Per-Job Environment
+
+`runtime.environment` is scoped to one job. `start` passes only the entries of
+the config file it was given, so a variable reaches exactly the GPUs declared
+in that file's `runtime.gpu.devices` and nothing else. Two jobs on the same
+machine can therefore differ: a job whose GPUs have working peer-to-peer keeps
+the section omitted, a job whose GPUs cross broken PCIe/NVLink paths declares
+the NCCL workarounds, and a job on an NVSwitch fabric can pin the NVLink fast
+path on.
+
+Because of that scoping, the effective pattern is one YAML per GPU group
+rather than one shared YAML that gets edited whenever the device list changes:
+
+```yaml
+# train-local.yaml: GPUs with working peer-to-peer; nothing to declare
+runtime:
+  gpu:
+    devices: [0, 1]
+
+# train-cross-pcie.yaml: this GPU set crosses broken paths, so only this job
+# declares the workarounds (confirm first; see docs/guides/distributed.md)
+runtime:
+  gpu:
+    devices: [4, 5, 6, 7]
+  environment:
+    NCCL_P2P_DISABLE: "1"
+    NCCL_NET_GDR_LEVEL: "0"
+```
+
+The same mechanism carries positive tuning, not just workarounds. On an
+NVSwitch node (Hopper-class GPUs with fabric manager running), NVLink SHARP
+multicast (NVLS) is the fast allreduce path and NCCL enables it automatically
+where supported. A job may pin it on explicitly and raise channel parallelism
+when benchmarks show the NVLink bandwidth is underused:
+
+```yaml
+# train-nvlink.yaml: NVSwitch node; keep the disables OUT and pin the fast
+# path on instead (verify support with NCCL_DEBUG=INFO first)
+runtime:
+  gpu:
+    devices: [0, 1, 2, 3]
+  environment:
+    NCCL_NVLS_ENABLE: "1"
+    NCCL_MIN_NCHANNELS: "8"
+    # NCCL_ALGO: NVLS   # force one algorithm; unsupported values fail loudly
+```
+
+NVLS requires NVSwitch multicast support; on plain NVLink bridges or PCIe-only
+sets, keep the section omitted and let NCCL pick Ring/Tree with P2P. Newer
+drivers list the actual interconnect and NVLS support directly in
+`nvidia-smi topo -m`, so check that before assuming.
+
+Confirm a variable is needed before adding it, and only in the YAML of the job
+that hits the problem:
+
+```bash
+nvidia-smi topo -m    # check P2P support between exactly the selected GPUs
+NCCL_DEBUG=INFO       # confirm NCCL transport errors before disabling them
+```
+
+See `docs/guides/distributed.md` for what each troubleshooting variable
+disables. The two directions are mutually exclusive: `NCCL_P2P_DISABLE` and
+`NCCL_NET_GDR_LEVEL` remove bandwidth and must never appear in the same
+environment as the NVLink entries above.
+
+Semantics:
+
+- Values must be scalars and are rendered with `str()`, so quote them
+  explicitly (`"1"`, `"0"`) instead of relying on YAML booleans or numbers.
+- A `null` value exports the name with an empty value.
+- This section is the only path for extra host variables into the trainer
+  container; variables exported in the host shell do not pass through Compose.
+
 ## Fixed Container Paths
 
 | Runtime path | Container path | Access |
@@ -72,8 +147,8 @@ runtime:
 | the selected YAML | `/run/astrai/train.yaml` | read-only |
 
 Training configuration must therefore use `data_root_path: /data`. The source
-code is baked into `/app`; `start` uses `--build`, so code changes rebuild the
+code is baked into `/app`; `start` reuses the existing image, so run
-image when necessary.
+`bash scripts/train.sh build [CONFIG]` after code changes.
 
 ## Operations
 
@@ -123,5 +198,7 @@ the Docker timeout expires.
 3. Do not force DDP for a model that requires FSDP; declare the mode explicitly.
 4. Do not use `kill -9` for routine shutdown; use `scripts/train.sh stop CONFIG`.
 5. The image user is built with the host UID/GID so mounted checkpoints retain usable ownership.
+6. Scope `runtime.environment` to the job YAML that needs it; do not copy NCCL
+   workarounds into every config.
 
-> Document Update Time: 2026-08-22
+> Document Update Time: 2026-08-29

docs/developer/internals.md

@@ -185,7 +185,7 @@ The extension package separates mechanism from policy:
 Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/backend/attention.py`):
 
 - **`CudaBackend`** (default when supported): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool).
-- **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
+- **`FlashAttnBackend`**: optional flash-attn dispatch; inference paths gather flat K/V from the pool via `req_to_token` and call `flash_attn_varlen_func` over the ragged batch (fp16/bf16 only); dense mask-free training calls use `flash_attn_func`.
 - **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
 - The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call.
 - Resolution precedence is: explicit `attn_backend(...)` context > `ASTR_BACKEND` env > default. An explicit `attn_backend(...)` selection is strict (incompatible calls raise); `ASTR_BACKEND` is a default-level override that falls back to a compatible backend when incapable. Training calls (`fwd=None`, no KV cache) resolve by capability: the CUDA cache kernels cannot run without a cache, so they fall back to flash (mask-free/causal calls only) and finally to torch SDPA.

docs/get-started.md

@@ -190,8 +190,9 @@ python scripts/tools/train.py \
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
+# Only if this host's NCCL transport is broken; see docs/guides/distributed.md:
-export NCCL_NET_GDR_LEVEL=0
+# export NCCL_P2P_DISABLE=1
+# export NCCL_NET_GDR_LEVEL=0
 
 python scripts/tools/train.py \
     --train_type=seq \

scripts/docker/serve_runtime.py

@@ -78,9 +78,10 @@ def load_runtime(config_path: str) -> dict[str, str]:
         raise ValueError("runtime.gpu.enabled must be a boolean")
 
     devices = gpu.get("devices", "all")
+    visible_devices = None
     if gpu_enabled:
         if devices == "all":
-            visible_devices = ""
+            pass
         elif isinstance(devices, list) and len(devices) == 1:
             text = str(devices[0])
             if not text.isdigit():
@@ -93,7 +94,6 @@ def load_runtime(config_path: str) -> dict[str, str]:
                 "runtime.gpu.devices must be 'all' or a single-device list such as [0]"
             )
     else:
-        visible_devices = ""
         if devices != "all":
             raise ValueError(
                 "runtime.gpu.devices is ignored when runtime.gpu.enabled is false"
@@ -110,9 +110,10 @@ def load_runtime(config_path: str) -> dict[str, str]:
         "SERVE_PARAM_DIR": _path(paths.get("param", "./params"), "param", path.parent),
         "SERVE_GPU_ENABLED": "true" if gpu_enabled else "false",
         "SERVE_DEVICE": device,
-        "CUDA_VISIBLE_DEVICES": visible_devices,
         "CUDA_TAG": str(container.get("cuda_tag", "cu128")),
     }
+    if visible_devices is not None:
+        values["CUDA_VISIBLE_DEVICES"] = visible_devices
 
     for name, value in environment.items():
         if not isinstance(name, str) or not ENV_NAME.fullmatch(name):

scripts/docker/train_runtime.py

@@ -53,9 +53,9 @@ def load_runtime(config_path: str) -> dict[str, str]:
         )
 
     devices = gpu.get("devices", "all")
+    visible_devices = None
     if devices == "all":
         gpu_count = "all"
-        visible_devices = ""
     elif isinstance(devices, list) and devices:
         normalized = []
         for device in devices:
@@ -102,7 +102,6 @@ def load_runtime(config_path: str) -> dict[str, str]:
             paths.get("checkpoints"), "checkpoints", path.parent
         ),
         "TRAIN_GPU_COUNT": gpu_count,
-        "CUDA_VISIBLE_DEVICES": visible_devices,
         "TRAIN_PARALLEL_MODE": parallel_mode,
         "CUDA_TAG": str(container.get("cuda_tag", "cu128")),
         "TRAIN_IPC_MODE": str(container.get("ipc", "host")),
@@ -111,6 +110,8 @@ def load_runtime(config_path: str) -> dict[str, str]:
         "CHECKPOINT_KEEP_LAST": str(container.get("checkpoint_keep_last", 5)),
         "TRAIN_MAX_DURATION_SECONDS": str(max_seconds),
     }
+    if visible_devices is not None:
+        values["CUDA_VISIBLE_DEVICES"] = visible_devices
 
     for name, value in environment.items():
         if not isinstance(name, str) or not ENV_NAME.fullmatch(name):

scripts/serve.sh

@@ -46,7 +46,7 @@ load_config() {
         die "PyYAML is required on the host (install python3-yaml)"
 
     local exports
-    exports="$(python3 "${ROOT_DIR}/scripts/tools/serve_runtime.py" exports "${CONFIG_FILE}")" ||
+    exports="$(python3 "${ROOT_DIR}/scripts/docker/serve_runtime.py" exports "${CONFIG_FILE}")" ||
         die "Failed to load runtime configuration"
     eval "${exports}"
     if [[ -n "${SERVE_JOB_NAME}" ]]; then
@@ -55,7 +55,12 @@ load_config() {
 }
 
 compose() {
-    ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
+    if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
+        ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
+    else
+        ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" \
+            env -u CUDA_VISIBLE_DEVICES "${COMPOSE_BASE[@]}" "$@"
+    fi
 }
 
 container_name() {
@@ -108,7 +113,7 @@ runtime_environment_args() {
     local pair
     while IFS= read -r -d '' pair; do
         RUNTIME_ENV_ARGS+=(--env "${pair}")
-    done < <(python3 "${ROOT_DIR}/scripts/tools/serve_runtime.py" environment "${CONFIG_FILE}")
+    done < <(python3 "${ROOT_DIR}/scripts/docker/serve_runtime.py" environment "${CONFIG_FILE}")
 }
 
 start_server() {
@@ -129,11 +134,11 @@ start_server() {
         "${RUNTIME_ENV_ARGS[@]}"
     )
     if [[ "${foreground}" == "true" ]]; then
-        compose "${PROFILE_ARGS[@]}" run --build --rm --service-ports \
+        compose "${PROFILE_ARGS[@]}" run --rm --service-ports \
             "${run_options[@]}" "$(service_name)" \
             python -m scripts.tools.server --config /run/astrai/serve.yaml "$@"
     else
-        compose "${PROFILE_ARGS[@]}" run -d --build --service-ports \
+        compose "${PROFILE_ARGS[@]}" run -d --service-ports \
             --name "${container}" "${run_options[@]}" "$(service_name)" \
             python -m scripts.tools.server --config /run/astrai/serve.yaml "$@"
         log_info "Server started; run scripts/serve.sh logs ${CONFIG_FILE} to follow it"

scripts/train.sh

@@ -51,14 +51,19 @@ load_config() {
         die "PyYAML is required on the host (install python3-yaml)"
 
     local exports
-    exports="$(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" exports "${CONFIG_FILE}")" ||
+    exports="$(python3 "${ROOT_DIR}/scripts/docker/train_runtime.py" exports "${CONFIG_FILE}")" ||
         die "Failed to load runtime configuration"
     eval "${exports}"
     validate_job_name "${TRAIN_JOB_NAME}"
 }
 
 compose() {
-    ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
+    if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
+        ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
+    else
+        ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" \
+            env -u CUDA_VISIBLE_DEVICES "${COMPOSE_BASE[@]}" "$@"
+    fi
 }
 
 checkpoint_dir() {
@@ -143,7 +148,7 @@ runtime_environment_args() {
     local pair
     while IFS= read -r -d '' pair; do
         RUNTIME_ENV_ARGS+=(--env "${pair}")
-    done < <(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" environment "${CONFIG_FILE}")
+    done < <(python3 "${ROOT_DIR}/scripts/docker/train_runtime.py" environment "${CONFIG_FILE}")
 }
 
 start_training() {
@@ -164,9 +169,9 @@ start_training() {
         "${RUNTIME_ENV_ARGS[@]}"
     )
     if [[ "${foreground}" == "true" ]]; then
-        compose run --build --rm "${run_options[@]}" trainer "$@"
+        compose run --rm "${run_options[@]}" trainer "$@"
     else
-        compose run -d --build --name "${container}" "${run_options[@]}" trainer "$@"
+        compose run -d --name "${container}" "${run_options[@]}" trainer "$@"
         schedule_timer
         log_info "Training started; run scripts/train.sh logs ${CONFIG_FILE} to follow it"
     fi

setup.py

@@ -6,7 +6,9 @@ import warnings
 from pathlib import Path
 
 from setuptools import setup
+from setuptools.command.build import build as _build
 from setuptools.command.build_ext import build_ext as _build_ext
+from setuptools.command.editable_wheel import editable_wheel as _editable_wheel
 
 sys.path.insert(0, str(Path(__file__).parent))
 os.makedirs("astrai/extension/lib", exist_ok=True)
@@ -93,10 +95,40 @@ class _CMakeBuildExt(_build_ext):
         if not arch:
             arch = _detect_cuda_arch()
         if arch:
+            try:
+                if int(str(arch)) < 89:
+                    warnings.warn(
+                        f"FP8 operator disabled: CUDA compute capability {arch} "
+                        "requires 89 or newer.",
+                        stacklevel=2,
+                    )
+            except ValueError:
+                warnings.warn(
+                    f"Could not parse ASTRAI_CUDA_ARCH={arch!r}; "
+                    "FP8 capability will be decided by CMake.",
+                    stacklevel=2,
+                )
             cfg.append(f"-DASTRAI_CUDA_ARCH={arch}")
         subprocess.run(cfg, check=True)
         subprocess.run([cmake, "--build", str(build_dir), "-j", parallel], check=True)
 
+        # After compilation finishes, verify mandatory CUDA kernels to confirm build succeeded.
+        # CMake may report partial‑target success even if some architecture‑specific kernels are skipped.
+        # Prevent editable install from reporting success when critical kernel shared objects are missing.
+        lib_dir = src / "astrai" / "extension" / "lib"
+        required = (
+            "attn_decode",
+            "attn_prefill",
+            "attn_paged_decode",
+            "attn_paged_prefill",
+            "rotary_emb",
+        )
+        missing = [name for name in required if not any(lib_dir.glob(f"{name}.*.so"))]
+        if missing:
+            raise RuntimeError(
+                "CUDA build completed without some required kernel modules!"
+            )
+
 
 def _cuda_toolkit_version():
     import shutil
@@ -148,6 +180,24 @@ class _NullBuildExt(_build_ext):
         pass
 
 
+class _Build(_build):
+    """Run the CMake kernel build as part of setuptools' build lifecycle."""
+
+    def run(self):
+        if _should_build():
+            self.run_command("build_ext")
+        super().run()
+
+
+class _EditableWheel(_editable_wheel):
+    """Run the CMake kernel build for PEP 660 editable installations."""
+
+    def run(self):
+        if _should_build():
+            self.run_command("build_ext")
+        super().run()
+
+
 cmdclass = {}
 
 if _should_build():
@@ -155,4 +205,10 @@ if _should_build():
 else:
     cmdclass["build_ext"] = _NullBuildExt
 
-setup(ext_modules=[], cmdclass=cmdclass)
+cmdclass["build"] = _Build
+cmdclass["editable_wheel"] = _EditableWheel
+
+setup(
+    ext_modules=[],
+    cmdclass=cmdclass,
+)

tests/extension/test_fp8_mma.py

@@ -1,9 +1,10 @@
 """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
-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
@@ -14,9 +15,8 @@ import torch.nn.functional as F
 
 import astrai.extension.fp8 as f8mod
 from astrai.extension.fp8 import (
-    DelayedScaling,
-    DynamicScaling,
     FP8Format,
+    FP8Recipe,
     FP8TensorMeta,
     _ScaleRing,
     fp8_autocast,
@@ -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, quantize_dual
-    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,91 +68,178 @@ 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,
+@skip_no_fp8
-            w,
+@pytest.mark.parametrize("bias_on", [False, True])
-            None,
+def test_mm_fp8_fused_bias(bias_on):
-            ring.scale,
+    """Epilogue-fused bias matches the unfused out + bias reference (single
-            sw,
+    fp32 rounding vs the reference's double rounding keeps it within 1 ulp),
-            "e4m3",
+    including N-tail columns and batched broadcast."""
-            None,
+    torch.manual_seed(31)
-            ring.state,
+    m, n, k = 19, 13, 37  # odd n exercises the guarded bias loads
-            idx,
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-            0,
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-        )
+    sa, sb = _scale(a), _scale(b)
-        torch.cuda.synchronize()
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
-        expected_hist = hist0.clone()
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
-        expected_hist[idx] = x.abs().amax().float()
+    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
-        torch.testing.assert_close(ring.hist, expected_hist)
+
-        expected_scale = (expected_hist.max() / 448.0).reshape(1)
+    out = mm_fp8(a8, b8, sa * sb, trans_b=True, bias=bias if bias_on else None)
-        torch.testing.assert_close(ring.scale, expected_scale, rtol=1e-6, atol=1e-12)
+    base = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
-        # counter re-armed to int32 zero
+    expected = base + bias if bias_on else base
-        assert ring.state[-1].view(torch.int32).item() == 0
+    # bias is O(1) against O(sqrt(k)) accumulators: absolute tolerance rules
-        hist0 = expected_hist.clone()
+    torch.testing.assert_close(out, expected, atol=0.13, rtol=0.01)
-        idx = (idx + 1) % 4
+
+    # Batched broadcast: bias applies to every batch slice (each slice gets
+    # its own reference from its own operand values).
+    ab = torch.randn(3, m, k, device="cuda", dtype=torch.bfloat16)
+    ab8, _ = quantize(ab, sa.reciprocal(), "e4m3")
+    outb = mm_fp8(ab8, b8, sa * sb, trans_b=True, bias=bias)
+    assert outb.shape == (3, m, n)
+    for i in range(3):
+        expected_b = (_quantize(ab[i], sa) @ _quantize(b, sb).t() * sa * sb).to(
+            torch.bfloat16
+        ) + bias
+        torch.testing.assert_close(outb[i], expected_b, atol=0.13, rtol=0.01)
+
+
+@skip_no_fp8
+@pytest.mark.parametrize("trans_a", [False, True])
+@pytest.mark.parametrize("trans_b", [False, True])
+def test_mm_fp8_batched(trans_a, trans_b):
+    """3D operands run as one bmm launch: all four layouts, odd shapes."""
+    torch.manual_seed(23)
+    batch, m, n, k = 4, 19, 13, 37
+    a = torch.randn(batch, m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(batch, n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+    a_op = a8.transpose(-2, -1).contiguous() if trans_a else a8
+    b_op = b8 if trans_b else b8.transpose(-2, -1).contiguous()
+
+    out = mm_fp8(a_op, b_op, sa * sb, trans_a=trans_a, trans_b=trans_b)
+    assert out.shape == (batch, m, n)
+    # flags + transposed buffers reconstruct the original operands: the math
+    # is always A_orig @ B_orig^T regardless of the layout combination.
+    expected = (_quantize(a, sa) @ _quantize(b, sb).transpose(-2, -1) * sa * sb).to(
+        torch.bfloat16
+    )
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
+
+
+@skip_no_fp8
+def test_mm_fp8_batched_broadcast():
+    """A size-1 batch broadcasts across the other operand (matmul rules),
+    and a 2D operand broadcasts across a 3D one."""
+    torch.manual_seed(29)
+    batch, m, n, k = 3, 16, 8, 32
+    a = torch.randn(batch, m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(1, n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+
+    out = mm_fp8(a8, b8, sa * sb, trans_b=True)
+    assert out.shape == (batch, m, n)
+    expected = (_quantize(a, sa) @ _quantize(b, sb).transpose(-2, -1) * sa * sb).to(
+        torch.bfloat16
+    )
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
+
+    # 2D weight broadcast over 3D activations
+    w8 = b8[0]
+    out2 = mm_fp8(a8, w8, sa * sb, trans_b=True)
+    assert out2.shape == (batch, m, n)
+    torch.testing.assert_close(out2, expected, atol=0.125, rtol=0.01)
+
+
+@skip_no_fp8
+def test_mm_fp8_col_major_view_zero_copy():
+    """An inner-transposed view (.t() of a contiguous buffer) folds into the
+    layout tag with no device copy — the only allocation is the output."""
+    torch.manual_seed(31)
+    m, n, k = 64, 64, 64
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+
+    torch.cuda.synchronize()
+    before = torch.cuda.memory_allocated()
+    out = mm_fp8(a8.t(), b8, sa * sb, trans_a=True, trans_b=True)
+    torch.cuda.synchronize()
+    grew = torch.cuda.memory_allocated() - before
+    assert grew == out.numel() * out.element_size()  # no operand copy
+
+    expected = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
 
 
 @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()
     state.reset()
-    state.default_recipe = DelayedScaling(history_len=1, margin=0)
+    state.default_recipe = FP8Recipe(history_len=1, margin=0)
     state.default_format = FP8Format.E4M3
     try:
         m, n, k = 32, 16, 64
         x1 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.5
         # Smaller amax than x1: the delayed scale (amax(x1)/448) still covers
-        # x2 without fp8 saturation, while the next-step scale would differ —
+        # 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 +262,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 = FP8Recipe(dynamic=True)
-    )
+    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 +328,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 +368,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 +391,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)
 # --------------------------------------------------------------------------
@@ -374,13 +399,13 @@ def test_mm_fp8_fp8_output():
 def test_recipe_scale_from_history():
     """Delayed: max over the window + margin; dynamic: current amax."""
     hist = torch.tensor([1.0, 2.0, 0.5])
-    d = DelayedScaling(history_len=3, margin=0)
+    d = FP8Recipe(history_len=3, margin=0)
     assert torch.allclose(d.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0))
-    d_m = DelayedScaling(history_len=3, margin=2)
+    d_m = FP8Recipe(history_len=3, margin=2)
     assert torch.allclose(
         d_m.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0 / 4.0)
     )
-    dyn = DynamicScaling()
+    dyn = FP8Recipe(dynamic=True)
     amax = torch.tensor([0.25])
     assert torch.allclose(
         dyn.scale_from_history(amax, "e4m3"), torch.tensor(0.25 / 448.0)
@@ -398,73 +423,107 @@ def test_fp8_format_enum():
 
 
 def test_fp8_autocast_context():
-    """fp8_autocast sets and restores recipe + format on the global state."""
+    """fp8_autocast pushes and restores the thread-local active config."""
     state = fp8_state()
-    prev = (state.enabled, state.recipe, state.fp8_format)
+    state.reset()
     try:
         with fp8_autocast(enabled=True, fp8_format="hybrid", update_interval=8):
-            assert state.enabled
+            cfg = f8mod._active_config.get()
-            assert isinstance(state.recipe, DelayedScaling)
+            assert cfg is not None and cfg.enabled
-            assert state.recipe.history_len == 8
+            assert not cfg.recipe.dynamic
-            assert state.fp8_format is FP8Format.HYBRID
+            assert cfg.recipe.history_len == 8
-            with fp8_autocast(enabled=True, recipe=DynamicScaling(), fp8_format="e4m3"):
+            assert cfg.fp8_format is FP8Format.HYBRID
-                assert isinstance(state.recipe, DynamicScaling)
+            with fp8_autocast(
-                assert state.fp8_format is FP8Format.E4M3
+                enabled=True, recipe=FP8Recipe(dynamic=True), fp8_format="e4m3"
-            assert state.fp8_format is FP8Format.HYBRID  # restored on exit
+            ):
-        assert not state.enabled
+                inner = f8mod._active_config.get()
+                assert inner.recipe.dynamic
+                assert inner.fp8_format is FP8Format.E4M3
+            assert f8mod._active_config.get() is cfg  # restored on exit
+        assert f8mod._active_config.get() is None
+        assert not fp8_linear_enabled()
     finally:
-        state.enabled, state.recipe, state.fp8_format = prev
+        state.reset()
 
 
 def test_fp8_tensor_meta_delayed_update():
     """Meta seeds from data; hist/scale are packed views of one state buffer."""
-    meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0))
+    recipe = FP8Recipe(history_len=4, margin=0)
+    meta = FP8TensorMeta(
+        _ScaleRing(torch.device("cpu"), recipe),
+        _ScaleRing(torch.device("cpu"), recipe),
+        _ScaleRing(torch.device("cpu"), recipe),
+    )
     w = torch.randn(8, 8)
     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 | 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()
-    # 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
 
+    # fold_args hands the kernel the buffer, the slot and the recipe constants
+    args = meta.w.fold_args("e4m3")
+    assert args["ring_state"] is meta.w.state and args["hist_idx"] == 1
+    assert args["fp8_max"] == 448.0 and args["pow2_margin"] == 1.0
 
-def test_quantize_bf16_cpu_fallback():
+
-    """CPU fallback of the quantize primitive (scale semantics + amax)."""
+@skip_no_fp8
-    x = torch.randn(16, 32, dtype=torch.bfloat16)
+@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"])
-    scale = torch.tensor([0.5])
+def test_quantize_dual_and_transposed_orientations(fmt):
-    x8, amax = quantize_bf16(x, scale, "e4m3")
+    """quantize_dual yields both orientations from one read; quantize's
-    assert x8.dtype == torch.float8_e4m3fn
+    transposed switch keeps the 2-tuple arity with the [cols][rows] layout."""
-    ref = (x.float() / 0.5).to(torch.float8_e4m3fn)
+    torch.manual_seed(11)
-    assert torch.equal(x8, ref)
+    x = torch.randn(37, 67, device="cuda", dtype=torch.bfloat16) * 3
+    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])
-    sa = torch.tensor([2.0])
+def test_quantize_ring_fold_matches_host_update(fmt, fmax, margin):
-    sb = torch.tensor([0.5])
+    """The in-kernel delayed-scaling fold matches a host-side reference."""
-    out = mm_fp8(a8, b8, sa, sb)
+    dev = torch.device("cuda")
-    ref = (a8.float() @ b8.float() * 2.0 * 0.5).to(torch.bfloat16)
+    n, idx = 4, 2
-    torch.testing.assert_close(out, ref)
+    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)
 
-def test_mm_fp8_fp8_output_cpu():
+    # Fused: same window, fold inside the quantize kernel's last block.
-    """CPU fallback with an FP8 output (out_dtype='e4m3' + out_scale)."""
+    ring = torch.zeros(n + 4, device=dev)
-    a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
+    ring[:n].fill_(1.0)
-    b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
+    x8, _ = quantize(
-    sa = torch.tensor([2.0])
+        x,
-    sb = torch.tensor([0.5])
+        mult,
-    os_ = torch.tensor([0.25])
+        fmt,
-    out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_)
+        ring_state=ring,
-    assert out8.dtype == torch.float8_e4m3fn
+        hist_idx=idx,
-    ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn)
+        fp8_max=fmax,
-    assert torch.equal(out8, ref)
+        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
 
 
 # --------------------------------------------------------------------------
@@ -488,7 +547,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"

tests/inference/test_scheduler.py

@@ -393,7 +393,9 @@ def test_decode_does_not_reuse_previous_batch_state():
     old_info = object()
     new_info = object()
     executor._decode_cache = DecodeSteadyState(("old",), [2], old_info)
-    executor._sample_logits = MagicMock(return_value=[3])
+    executor._sample_logits = MagicMock(
+        return_value=([3], torch.tensor([3], dtype=torch.long))
+    )
 
     task = Task("new", list(range(8)), temperature=0)
     task.input_tokens = 8
@@ -412,3 +414,46 @@ def test_decode_does_not_reuse_previous_batch_state():
     args, kwargs = executor._sample_logits.call_args
     assert args[1:] == ([task], False)
     assert kwargs["info"] is new_info
+
+
+def test_decode_fills_input_ids_from_device_on_matching_signature():
+    """Steady-state decode copies cached device tokens, skipping the host."""
+    executor = object.__new__(Executor)
+    executor.device = torch.device("cpu")
+    executor.task_cache = MagicMock()
+    executor.task_cache.bind_was_steady = True
+    executor.task_cache.bind.return_value = MagicMock()
+    executor._graph_supported = False
+    executor._graph_ctx = SimpleNamespace(enabled=False)
+
+    workspace = MagicMock()
+    workspace.position_ids = torch.tensor([2], dtype=torch.long)
+    workspace.fill_input_ids_from_device.return_value = torch.tensor(
+        [9], dtype=torch.long
+    )
+    executor._workspace = workspace
+    executor.model = MagicMock(
+        return_value={"logits": torch.zeros(1, 1, 10, dtype=torch.float32)}
+    )
+
+    info = object()
+    tokens = torch.tensor([3], dtype=torch.long)
+    executor._decode_cache = DecodeSteadyState(("t1",), [2], info, last_tokens=tokens)
+    executor._sample_logits = MagicMock(return_value=([3], tokens))
+
+    task = Task("t1", list(range(8)), temperature=0)
+    task.input_tokens = 8
+    task.output_ids = [7]
+    task.mark_prefill_done()
+
+    with patch(
+        "astrai.inference.runtime.executor._build_sampling_batch_info",
+        return_value=info,
+    ):
+        assert executor.execute_decode([task]) == [3]
+
+    workspace.fill_input_ids.assert_not_called()
+    workspace.fill_input_ids_from_device.assert_called_once_with(tokens)
+    assert workspace.position_ids.tolist() == [3]
+    assert executor._decode_cache.task_sig == ("t1",)
+    assert executor._decode_cache.last_tokens is tokens

tests/test_serve_runtime.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from scripts.tools.serve_runtime import load_runtime
+from scripts.docker.serve_runtime import load_runtime
 
 
 def _write(tmp_path, body: str) -> str:
@@ -26,7 +26,7 @@ def test_runtime_exports_defaults(tmp_path):
     assert runtime["SERVE_CONTAINER_PORT"] == "8000"
     assert runtime["SERVE_PARAM_DIR"] == str((tmp_path / "params").resolve())
     assert runtime["SERVE_GPU_ENABLED"] == "true"
-    assert runtime["CUDA_VISIBLE_DEVICES"] == ""
+    assert "CUDA_VISIBLE_DEVICES" not in runtime
     assert runtime["SERVE_DEVICE"] == "cuda"
     assert runtime["CUDA_TAG"] == "cu128"
     assert runtime["SERVE_JOB_NAME"] == ""