perf: overlap direct crosswise loads with fp8 gemm mma

- Rotate kStages+1 canonical buffers for direct-crosswise operands so their load issues right after barrier 1 and its LDG+PRMT latency overlaps the MMA phase, instead of stalling the post-compute inter-barrier window (ncu on the production 24L/dim1536/ffn6912 shapes: barrier 3.7-4.1 + long-scoreboard 1.6-1.8 stalls per issue before).
- Split the stage loaders into load_async (cp.async operands, committed after the post-compute barrier) and load_direct (ring-indexed by tile, not stage, since the rings differ in depth).

On the production shapes vs cuBLASLt _scaled_mm: dX 1.11-1.24x (from 1.27-1.42x), dW 1.32-1.35x (from 1.26-1.49x); generic 4096/11008 shapes also improve (dX 37.6->40.8 TF, dW 36.0->37.3 TF). fp8 train step vs bf16: ~1.2x at both 512 and 2048 tokens.

astrai/extension/fp8.py

@@ -1,35 +1,37 @@
 """FP8 training: scaling recipes, per-tensor state, and aten::linear dispatch.
 
-Layered (see also ``ops/fp8.py`` for the CUDA interface adapter):
-
-1. Kernel interface: ``ops.fp8`` — the only module touching the pybind.
-2. Training state (this module): scaling *recipes* (TE-style delayed scaling
-   or dynamic current-amax scaling), per-tensor scales + amax history, and
-   the ``fp8_autocast`` context (like ``torch.autocast``).
-3. aten::linear integration (this module): registers the CUDA impl and the
-   dtype guard.
+Layered (see ``ops/fp8.py`` for the CUDA interface adapter):
+1. ``ops.fp8`` — the only module touching the pybind.
+2. This module (strategy layer): scaling *recipes* (TE-style delayed scaling
+   or dynamic current-amax scaling), per-tensor scales + amax history, and the
+   ``fp8_autocast`` context manager (like ``torch.autocast``).
+3. aten::linear integration: registers the CUDA + AutogradCUDA impls.
 
 Usage::
 
     from astrai.extension.fp8 import fp8_autocast
-
     with fp8_autocast(enabled=True, fp8_format="hybrid"):
         logits = model(input_ids)
-    loss.backward()  # fp8 backward runs wherever it is called: the
-    # forward captures the fmt/recipe/meta on the autograd node
+    loss.backward()  # fp8 backward runs anywhere; fwd captured state on the node
 
-Importing this module registers the aten::linear CUDA and AutogradCUDA
-implementations.
+Format defaults follow the ecosystem consensus: E4M3 forward / E5M2 backward
+("hybrid"); every operand's scale is a quantization step derived from its amax
+history by the active recipe.
 
-Format defaults follow the ecosystem consensus: E4M3 for the forward pass,
-E5M2 for the backward (gradient) pass ("hybrid"); every operand's scale is a
-quantization step derived from its amax history by the active recipe.
+The context mirrors ``torch.autocast`` (``autocast_mode.py``): the active
+``(enabled, recipe, fp8_format)`` triple is thread-local (a ``contextvars``
+``ContextVar``, absent outside any region), and the manager is class-based and
+reentrant with nested ``enabled=False`` disabling dispatch inside it. The module
+targets *training*: every step quantizes x/w/g fresh (no weight-cast cache — the
+optimizer bumps the weight version each step, so a torch-style cached_cast would
+miss anyway), and the per-operand scales come from the delayed/dynamic recipe.
 """
 
-from contextlib import contextmanager
+import functools
+from contextvars import ContextVar, Token
 from dataclasses import dataclass
 from enum import Enum
-from typing import Optional
+from typing import Dict, List, Optional
 
 import torch
 from torch.library import Library
@@ -58,73 +60,62 @@ class FP8Format(str, Enum):
 
 
 class FP8Recipe:
-    """Scale-from-amax policy; the scale computation is the injection point.
+    """Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
 
-    ``scale_from_history`` receives the amax tensor for this operand (a ring
-    window for delayed scaling, the current amax for dynamic scaling) and
-    returns the quantization step: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
+    ``scale_from_history`` receives the operand's amax tensor (a ring window for
+    delayed scaling, the current amax for dynamic scaling) and returns the
+    quantization step. Subclasses set ``history_len`` / ``margin``.
     """
 
     history_len: int = 16
     margin: int = 0
 
     def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
-        raise NotImplementedError
+        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.
-
-    The scale is computed from amax measured in *previous* steps (delayed one
-    step); the window length trades responsiveness against stability.
-    """
+    """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
 
-    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 DynamicScaling(FP8Recipe):
-    """Current-amax scaling (torchao DYNAMIC): measure, then quantize.
-
-    No history — the scale is derived from the amax of the tensor being
-    quantized in the same step, at the cost of an extra reduction pass.
-    """
+    """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
 
-    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)
-
 
 class _ScaleRing:
-    """One operand's delayed-scaling state: amax history ring + derived scale.
-
-    The ring captures its recipe at construction; ``update`` records a fresh
-    amax and refreshes the scale for the *next* step (delayed one step).
+    """One operand's delayed-scaling state: a float32 buffer
+    ``[hist[n] | scale | counter]`` (views). The quantize kernel's last-finishing
+    block records the measured amax into ``hist[idx]``, reduces the window and
+    publishes the next scale entirely on device — the Python-side write/max/write
+    chain is gone. The counter slot stays int32-zero (float bits) between
+    launches; ``idx`` advances host-side each step.
     """
 
-    __slots__ = ("recipe", "hist", "idx", "scale", "initialized")
+    __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
 
     def __init__(self, device: torch.device, recipe: FP8Recipe):
         self.recipe = recipe
         n = recipe.history_len
-        self.hist = torch.ones(n, device=device, dtype=torch.float32)
+        self.state = torch.zeros(n + 2, device=device, dtype=torch.float32)
+        self.hist = self.state[:n]
+        self.scale = self.state[n : n + 1]
         self.idx = 0
-        self.scale = torch.ones(1, device=device, dtype=torch.float32)
         self.initialized = False
 
-    def update(self, amax: torch.Tensor, fmt: str) -> None:
-        self.hist[self.idx] = amax.reshape(())
+    def advance(self) -> None:
+        """Rotate to the next history slot after an in-kernel finalize."""
         self.idx = (self.idx + 1) % self.hist.numel()
-        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
 
     def seed(self, t: torch.Tensor, fmt: str) -> None:
         amax = t.abs().amax().to(torch.float32).clamp_min(1e-12)
@@ -134,12 +125,10 @@ class _ScaleRing:
 
 
 class FP8TensorMeta:
-    """Per-weight delayed-scaling state: one ring per operand role.
-
-    Holds the ``w`` / ``x`` / ``g`` rings; fused kernels record the amax
-    while quantizing, so the scale used at step N reflects amax from steps
-    < N. DynamicScaling never allocates a meta — it measures the current
-    amax inline (``_dynamic_scale``), so it needs no history storage.
+    """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/
+    ``g``). Fused kernels record amax while quantizing, so the scale used at step
+    N reflects amax from steps < N. DynamicScaling never allocates a meta — it
+    measures the current amax inline.
     """
 
     __slots__ = ("w", "x", "g")
@@ -150,14 +139,70 @@ class FP8TensorMeta:
         self.g = _ScaleRing(device, recipe)
 
 
+@dataclass(frozen=True)
+class _ActiveConfig:
+    """The immutable (enabled, recipe, format) triple of one open region."""
+
+    enabled: bool
+    recipe: FP8Recipe
+    fp8_format: FP8Format
+
+
+# Thread-local active configuration (torch's autocast TLS analog): set by
+# fp8_autocast on __enter__, absent outside any region. Autograd engine
+# threads run backwards with their own empty context — fine, since backward
+# only reads state captured on ctx at forward time.
+_active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
+    "astrai_fp8_active_config", default=None
+)
+
+
 class FP8State:
-    """Global fp8 training state: active recipe + per-tensor metas."""
+    """Global fp8 training state: per-tensor metas + out-of-region defaults.
+
+    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set
+    by ``fp8_autocast``. The properties below read that active config when a
+    region is open and the global defaults otherwise; the setters (and
+    ``fp8_linear_enable``) write the global defaults — the persistent switch
+    applying outside any region. The metas registry is shared across threads
+    (GIL-protected); fp8 backward runs on autograd engine threads and only
+    touches metas captured on ``ctx`` at forward time.
+    """
 
     def __init__(self):
-        self.enabled = False
-        self.recipe: FP8Recipe = DelayedScaling()
-        self.fp8_format: FP8Format = FP8Format.HYBRID
-        self._metas: dict[tuple, FP8TensorMeta] = {}
+        self.default_enabled = False
+        self.default_recipe: FP8Recipe = DelayedScaling()
+        self.default_format: FP8Format = FP8Format.HYBRID
+        self._metas: Dict[tuple, FP8TensorMeta] = {}
+
+    # Active-config views (region config if open, else the defaults).
+    @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)
@@ -168,13 +213,11 @@ class FP8State:
         return meta
 
     def reset(self) -> None:
-        self.enabled = False
+        self.default_enabled = False
         self._metas.clear()
 
 
-# Global singleton: autograd backward runs on the engine worker threads, so
-# thread-local state would lose the fp8 flag during loss.backward(). The GIL
-# protects Python-side mutation; the CUDA kernels take their own mutex.
+# Process-wide singleton; per-thread/per-region state lives in _active_config.
 _state = FP8State()
 
 
@@ -182,43 +225,76 @@ def fp8_state() -> FP8State:
     return _state
 
 
-@contextmanager
-def fp8_autocast(
-    enabled: bool = True,
-    update_interval: int = 16,
-    recipe: Optional[FP8Recipe] = None,
-    fp8_format: str = "hybrid",
-    margin: int = 0,
-):
+def _active() -> Optional[_ActiveConfig]:
+    """The active config when fp8 dispatch is on, else ``None`` (fast guard).
+
+    A region config wins (honoring nested ``enabled=False`` regions); with no
+    region open this falls back to the persistent global switch
+    (``fp8_linear_enable``), so that flag still routes aten::linear to fp8.
+    """
+    cfg = _active_config.get()
+    if cfg is not None:
+        return cfg if cfg.enabled else None
+    if _state.default_enabled:
+        return _ActiveConfig(True, _state.default_recipe, _state.default_format)
+    return None
+
+
+def _current_config() -> _ActiveConfig:
+    """Like ``_active()`` but always returns a config (disabled regions and
+    out-of-region direct calls resolve to the global defaults)."""
+    cfg = _active_config.get()
+    if cfg is not None:
+        return cfg
+    return _ActiveConfig(
+        _state.default_enabled, _state.default_recipe, _state.default_format
+    )
+
+
+class fp8_autocast:
     """Autocast-style context: fp8 linear dispatch on this thread.
 
-    Usage::
+    Mirrors ``torch.autocast`` — a class-based, reentrant, nestable context
+    over thread-local state::
 
         with fp8_autocast(enabled=True, fp8_format="hybrid"):
             logits = model(input_ids)   # aten::linear -> fp8 path
         loss.backward()  # fp8 backward; state was captured at forward time
 
-    Args:
-        enabled: toggle fp8 dispatch for aten::linear.
-        update_interval: legacy alias for the delayed-scaling history window
-            (used only when ``recipe`` is not given).
-        recipe: scaling policy; defaults to ``DelayedScaling(update_interval)``.
-        fp8_format: ``"e4m3"`` / ``"e5m2"`` / ``"hybrid"`` (default) — hybrid
-            means E4M3 forward, E5M2 backward.
-        margin: scale headroom (``scale = (amax / FP8_MAX) / 2^margin``) used
-            with the default delayed recipe.
+    Nesting follows torch: each ``__enter__`` pushes the new active config, each
+    ``__exit__`` restores the previous one, and a nested ``enabled=False`` region
+    simply disables dispatch inside it. The instance doubles as a decorator.
     """
-    state = fp8_state()
-    prev = (state.enabled, state.recipe, state.fp8_format)
-    if recipe is None:
-        recipe = DelayedScaling(history_len=update_interval, margin=margin)
-    state.enabled = enabled
-    state.recipe = recipe
-    state.fp8_format = FP8Format(fp8_format)
-    try:
-        yield
-    finally:
-        state.enabled, state.recipe, state.fp8_format = prev
+
+    def __init__(
+        self,
+        enabled: bool = True,
+        update_interval: int = 16,
+        recipe: Optional[FP8Recipe] = None,
+        fp8_format: str = "hybrid",
+        margin: int = 0,
+    ):
+        if recipe is None:
+            recipe = DelayedScaling(history_len=update_interval, margin=margin)
+        self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
+        self._tokens: List[Token] = []
+
+    def __enter__(self) -> "fp8_autocast":
+        self._tokens.append(_active_config.set(self._config))
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        token = self._tokens.pop()
+        _active_config.reset(token)
+        return False
+
+    def __call__(self, func):
+        @functools.wraps(func)
+        def decorate(*args, **kwargs):
+            with self:
+                return func(*args, **kwargs)
+
+        return decorate
 
 
 # ---------------------------------------------------------------------------
@@ -231,54 +307,96 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor
     return recipe.scale_from_history(amax, fmt)
 
 
-def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
+_zero_bias: Dict[Optional[int], torch.Tensor] = {}
+
+
+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(
+    x: torch.Tensor, w: torch.Tensor, bias=None, cfg: Optional[_ActiveConfig] = None
+):
     """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 pre-quantized GEMM, and feed the freshly measured amax back into the
-    delayed-scaling ring (dynamic scaling measures the current amax itself).
+    Pure FP8 path for both recipes: quantize x/w with the active scales, run the
+    pre-quantized GEMM. Delayed scaling finalizes the rings inside the quantize
+    kernels (amax folded into the window, next scale published on device);
+    dynamic scaling measures the current amax itself. Training quantizes the
+    weight every step (the optimizer bumps its version, so there is no cast
+    cache, matching ``cached_cast``-less behavior).
     """
-    if bias is None:
-        bias = torch.empty(0, device=x.device, dtype=x.dtype)
     state = fp8_state()
-    fmt = state.fp8_format.fwd()
-    if isinstance(state.recipe, DynamicScaling):
-        meta = None
-        sx = _dynamic_scale(x.reshape(-1, w.size(1)), state.recipe, fmt)
-        sw = _dynamic_scale(w, state.recipe, fmt)
+    if cfg is None:
+        cfg = _current_config()
+    fmt = cfg.fp8_format.fwd()
+    margin = cfg.recipe.margin
+    if bias is None:
+        bias = _empty_bias(x)
+    if isinstance(cfg.recipe, DynamicScaling):  # measure-then-quantize, no state
+        sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
+        sw = _dynamic_scale(w, cfg.recipe, fmt)
+        out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt)
+        return out
+
+    meta = state.get_weight_meta(w)
+    if not meta.w.initialized:
+        meta.w.seed(w, fmt)
+    if not meta.x.initialized:
+        meta.x.seed(x, fmt)
+    # The kernel finalizes each ring in-kernel (overwriting the scale slot), so
+    # the w/x scales are taken from the ring before the quantize.
+    if w.dtype is not torch.bfloat16:  # static pre-quantized weight
+        w_arg, sw_arg, w_ring = w, meta.w.scale, None
     else:
-        meta = state.get_weight_meta(w)
-        if not meta.w.initialized:
-            meta.w.seed(w, fmt)
-        if not meta.x.initialized:
-            meta.x.seed(x, fmt)
-        sx, sw = meta.x.scale, meta.w.scale
-    out, amax_x, amax_w = linear_forward_fp8(x, w, bias, sx, sw, fmt)
-    if meta is not None:
-        meta.x.update(amax_x, fmt)
-        meta.w.update(amax_w, fmt)
+        w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state
+    out, _x8, _w8, _ax, _aw = linear_forward_fp8(
+        x,
+        w_arg,
+        bias,
+        meta.x.scale,
+        sw_arg,
+        fmt,
+        None,
+        meta.x.state,
+        meta.x.idx,
+        margin,
+        w_ring,
+        meta.w.idx,
+        margin,
+    )
+    meta.x.advance()
+    if w_ring is not None:
+        meta.w.advance()
     return out
 
 
 class _LinearFp8(torch.autograd.Function):
     """The fp8 linear forward/backward pair (standard Function style).
 
-    The forward runs inside the ``fp8_autocast`` region and captures the
-    active fmt/recipe/meta on ``ctx``; the backward reads only the captured
-    state, so ``loss.backward()`` may run after the context exits. The
-    gradient is quantized once (E5M2 in hybrid mode) and the dX / dW GEMMs
-    share that quantization; output masks come from ``needs_input_grad``.
+    The forward runs inside ``fp8_autocast`` and captures the active
+    fmt/recipe/meta on ``ctx``; the backward reads only that captured state, so
+    ``loss.backward()`` may run after the context exits. The gradient is
+    quantized once (E5M2 in hybrid) and both dX/dW GEMMs share it; the output
+    masks come from ``needs_input_grad``.
     """
 
     @staticmethod
     def forward(ctx, x, w, bias):
-        out = fp8_linear_forward(x, w, bias)
-        state = fp8_state()
+        cfg = _current_config()
+        out = fp8_linear_forward(x, w, bias, cfg)
         ctx.save_for_backward(x, w)
-        ctx.fmt_bwd = state.fp8_format.bwd()
-        ctx.recipe = state.recipe
-        ctx.is_dynamic = isinstance(state.recipe, DynamicScaling)
-        ctx.meta = None if ctx.is_dynamic else state.get_weight_meta(w)
+        ctx.fmt_bwd = cfg.fp8_format.bwd()
+        ctx.recipe = cfg.recipe
+        ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
+        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w)
         return out
 
     @staticmethod
@@ -286,22 +404,34 @@ class _LinearFp8(torch.autograd.Function):
     def backward(ctx, g):
         x, w = ctx.saved_tensors
         fmt = ctx.fmt_bwd
+        # Per-recipe scale/ring selection; both branches share one call below.
         if ctx.is_dynamic:
             sg = _dynamic_scale(g, ctx.recipe, fmt)
             sw = _dynamic_scale(w, ctx.recipe, fmt)
             sx = _dynamic_scale(x, ctx.recipe, fmt)
+            ring, idx = None, 0
         else:
             meta = ctx.meta
             if not meta.g.initialized:
                 meta.g.seed(g, fmt)
-            sg, sw, sx = meta.g.scale, meta.w.scale, meta.x.scale
-        masks = list(ctx.needs_input_grad)
-        grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
-            g, x, w, masks, sg, sw, sx, fmt
+            sg, ring, idx = meta.g.scale, meta.g.state, meta.g.idx
+            sw, sx = meta.w.scale, meta.x.scale
+        grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8(
+            g,
+            x,
+            w,
+            list(ctx.needs_input_grad),
+            sg,
+            sw,
+            sx,
+            fmt,
+            ring,
+            idx,
+            ctx.recipe.margin,
         )
         if not ctx.is_dynamic:
-            ctx.meta.g.update(amax_g, fmt)
-        return grad_x, grad_w, grad_b if masks[2] else None
+            meta.g.advance()  # the g quantize kernel finalized the ring in-kernel
+        return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
 
 
 # ---------------------------------------------------------------------------
@@ -310,31 +440,29 @@ class _LinearFp8(torch.autograd.Function):
 
 
 def fp8_linear_enable(enabled: bool = True) -> None:
-    """Toggle fp8 dispatch for aten::linear (global; backward runs on engine
-    worker threads, so a thread-local flag would be lost during backward)."""
-    fp8_state().enabled = enabled
+    """Toggle fp8 dispatch for aten::linear globally (the out-of-region default;
+    ``fp8_autocast`` regions override it thread-locally)."""
+    fp8_state().default_enabled = enabled
 
 
 def fp8_linear_enabled() -> bool:
-    return fp8_state().enabled
+    """Whether fp8 dispatch is active right now (region config or global)."""
+    return _active() is not None
 
 
 def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool:
-    """Shape guard for the fp8 linear path.
-
-    Unlike a strict 16-alignment requirement, the fp8 kernels handle unaligned
-    M/N via boundary checks (slower but correct) — so no whole-call bf16
-    fallback for small decode batches. Only the K-dimension contraction must
-    match, and the weight must be 2D.
-    """
+    """Shape guard for the fp8 path. Unlike a strict 16-alignment requirement,
+    the kernels handle unaligned M/N via boundary checks (slower but correct) —
+    so no whole-call bf16 fallback for small decode batches. Only the K-dimension
+    contraction must match and the weight must be 2D."""
     return x.dim() >= 2 and w.dim() == 2 and x.size(-1) == w.size(1)
 
 
 def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
     if (
-        fp8_linear_enabled()
-        and x.dtype == torch.bfloat16
-        and w.dtype == torch.bfloat16
+        _active() is not None
+        and x.dtype is torch.bfloat16
+        and w.dtype is torch.bfloat16
         and _fp8_supported(x, w)
     ):
         return _LinearFp8.apply(x, w, bias)
@@ -349,9 +477,8 @@ def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
 _lib = Library("aten", "IMPL", "CUDA")
 _lib.impl("linear", _linear_cuda_impl)
 # Also replace torch's generated linear autograd formula (which would call
-# aten::linear_backward after the fp8_autocast region exits). The fp8
-# backward is owned by _LinearFp8 with its state captured at forward time,
-# so loss.backward() works wherever it is called; the same CUDA registration
-# still covers inference_mode, where autograd keys are skipped entirely.
+# aten::linear_backward after the fp8_autocast region exits). The fp8 backward
+# is owned by _LinearFp8 with state captured at forward time, so loss.backward()
+# works wherever it is called; the CUDA registration still covers inference_mode.
 _lib_autograd = Library("aten", "IMPL", "AutogradCUDA")
 _lib_autograd.impl("linear", _linear_cuda_impl)

astrai/extension/ops/fp8.py

@@ -4,7 +4,7 @@ Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
 
 - ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → FP8 with fused amax
 - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
-- ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, amax_x, amax_w)``
+- ``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
@@ -16,6 +16,8 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
 this module is stateless.
 """
 
+from typing import List, Optional, Tuple
+
 import torch
 from torch.library import custom_op
 
@@ -39,7 +41,7 @@ def _fmt_dtype(fmt: str) -> torch.dtype:
 @custom_op("custom::fp8_quantize", mutates_args=())
 def fp8_quantize(
     x: torch.Tensor, scale: torch.Tensor, fmt: int
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``."""
 
 
@@ -73,7 +75,7 @@ def fp8_gemm(
     sa: torch.Tensor,
     sb: torch.Tensor,
     out_dtype: int = 0,
-    out_scale: torch.Tensor | None = None,
+    out_scale: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation.
 
@@ -106,7 +108,9 @@ def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None):
     return acc.to(torch.bfloat16)
 
 
-def quantize_bf16(x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"):
+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
@@ -122,7 +126,7 @@ def mm_fp8(
     sa: torch.Tensor,
     sb: torch.Tensor,
     out_dtype: str = "bf16",
-    out_scale: torch.Tensor | None = None,
+    out_scale: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``.
 
@@ -138,14 +142,36 @@ def mm_fp8(
     return fp8_gemm(a, b, sa, sb, int(out_dtype == "e4m3"), out_scale)
 
 
-def linear_forward_fp8(x, w, bias, sx, sw, fmt: str = "e4m3", bias_scale=None):
+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, amax_x, amax_w)``. ``bias`` may be ``None``. For static
-    fp8 inference, ``w`` and ``bias`` may arrive pre-quantized to ``fmt``
+    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):
@@ -155,16 +181,42 @@ def linear_forward_fp8(x, w, bias, sx, sw, fmt: str = "e4m3", bias_scale=None):
     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,
+        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, x, w, masks, sg, sw, sx, fmt: str = "e5m2"):
+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)``.
 
     The gradient (and the transposed w/x operands) are quantized to ``fmt``
     (default E5M2 — larger dynamic range for gradients) and 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`).
     """
     if not (
         g.dtype == torch.bfloat16
@@ -175,5 +227,15 @@ def linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt: str = "e5m2"):
             f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
         )
     return get_module("fp8_ops").linear_backward_fp8(
-        g, x, w, list(masks), sg, sw, sx, _fmt_int(fmt)
+        g,
+        x,
+        w,
+        list(masks),
+        sg,
+        sw,
+        sx,
+        _fmt_int(fmt),
+        g_ring,
+        g_ring_idx,
+        g_ring_margin,
     )

csrc/kernels/fp8/common.h

@@ -61,18 +61,55 @@ struct Fp8GemmTraits {
     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.
+    // The shared-memory budget is layout-aware (crosswise operands add K-
+    // major staging + a canonical buffer), so it lives in Fp8GemmSmem in
+    // gemm.cuh together with the resident-CTA hint for __launch_bounds__.
+    static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
+};
+
+// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional
+// delayed-scaling ring finalization. Separate from FP8Params so each
+// operator owns exactly the fields it touches (the GEMM never reads amax /
+// ring state). Same NSDMI rationale: amax / ring_state gate optional paths
+// via null checks. Still an aggregate, still trivially copyable.
+struct FP8QuantizeParams {
+    // BF16 input and FP8 output buffers; scale_a is the quantization step
+    // (device scalar). amax_a (may be null) is zero-initialized by the
+    // binding and receives the raw-domain absolute maximum.
+    const void* __restrict__ a_ptr = nullptr;
+    void* __restrict__ out_ptr = nullptr;
+    const float* __restrict__ scale_a = nullptr;
+    float* __restrict__ amax_a = nullptr;
+
+    // Optional delayed-scaling ring finalization. ring_state packs
+    // [hist[ring_len] | scale | counter] with ring_len = numel - 2. When
+    // non-null and amax_a is set, the last-finishing block records the
+    // measured amax into hist[ring_idx], reduces the window and publishes
+    // the next step's scale (max(hist) / fp8_max / 2^ring_margin) — the
+    // fused replacement for the eager hist-write / max / scale-write chain,
+    // at zero extra launches. The counter slot is a persistent zero-armed
+    // int32 (float bits) electing the last block each launch.
+    float* ring_state = nullptr;
+    int ring_len = 0;
+    int ring_idx = 0;
+    int ring_margin = 0;
+
+    // Element count (only the elementwise quantize kernel uses it).
+    int total = 0;
 };
 
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
-// through quantize / fused / pre-quantized kernels. Each kernel touches only
-// the fields it needs; buffers are raw pointers packed by the torch binding.
+// through the pre-quantized GEMM kernels. Each kernel touches only the
+// fields it needs; buffers are raw pointers packed by the torch binding.
 // Pointer members default to null (same NSDMI rationale as AttentionParams:
-// bias / amax / out_scale gate optional paths via null checks, so a partially
+// 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 BF16 for the fused (quantize-in-GEMM) path, FP8 for
-    // the pre-quantized path. Scales are quantization steps (device scalars).
+    // Inputs: a/b are FP8 for the pre-quantized path. Scales are
+    // quantization steps (device scalars).
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
     const void* __restrict__ bias = nullptr;
@@ -84,23 +121,16 @@ struct FP8Params {
     void* __restrict__ out_ptr = nullptr;
     const float* __restrict__ out_scale = nullptr;
 
-    // Fused forward extras: bias (may be null) and amax slots (may be null).
-    float* __restrict__ amax_a = nullptr;
-    float* __restrict__ amax_b = nullptr;
-
-    // Shapes. total is only used by the elementwise quantize kernel. `int`
-    // covers every realistic LLM shape; the kernels promote to int64 for all
-    // pointer arithmetic.
+    // Shapes. `int` covers every realistic LLM shape; the kernels promote
+    // to int64 for all pointer arithmetic.
     int m, n, k;
 
-    // Physical leading dimensions (column count, i.e. row stride) of A and B.
-    // For a non-transposed operand the stride equals the contract dim; for a
-    // transposed operand it is the operand's own column count. The binding
-    // packs these so the kernel reads both buffers either naturally or
-    // transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
+    // Physical leading dimensions (column count, i.e. row stride) of A and
+    // B. For a non-transposed operand the stride equals the contract dim;
+    // for a transposed operand it is the operand's own column count. The
+    // binding packs these so the kernel reads both buffers either naturally
+    // or transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
     int a_ld, b_ld;
-
-    int total;
 };
 
 }  // namespace fp8

csrc/kernels/fp8/gemm.cuh

@@ -70,7 +70,7 @@ __device__ __forceinline__ unsigned quantize2(unsigned pair, float inv,
 }
 
 template <FP8Format Fmt>
-__global__ void fp8_quantize_kernel(FP8Params p) {
+__global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
     const float inv = 1.0f / *p.scale_a;
     const auto* x = reinterpret_cast<const __nv_bfloat16*>(p.a_ptr);
     void* x8 = p.out_ptr;
@@ -122,6 +122,51 @@ __global__ void fp8_quantize_kernel(FP8Params p) {
             atomic_max_float(amax, v);
         }
     }
+    if (p.ring_state && amax) {
+        // Delayed-scaling ring finalization as a last-block epilogue (the
+        // CUDA threadFenceReduction pattern): the fence + counter elect the
+        // final block once every block's atomic_max above is visible; warp 0
+        // folds the fresh amax into the window, reduces it and publishes the
+        // next step's scale, then re-arms the counter for the next launch.
+        // __fdiv_rn / ldexpf keep the scale bit-identical to the eager
+        // (peak / fp8_max) / 2^margin fp32 chain despite --use_fast_math.
+        __threadfence();
+        __shared__ bool ring_last;
+        if (threadIdx.x == 0)
+            ring_last = atomicAdd(reinterpret_cast<int*>(p.ring_state +
+                                                        p.ring_len + 1),
+                                  1) == gridDim.x - 1;
+        __syncthreads();
+        if (ring_last && threadIdx.x < 32) {
+            float* hist = p.ring_state;
+            const int lane = threadIdx.x;
+            float v = 0.0f;
+            if (lane < p.ring_len) v = hist[lane];
+            if (lane == p.ring_idx) {
+                v = *amax;  // the global amax is final now
+                hist[lane] = v;
+            }
+            // Windows longer than one warp (atypical) fold the tail.
+            for (int i = lane + 32; i < p.ring_len; i += 32) {
+                float h = hist[i];
+                if (i == p.ring_idx) {
+                    h = *amax;
+                    hist[i] = h;
+                }
+                v = fmaxf(v, h);
+            }
+            const float peak = warp_reduce_max(v);
+            if (lane == 0) {
+                constexpr float kFmtMax =
+                    Fmt == FP8Format::E5M2 ? 57344.0f : 448.0f;
+                p.ring_state[p.ring_len] = fmaxf(
+                    ldexpf(__fdiv_rn(peak, kFmtMax), -p.ring_margin), 1e-12f);
+                __threadfence();
+                // Re-arm the counter (0.0f bits == int32 0).
+                p.ring_state[p.ring_len + 1] = 0.0f;
+            }
+        }
+    }
 }
 
 // Swizzled address inside a flat [rows * K] staging tile: the 16-byte chunk
@@ -146,20 +191,11 @@ __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) {
            ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + (col & 15));
 }
 
-// Stage-load one GEMM operand into the canonical flat [rows * K] shared tile
-// (addressing via tile_at, so stores land in the swizzled layout). The
-// transpose is folded into the staging step via a CUTLASS-style crosswise
-// layout: RowMajor (stored [rows][contract]) copies 16-byte K-contiguous runs
-// with cp.async, while ColMajor (stored [contract][rows]) reads 16-byte runs
-// along the operand's contiguous non-contract dim and scatters them across
-// the tile's rows. Crosswise runs cannot use cp.async (the 16 destination
-// bytes land on 16 different rows), so their global loads are plain LDGs —
-// issued as one batch per row group before the first scatter so their
-// latencies overlap instead of serializing behind the shared stores.
-// RowsTile is the tile's row capacity (kBlockM / kBlockN) and kThreads the
-// CTA size; the runtime `rows` bound may be smaller (tail predication).
-// `block_row` is this block's origin in the operand's row dim.
-template <typename T8, int K, typename Layout, int RowsTile, int kThreads>
+// Stage-load a CONGRUOUS operand (stored [rows][contract], contract-
+// contiguous — the only cp.async-able shape for the canonical tile) into the
+// flat [rows * K] shared tile via tile_at's swizzle. Crosswise operands go
+// through stage_crosswise_tile + transpose_crosswise_tile instead.
+template <typename T8, int K, int RowsTile, int kThreads>
 __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,
@@ -168,101 +204,29 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
     static_assert(RowsTile * kChunks % kThreads == 0,
                   "tile chunks must divide evenly across threads");
     constexpr int kCpt = RowsTile * kChunks / kThreads;  // chunks per thread
-    if constexpr (std::is_same_v<Layout, ColMajor>) {
-        // Operand stored [contract][rows]: contiguous along the non-contract
-        // dim. Each thread scatters one 16-byte run per K/32 pass; when the
-        // tile has more 16-row groups than warps (RowsTile > kThreads/2),
-        // each thread covers several groups.
-        constexpr int kWarpsTile = kThreads / 32;
-        constexpr int kGroups = RowsTile / 16;
-        constexpr int kPasses = K / 32;
-        static_assert(kGroups % kWarpsTile == 0,
-                      "row groups must divide evenly across warps");
-        const int kl = tid & 31;  // byte column within a 32B pass
-        // r0 is always a multiple of 16 (block_row is a multiple of RowsTile
-        // and each group covers 16 rows), so every run shares the base+ld
-        // alignment: one uniform check instead of one per pass.
-        const bool run_aligned =
-            ((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
+    // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of
+    // one row (K=64: a contiguous 32B pair; K=32: a single chunk).
+    constexpr int kCpr = kChunks / kCpt;  // chunks per row slice
+    const int r = tid / kCpr;
+    const int c0 = (tid % kCpr) * kCpt * 16;
+    const int64_t row = block_row + r;
+    const bool row_ok = row < rows;
+    // k_base and every c are multiples of 16, so the per-chunk sources
+    // share the row base's alignment.
+    const auto* src = operand + row * ld + k_base;
+    const bool chunk_aligned = (reinterpret_cast<uintptr_t>(src) & 15) == 0;
 #pragma unroll
-        for (int g = 0; g < kGroups / kWarpsTile; ++g) {
-            const int rg = (tid >> 5) + g * kWarpsTile;
-            const int64_t r0 = block_row + rg * 16;
-            const bool rows_full = r0 + 15 < rows;  // pass-invariant
-            // Batch every 16B run load of this row group before the first
-            // scatter: the LDGs are independent, and the byte-granular
-            // shared stores would otherwise serialize behind each one.
-            uint4 v[kPasses];
-            bool fast[kPasses];
+    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, true);
+        } else {
+            // Tail chunk (or misaligned base): predicated scalar fill.
 #pragma unroll
-            for (int pass = 0; pass < kPasses; ++pass) {
-                const int64_t k_idx = k_base + kl + pass * 32;
-                fast[pass] = rows_full && run_aligned && k_idx < contract;
-                if (fast[pass])
-                    v[pass] =
-                        *reinterpret_cast<const uint4*>(operand + k_idx * ld + r0);
-            }
-#pragma unroll
-            for (int pass = 0; pass < kPasses; ++pass) {
-                const int col = kl + pass * 32;
-                if (fast[pass]) {
-                    const auto* bytes = reinterpret_cast<const T8*>(&v[pass]);
-                    // Scatter 16 bytes along the tile rows through tile_at's
-                    // swizzle. Rows sharing a physical chunk form groups of
-                    // (8 / kChunks) consecutive rows (see tile_at), so each
-                    // group is one tile_at address plus a K-byte row stride.
-                    constexpr int kGrp = 8 / kChunks;
-#pragma unroll
-                    for (int j = 0; j < 16 / kGrp; ++j) {
-                        T8* p = tile_at<K>(tile, rg * 16 + j * kGrp, col);
-#pragma unroll
-                        for (int i = 0; i < kGrp; ++i)
-                            p[i * K] = bytes[j * kGrp + i];
-                    }
-                } else if (k_base + col < contract) {
-                    // Row-tail or misaligned run: byte-granular gather with
-                    // per-row predication (the k column itself is in range).
-#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);
-                    }
-                } else {
-                    // Contract tail: straight zero-fill, no global traffic.
-#pragma unroll
-                    for (int i = 0; i < 16; ++i)
-                        *tile_at<K>(tile, rg * 16 + i, col) = T8(0.0f);
-                }
-            }
-        }
-    } else {
-        // Operand stored [rows][contract]: contiguous along the contract dim.
-        // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of
-        // one row (K=64: a contiguous 32B pair; K=32: a single chunk).
-        constexpr int kCpr = kChunks / kCpt;  // chunks per row slice
-        const int r = tid / kCpr;
-        const int c0 = (tid % kCpr) * kCpt * 16;
-        const int64_t row = block_row + r;
-        const bool row_ok = row < rows;
-        // k_base and every c are multiples of 16, so the per-chunk sources
-        // share the row base's alignment.
-        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, true);
-            } else {
-                // Tail chunk (or misaligned base): predicated 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);
-            }
+            for (int i = 0; i < 16; ++i)
+                dst[i] =
+                    row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f);
         }
     }
 }
@@ -275,15 +239,200 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
 // ---------------------------------------------------------------------------
 
 // Swizzled 16B-chunk address (tile_at's layout) as a raw shared-memory
-// pointer for ldmatrix. Valid for kK in {32, 64} (the swizzle itself lives
-// only in tile_at; this wrapper just converts the element address).
+// pointer for ldmatrix. Valid for kK in {32, 64, 128} (the swizzle itself
+// lives only in tile_at; this wrapper just converts the element address).
 template <typename T8, int kK>
 __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk) {
-    static_assert(kK == 32 || kK == 64,
-                  "fragment swizzle offsets assume kK in {32, 64}");
+    static_assert(kK == 32 || kK == 64 || kK == 128,
+                  "fragment swizzle offsets assume kK in {32, 64, 128}");
     return __cvta_generic_to_shared(tile_at<kK>(tile, row, chunk << 4));
 }
 
+// Crosswise operands (stored [contract][rows], rows-contiguous) cannot be
+// cp.async'd into the canonical [rows][contract] tile — a 16B global run
+// holds one contract byte for each of 16 rows. They stage K-major instead
+// (byte (p, r) at p*RowsTile + r), where the very same runs land contiguously
+// and cp.async applies unchanged; a per-tile smem->smem transpose (below)
+// then produces the canonical swizzled tile the MMA fragments read. This
+// keeps the whole global→shared path asynchronous — the synchronous
+// LDG+byte-scatter staging this replaces left the kernel long-scoreboard
+// bound (ncu: 4.6 stalled loads per issue vs 0.4 on the congruous path).
+template <typename T8, int K, int RowsTile, int kThreads>
+__device__ __forceinline__ void
+stage_crosswise_tile(T8* staging, 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 kRuns = K * RowsTile / 16;  // 16B runs per tile
+    // r0 is a multiple of 16 and p*ld keeps 16B alignment whenever ld has it,
+    // so one uniform verdict covers every run.
+    const bool run_aligned =
+        ((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
+    for (int run = tid; run < kRuns; run += kThreads) {
+        const int pl = run % K;  // local contract byte (column of the run)
+        const int rg = run / K;  // 16-row group
+        const int64_t r0 = block_row + (int64_t)rg * 16;
+        T8* dst = staging + pl * RowsTile + rg * 16;
+        if (run_aligned && r0 + 15 < rows && k_base + pl < contract)
+            astrai::cp_async_16(dst, operand + (k_base + pl) * ld + r0, true);
+        else {
+            // Row tail, contract tail or misaligned base: predicated fill.
+#pragma unroll
+            for (int i = 0; i < 16; ++i) {
+                const int64_t r = r0 + i;
+                dst[i] = r < rows && k_base + pl < contract
+                             ? operand[(k_base + pl) * ld + r]
+                             : T8(0.0f);
+            }
+        }
+    }
+}
+
+// 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. Used for crosswise operands whose global data is
+// typically L2-resident (the A side of dW): the staging detour's extra
+// shared-memory round trip costs more than the latency it hides there,
+// while crosswise B operands (DRAM-streamed weights of dX) take the
+// asynchronous stage_crosswise_tile path instead.
+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 always a multiple of 16 (block_row is a multiple of RowsTile and
+    // each group covers 16 rows), and p*ld keeps the base 16B-aligned
+    // whenever ld is, 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)]. Byte i of a uint4
+                // lives in its (i>>2)-th 32-bit register.
+                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);
+                }
+            }
+        }
+    }
+}
+
+// K-major staging -> canonical [rows][kK] swizzled tile, one chunk at a time.
+// Each chunk (indexed within a k_seg region of `quads_per_seg` quads) covers
+// 4 consecutive contract bytes x 16 rows: four LDS.128 grab the staging runs,
+// PRMT byte selects transpose them in registers, and sixteen STS.32 land the
+// row quads through tile_at's swizzle — 4x fewer store instructions than a
+// byte-granular scatter. Chunk-at-a-time lets the caller pool work across
+// operands; the region restriction lets the main loop overlap one region's
+// transpose with another region's MMAs (a whole-tile serial transpose put
+// the crosswise GEMMs at 25% tensor utilization).
+template <typename T8, int K, int RowsTile>
+__device__ __forceinline__ void
+transpose_crosswise_region(T8* tile, const T8* staging, int idx, int quad0) {
+    constexpr int kGroups = RowsTile / 16;
+    const int quad = quad0 + idx / kGroups;
+    const int rg = idx % kGroups;
+    // The four runs sit RowsTile bytes apart (one per contract byte of the
+    // quad); each run is 16 contiguous staging bytes = 16 rows.
+    const char* run0 = reinterpret_cast<const char*>(
+        staging + quad * 4 * RowsTile + rg * 16);
+    uint4 v[4];
+#pragma unroll
+    for (int s = 0; s < 4; ++s)
+        v[s] = *reinterpret_cast<const uint4*>(run0 + s * RowsTile);
+    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)]. Byte i of a uint4 lives in
+        // its (i>>2)-th 32-bit register.
+        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);
+    }
+}
+
+// Layout-aware shared-memory budget and occupancy hint. A congruous operand
+// needs its kStages rotating canonical buffers; a staged-crosswise operand
+// (crosswise B with kBStaged) needs kStages K-major staging buffers plus ONE
+// canonical buffer (rewritten every tile by the in-kernel transpose); a
+// direct-crosswise operand rotates kStages+1 canonical buffers so its load
+// can run ahead of the compute phase (see the kernel's pipelining note).
+// The 48KB static-smem watermark picks the resident-CTA hint for
+// __launch_bounds__ (sm_89: 100KB smem per SM, so two CTAs fit while each
+// stays within the static budget).
+template <typename Traits, typename LayoutA, typename LayoutB, bool StagedB>
+struct Fp8GemmSmem {
+    // Crosswise = the stage-load's view: A's tag directly, B's transposed.
+    // A-crosswise always loads direct (L2-typical activations); B-crosswise
+    // stages only when its contract dim is long enough to stream DRAM.
+    static constexpr bool kCrossA = std::is_same_v<LayoutA, ColMajor>;
+    static constexpr bool kCrossB = std::is_same_v<LayoutB, RowMajor>;
+    static constexpr bool kBStagePath = kCrossB && StagedB;
+    static constexpr bool kDirectA = kCrossA;
+    static constexpr bool kDirectB = kCrossB && !kBStagePath;
+    static constexpr int kBytes =
+        (kDirectA ? Traits::kStages + 1 : Traits::kStages) *
+            Traits::kBlockM * Traits::kK +
+        (kDirectB || kBStagePath ? Traits::kStages + 1 : Traits::kStages) *
+            Traits::kBlockN * Traits::kK;
+    static constexpr int kMinCtas = kBytes <= 48 * 1024 ? 2 : 1;
+};
+
 // LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h):
 // A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] /
 // ColMajor = [N][K]. The kernel always computes
@@ -297,25 +446,51 @@ __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk
 // (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and
 // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the
 // launcher dispatches to it there (see launch_fp8_gemm).
-template <typename Traits, bool OutFp8 = false, typename LayoutA = RowMajor,
-          typename LayoutB = RowMajor>
-__global__ void
-    __launch_bounds__((Traits::kBlockM / 64) * (Traits::kBlockN / 32) * 32, 2)
+template <typename Traits, bool OutFp8 = false, 
+        typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
+        bool kBStaged = true>
+__global__ void __launch_bounds__(Traits::kCtaThreads,
+                                  Fp8GemmSmem<Traits, LayoutA, LayoutB,
+                                              kBStaged>::kMinCtas) 
     fp8_gemm_kernel(FP8Params p) {
     using T8 = std::conditional_t<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>;
     constexpr int kBlockM = Traits::kBlockM;
     constexpr int kBlockN = Traits::kBlockN;
     constexpr int kK = Traits::kK;
     constexpr int kStages = Traits::kStages;
-    constexpr int kCtaThreads = (kBlockM / 64) * (kBlockN / 32) * 32;
+    constexpr int kCtaThreads = Traits::kCtaThreads;
+    constexpr bool kCrossA = Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kCrossA;
+    constexpr bool kCrossB = Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kCrossB;
+    constexpr bool kBStagePath =
+        Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kBStagePath;
+    constexpr bool kDirectA = Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kDirectA;
+    constexpr bool kDirectB = Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kDirectB;
     static_assert(kStages >= 1 && kStages <= 8,
-                  "FP8 GEMM stages must be in the range [1, 8]");
+                  "FP8 GEMM stages must be in [1, 8]");
     // Tiles are flat [rows * kK] with a 16B-chunk XOR swizzle (tile_at):
     // ldmatrix reads whole 16B chunks through the same mapping the staging
     // writes, and the swizzle removes the bank conflict the unswizzled
-    // 8-word row stride caused (see tile_at).
-    __shared__ __align__(16) T8 a_smem[kStages][kBlockM * kK];
-    __shared__ __align__(16) T8 b_smem[kStages][kBlockN * kK];
+    // 8-word row stride caused (see tile_at). The stages live in dynamic
+    // shared memory so deep pipelines (kStages * (kBlockM + kBlockN) * kK >
+    // 48KB static limit) opt in via cudaFuncSetAttribute in the launcher.
+    extern __shared__ __align__(16) char fp8_gemm_smem[];
+    // Per operand: congruous = kStages rotating canonical buffers; direct-
+    // crosswise = kStages+1 of them (the load for tile i+kStages targets
+    // buffer (i-1)%(kStages+1) — the one compute(i-1) finished reading at
+    // the previous barrier — so it issues right after barrier 1 and its
+    // global-load latency overlaps the MMA phase below); staged-crosswise
+    // (B) = kStages K-major staging buffers (filled by cp.async, one per
+    // tile in flight) followed by one canonical buffer the per-tile
+    // transpose rewrites.
+    constexpr int kAStageBytes = kBlockM * kK;
+    constexpr int kBStageBytes = kBlockN * kK;
+    constexpr int kARing = kDirectA ? kStages + 1 : kStages;  // A canonic ring
+    constexpr int kBRing = kDirectB ? kStages + 1 : kStages;  // B canonic ring
+    constexpr int kStB = kStages;  // B staging ring size (see above)
+    T8* const a_base = reinterpret_cast<T8*>(fp8_gemm_smem);
+    T8* const b_base =
+        reinterpret_cast<T8*>(fp8_gemm_smem + kARing * kAStageBytes);
+    T8* const b_canon = b_base + kStB * kBStageBytes;  // staged B only
 
     const auto* a = reinterpret_cast<const T8*>(p.a_ptr);
     const auto* b = reinterpret_cast<const T8*>(p.b_ptr);
@@ -334,12 +509,14 @@ __global__ void
     // before advancing along N. All CTAs of one group share the same B column
     // stripe, so B tiles stay hot in L2 across the wave (the default
     // N-fastest order makes each wave touch every B tile instead).
-    // Measured win for the A-crosswise layouts (10-21% at K>=2048) and loss
-    // for A-congruous (-17..20%, A's cp.async stream prefers the N-fastest
-    // order) — so the branch follows LayoutA.
+    // kGroupRaster is a template knob (the launcher defaults it to the
+    // measured best per layout: grouped for A-crosswise (dW) and for the
+    // congruous NT forward — whose big B operand gains the most from the
+    // shared stripe — plain for dX's crosswise-B layouts, where it measured
+    // neutral).
     constexpr int kGroupM = 8;
     int block_m, block_n;
-    if constexpr (std::is_same_v<LayoutA, ColMajor>) {
+    if constexpr (kGroupRaster) {
         const int blocks_m = gridDim.y;
         const int bid = blockIdx.y * gridDim.x + blockIdx.x;
         const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM;
@@ -364,18 +541,63 @@ __global__ void
     const float sb = *p.scale_b;
     float acc[4][4][4] = {};  // [nt][mt][acc]
 
-    // Both operands are staged into the canonical [M][kK] / [N][kK] shared
-    // tiles regardless of their global layout (see load_operand_tile), so the
-    // MMA fragment reads below stay unchanged across the four layout
-    // combinations. A's tag already names the operand view ([M][K] =
-    // [rows][contract]); B's tag is relative to the canonical [K][N], so the
-    // stage-load sees its transpose (transpose_layout_t, see common.h).
-    auto load_tile = [&](int stage, int64_t k_base) {
-        load_operand_tile<T8, kK, LayoutA, kBlockM, kCtaThreads>(
-            a_smem[stage], a, m, k, a_ld, tid, k_base, (int64_t)block_m * kBlockM);
-        load_operand_tile<T8, kK, transpose_layout_t<LayoutB>, kBlockN,
-                          kCtaThreads>(b_smem[stage], b, n, k, b_ld, tid, k_base,
-                                       (int64_t)block_n * kBlockN);
+    // Both operands end up in the canonical [M][kK] / [N][kK] shared tiles
+    // the MMA fragments read, regardless of their global layout. A's tag
+    // already names the operand view ([M][K] = [rows][contract]); B's tag is
+    // relative to the canonical [K][N], so the stage-load sees its transpose
+    // (transpose_layout_t, see common.h). Congruous operands cp.async
+    // straight into their rotating canonical buffers; crosswise operands
+    // cp.async into K-major staging (zero transformation) and get a per-tile
+    // smem->smem transpose below.
+    // Asynchronous loads for tile `tile`: congruous operands cp.async into
+    // their canonical rings, a staged B cp.asyncs into its K-major staging
+    // ring. Called after the post-compute barrier, alongside the commit.
+    auto load_async = [&](int64_t tile) {
+        const int64_t k_base = tile * kK;
+        if constexpr (!kDirectA)
+            load_operand_tile<T8, kK, kBlockM, kCtaThreads>(
+                a_base + (tile % kARing) * kAStageBytes, a, m, k, a_ld, tid,
+                k_base, (int64_t)block_m * kBlockM);
+        if constexpr (kBStagePath)
+            stage_crosswise_tile<T8, kK, kBlockN, kCtaThreads>(
+                b_base + (tile % kStB) * kBStageBytes, b, n, k, b_ld, tid,
+                k_base, (int64_t)block_n * kBlockN);
+        if constexpr (!kDirectB && !kBStagePath)
+            load_operand_tile<T8, kK, kBlockN, kCtaThreads>(
+                b_base + (tile % kBRing) * kBStageBytes, b, n, k, b_ld, tid,
+                k_base, (int64_t)block_n * kBlockN);
+    };
+    // Synchronous direct-crosswise loads for tile `tile` into the operand's
+    // (kStages+1)-deep canonical ring. In the steady state this runs right
+    // after barrier 1, so the LDG latency and the PRMT transpose overlap the
+    // MMA phase of the current tile instead of stalling the inter-barrier
+    // window (which dominated the dX/dW stall profile: barrier 3.7-4.1 +
+    // long-scoreboard 1.6-1.8 stalls per issue on the production shapes).
+    // Ring safety: the write targets buffer (i+kStages)%(kStages+1) =
+    // (i-1)%(kStages+1), which compute(i-1) finished reading before the
+    // previous barrier and compute(i+kStages) does not touch until several
+    // barriers later.
+    auto load_direct = [&](int64_t tile) {
+        const int64_t k_base = tile * kK;
+        if constexpr (kDirectA)
+            load_crosswise_direct<T8, kK, kBlockM, kCtaThreads>(
+                a_base + (tile % kARing) * kAStageBytes, a, m, k, a_ld, tid,
+                k_base, (int64_t)block_m * kBlockM);
+        if constexpr (kDirectB)
+            load_crosswise_direct<T8, kK, kBlockN, kCtaThreads>(
+                b_base + (tile % kBRing) * kBStageBytes, b, n, k, b_ld, tid,
+                k_base, (int64_t)block_n * kBlockN);
+    };
+    // smem->smem transpose of one k_seg region (kSegQuads contract quads) of
+    // this tile's staged-crosswise B into its single canonical buffer.
+    auto transpose_tile = [&](int tile, int seg) {
+        if constexpr (!kBStagePath) return;
+        constexpr int kSegQuads = kK / 4 / (kK / kMmaK);  // quads per k_seg
+        constexpr int kBRegion = kSegQuads * (kBlockN / 16);
+        const T8* b_stg = b_base + (tile % kStB) * kBStageBytes;
+        for (int idx = tid; idx < kBRegion; idx += kCtaThreads)
+            transpose_crosswise_region<T8, kK, kBlockN>(b_canon, b_stg, idx,
+                                                        seg * kSegQuads);
     };
 
     const int64_t tile_count = (k + kK - 1) / kK;
@@ -397,16 +619,18 @@ __global__ void
 
     // Prime the pipeline. Each committed group occupies one circular shared
     // memory stage; the loop also handles K dimensions smaller than kStages.
+    // Direct loads run synchronously here (back to back with their commit);
+    // the steady state below overlaps them with the compute phase.
 #pragma unroll
     for (int stage = 0; stage < kStages; ++stage) {
         if (stage < tile_count) {
-            load_tile(stage, static_cast<int64_t>(stage) * kK);
+            load_async(stage);
+            load_direct(stage);
             astrai::cp_async_commit_group();
         }
     }
 
     for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) {
-        const int stage = static_cast<int>(tile_index % kStages);
         const int64_t remaining = tile_count - tile_index - 1;
 
         // Keep up to kStages - 1 younger groups in flight while making the
@@ -418,48 +642,83 @@ __global__ void
         // 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 (tile_index + kStages < tile_count)
+            load_direct(tile_index + kStages);
+
+        // Staged-crosswise B: produce the canonical tile one k_seg region at
+        // a time so each region's transpose overlaps the previous region's
+        // MMA sequence (the transposes are pure shared-memory traffic — B's
+        // global path stayed fully asynchronous above).
+        constexpr int kSegs = kK / kMmaK;
+        if constexpr (kBStagePath) {
+            transpose_tile(tile_index, 0);
+            // Barrier 2: region 0 visible to every thread before its
+            // fragment loads. (Compiled out for congruous/direct layouts.)
+            __syncthreads();
+        }
+
+        const T8* a_tile = a_base + (size_t)(tile_index % kARing) * kAStageBytes;
+        const T8* b_tile = kBStagePath
+                               ? b_canon
+                               : b_base + (size_t)(tile_index % kBRing) * kBStageBytes;
+
         // 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg —
         // 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in
         // the 128x64-tile version (the kernel was LSU-issue-bound there).
-        constexpr int kSegs = kK / kMmaK;
-        // B fragments double-buffered across k_segs: the next k_seg's B load
-        // is issued before the current k_seg's MMA sequence, so its LDS
-        // latency hides behind the A pipeline + tensor-pipe work (same trick
-        // as the A mt+1 prefetch below; costs kSegs x 8 registers).
+        // B fragments double-buffer across k_segs while B is congruous (no
+        // region writes in flight); a crosswise B reloads per k_seg after
+        // the region's transpose became visible.
         unsigned b_frag[2][4][2];
+        if constexpr (!kBStagePath) {
 #pragma unroll
-        for (int nt = 0; nt < 4; ++nt) {
-            const int row = b_row0 + nt * 8 + r7;
-            astrai::ldmatrix_x2_lane(b_frag[0][nt],
-                                     frag_addr<T8, kK>(b_smem[stage], row, rh8));
+            for (int nt = 0; nt < 4; ++nt) {
+                const int row = b_row0 + nt * 8 + r7;
+                astrai::ldmatrix_x2_lane(b_frag[0][nt],
+                                         frag_addr<T8, kK>(b_tile, row, rh8));
+            }
         }
 #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) {
+            if constexpr (kBStagePath) {
+#pragma unroll
+                for (int nt = 0; nt < 4; ++nt) {
+                    const int row = b_row0 + nt * 8 + r7;
+                    astrai::ldmatrix_x2_lane(
+                        b_frag[bcur][nt],
+                        frag_addr<T8, kK>(b_tile, row, k_seg * 2 + rh8));
+                }
+            } else if (k_seg + 1 < kSegs) {
 #pragma unroll
                 for (int nt = 0; nt < 4; ++nt) {
                     const int row = b_row0 + nt * 8 + r7;
                     astrai::ldmatrix_x2_lane(
                         b_frag[bnext][nt],
-                        frag_addr<T8, kK>(b_smem[stage], row,
-                                          (k_seg + 1) * 2 + rh8));
+                        frag_addr<T8, kK>(b_tile, row, (k_seg + 1) * 2 + rh8));
                 }
             }
+            // Region k_seg+1's transpose overlaps this region's MMA work
+            // (disjoint canonical regions, no race).
+            if constexpr (kBStagePath) {
+                if (k_seg + 1 < kSegs)
+                    transpose_tile(tile_index, 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 fixed
             // latency hides behind tensor-pipe work (cuts the `wait` stall,
             // ~2.3 cycles/issue before this). Costs 4 extra registers.
             unsigned a_frag[5][4];
             astrai::ldmatrix_x4_lane(
-                a_frag[0], frag_addr<T8, kK>(a_smem[stage], a_row0 + rh8 * 8 + r7,
+                a_frag[0], frag_addr<T8, kK>(a_tile, a_row0 + rh8 * 8 + r7,
                                              k_seg * 2 + rh16));
 #pragma unroll
             for (int mt = 0; mt < 4; ++mt) {
                 if (mt < 3)
                     astrai::ldmatrix_x4_lane(
                         a_frag[mt + 1],
-                        frag_addr<T8, kK>(a_smem[stage],
+                        frag_addr<T8, kK>(a_tile,
                                           a_row0 + (mt + 1) * 16 + rh8 * 8 + r7,
                                           k_seg * 2 + rh16));
 #pragma unroll
@@ -467,12 +726,18 @@ __global__ void
                     astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt], b_frag[bcur][nt],
                                          acc[nt][mt]);
             }
+            // Barrier 3: region k_seg+1's transposes complete and become
+            // visible before the next k_seg reads them.
+            if constexpr (kBStagePath) {
+                if (k_seg + 1 < kSegs) __syncthreads();
+            }
         }
-        // Barrier 2: every thread finished reading this stage's tiles before
-        // the prefetch for the (i+kStages)-th tile overwrites them.
+        // Barrier 4: every thread finished reading this stage's tiles before
+        // the prefetch for the (i+kStages)-th tile overwrites them (and the
+        // next iteration's transposes rewrite the canonical buffer).
         __syncthreads();
         if (tile_index + kStages < tile_count) {
-            load_tile(stage, (tile_index + kStages) * kK);
+            load_async(tile_index + kStages);
             astrai::cp_async_commit_group();
         }
     }
@@ -539,7 +804,7 @@ __global__ void
 // ---------------------------------------------------------------------------
 
 template <FP8Format Fmt>
-void launch_fp8_quantize(const FP8Params& p, cudaStream_t stream) {
+void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
     constexpr int kThreads = 256;
     // One block per 256 vectors (8 elements each); at least one block so the
     // scalar tail of a tiny / misaligned tensor is still covered.
@@ -548,26 +813,74 @@ void launch_fp8_quantize(const FP8Params& p, cudaStream_t stream) {
     fp8_quantize_kernel<Fmt><<<blocks, kThreads, 0, stream>>>(p);
 }
 
+// Launch one kernel instantiation with its shared-memory budget: stages live
+// in dynamic smem, so budgets beyond the 48KB static limit opt in once per
+// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
+// memory"). Templated on the kernel *value* (auto NTTP) so every
+// instantiation owns its own armed flag — same-signature kernels must not
+// share it (the attribute is per-function).
+template <auto Kernel, typename... Args>
+void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
+                      cudaStream_t stream, Args... args) {
+    if (smem_bytes > 48 * 1024) {
+        static bool armed = false;  // per instantiation
+        if (!armed) {
+            cudaFuncSetAttribute(Kernel, 
+                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                 smem_bytes);
+            armed = true;
+        }
+    }
+    Kernel<<<grid, block, smem_bytes, stream>>>(args...);
+}
+
 // Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles).
-// kK selects the K tile (32 or 64; 64 halves the __syncthreads count per K
-// and doubles the MMA work per stage, at 2x the smem per stage — measured
-// 10-35% across shapes, so 64 is the default). Stages=2 with kK=64 keeps the
-// pipeline at 32KB smem; deeper pipelines only win on K >= 4096 squares and
-// lose elsewhere. LayoutA/LayoutB mirror the kernel template (defaults keep
-// the NN layout: out = a @ b). m <= 64 dispatches to the 64x128 CTA — a
-// 128-row CTA would waste half its MMA work on predicated-off rows.
+// kK selects the K tile (32 / 64 / 128; larger kK halves the __syncthreads
+// count per K and doubles the MMA work per stage at more smem per stage).
+// Stages is the cp.async pipeline depth (smem = Stages * (BM + BN) * kK
+// bytes for congruous layouts; deep pipelines are dynamic-smem backed, 1
+// CTA/SM past 48KB). GroupRaster defaults to the historically-measured best
+// per LayoutA (grouped for A-crosswise, plain for A-congruous). m <= 64
+// dispatches to the 64x128 CTA — a 128-row CTA would waste half its MMA work
+// on predicated-off rows.
+// Crosswise B takes the asynchronous staging+transpose pipeline only when
+// the contract dim is long enough that B streams from DRAM (dX-class GEMMs,
+// k = N_ffn); short-K crosswise GEMMs (dW: k = M tokens) read L2-resident
+// operands, where the staging round trip costs more shared-memory traffic
+// than the latency it hides (measured: dW ~37 TF direct vs ~29 TF staged,
+// dX ~39 TF staged vs ~38 direct).
+constexpr int64_t kCrossStageMinK = 8192;
+
 template <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor,
-          typename LayoutB = RowMajor, int kK = 64, int Stages = 2>
+          typename LayoutB = RowMajor, int kK = 64, int Stages = 2,
+          bool GroupRaster = std::is_same_v<LayoutA, ColMajor> || std::is_same_v<LayoutB, ColMajor>>
 void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
     dim3 grid((p.n + 127) / 128, (p.m + 127) / 128);
+    const bool b_staged = p.k >= kCrossStageMinK;
     if (p.m <= 64) {
         using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
-        fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB>
-            <<<grid, (64 / 64) * (128 / 32) * 32, 0, stream>>>(p);
+        if (b_staged)
+            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+                                             GroupRaster, true>>(
+                Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
+                dim3(Traits::kCtaThreads), stream, p);
+        else
+            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+                                             GroupRaster, false>>(
+                Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
+                dim3(Traits::kCtaThreads), stream, p);
     } else {
         using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
-        fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB>
-            <<<grid, (128 / 64) * (128 / 32) * 32, 0, stream>>>(p);
+        if (b_staged)
+            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+                                             GroupRaster, true>>(
+                Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
+                dim3(Traits::kCtaThreads), stream, p);
+        else
+            launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
+                                             GroupRaster, false>>(
+                Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
+                dim3(Traits::kCtaThreads), stream, p);
     }
 }
 

csrc/kernels/fp8/ops.cu

@@ -71,30 +71,34 @@ void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
     p.out_scale = out_scale ? out_scale->data_ptr<float>() : nullptr;
     p.bias = bias;
     p.bias_scale = bias_scale ? bias_scale->data_ptr<float>() : nullptr;
-    p.amax_a = nullptr;
-    p.amax_b = nullptr;
     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);
-    p.total = 0;
 }
 
-void pack_quantize_params(FP8Params& p, const void* x, void* x8,
+// Pack the quantize params, optionally wiring the delayed-scaling ring.
+// ring (may be null) packs [hist[len] | scale | counter]; len/margin come
+// from the active recipe and idx is the caller's slot for this step.
+void pack_quantize_params(FP8QuantizeParams& p, const void* x, void* x8,
                           const torch::Tensor& scale, torch::Tensor* amax,
-                          int64_t total) {
+                          const torch::Tensor* ring, int64_t ring_idx,
+                          int64_t ring_margin, int64_t total) {
     p.a_ptr = x;
-    p.b_ptr = nullptr;
     p.out_ptr = x8;
     p.scale_a = scale.data_ptr<float>();
-    p.scale_b = nullptr;
-    p.out_scale = nullptr;
-    p.bias = nullptr;
     p.amax_a = amax ? amax->data_ptr<float>() : nullptr;
-    p.amax_b = nullptr;
-    p.m = p.n = p.k = 0;
-    p.a_ld = p.b_ld = 0;
+    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);
 }
 
@@ -152,11 +156,14 @@ std::tuple<torch::Tensor, torch::Tensor> quantize_bf16(torch::Tensor x,
     auto x_c = x.contiguous();
     auto x8 = torch::empty_like(
         x_c, x_c.options().dtype(fmt ? torch::kFloat8_e5m2
-                                     : torch::kFloat8_e4m3fn));
-    auto amax = torch::zeros({1}, x_c.options().dtype(torch::kFloat32));
-    FP8Params p;
+                                      : 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,
-                         x_c.numel());
+                         nullptr, 0, 0, x_c.numel());
     if (fmt) {
         launch_fp8_quantize<FP8Format::E5M2>(p, stream.stream());
     } else {
@@ -226,16 +233,28 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
     return out;
 }
 
-std::tuple<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) {
+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.
-    // 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.
+    // 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;
@@ -267,14 +286,25 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
         if (b_prequant) check_scale(*bias_scale, x, "bias_scale");
     }
     auto x8 = torch::empty({m, k}, x_c.options().dtype(f8opt));
-    auto amax_x = torch::zeros({1}, x.options().dtype(torch::kFloat32));
-    auto amax_w = torch::zeros({1}, x.options().dtype(torch::kFloat32));
+    // 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) {
-        FP8Params qp;
+                        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());
@@ -282,13 +312,14 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
             launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
         }
     };
-    quantize(x_c, x8, sx, &amax_x);
+    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);
+    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]
@@ -309,16 +340,22 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
 
     std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
     shape.push_back(n);
-    return {out.reshape(shape), amax_x, amax_w};
+    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) {
+                    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).
+    // 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");
     TORCH_CHECK(g.scalar_type() == torch::kBFloat16 &&
                     x.scalar_type() == torch::kBFloat16 &&
@@ -341,14 +378,23 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     auto grad_input = torch::empty_like(x);
     auto grad_weight = torch::empty_like(w);
     auto grad_bias = torch::empty({0}, g.options());
-    auto amax_g = torch::zeros({1}, g.options().dtype(torch::kFloat32));
+    // 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 f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2)
                      : g.options().dtype(torch::kFloat8_e4m3fn);
 
     auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst,
-                        const torch::Tensor& scale, torch::Tensor* amax) {
-        FP8Params qp;
+                        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());
@@ -375,13 +421,13 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     torch::Tensor g8;
     if (masks[0] || masks[1]) {
         g8 = torch::empty({m, n}, f8opt);
-        quantize(g_c, g8, sg, &amax_g);
+        quantize(g_c, g8, sg, &amax_g, g_ring, g_ring_idx, g_ring_margin);
     }
     // dX = g @ w: A = g8 [M,N] (contract over N), B = w8 [N,K] read transposed
     // (b[p*b_ld + n] = w[p,n]); out = [M,K], a_ld = N, b_ld = K, contract = N.
     if (masks[0]) {
         auto w8 = torch::empty({n, k}, f8opt);
-        quantize(w_c, w8, sw, nullptr);
+        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(),
@@ -394,7 +440,7 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
     // b_ld = K, contract = M.
     if (masks[1]) {
         auto x8 = torch::empty({m, k}, f8opt);
-        quantize(x_c, x8, sx, nullptr);
+        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,
@@ -423,12 +469,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("linear_forward_fp8", &linear_forward_fp8, py::arg("x"),
           py::arg("w"), py::arg("bias"), py::arg("sx"), py::arg("sw"),
           py::arg("fmt") = 0, py::arg("bias_scale") = py::none(),
+          py::arg("x_ring") = py::none(), py::arg("x_ring_idx") = 0,
+          py::arg("x_ring_margin") = 0, py::arg("w_ring") = py::none(),
+          py::arg("w_ring_idx") = 0, py::arg("w_ring_margin") = 0,
           "Pure FP8 linear forward: quantize x/w, pre-quantized GEMM with the "
           "bias fused into the epilogue; w and bias may be pre-quantized fp8 "
           "matching fmt (static inference path; fp8 bias requires bias_scale);"
-          " returns (out, amax_x, amax_w)");
+          " 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"),
-          "FP8 linear backward; returns (grad_input, grad_weight, grad_bias, amax_g)");
+          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)");
 }

tests/extension/test_fp8_mma.py

@@ -6,15 +6,22 @@ policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks
 of the custom ops) run without a GPU.
 """
 
+import threading
+
 import pytest
 import torch
+import torch.nn.functional as F
 
+import astrai.extension.fp8 as f8mod
 from astrai.extension.fp8 import (
     DelayedScaling,
     DynamicScaling,
     FP8Format,
     FP8TensorMeta,
+    _ScaleRing,
     fp8_autocast,
+    fp8_linear_enable,
+    fp8_linear_enabled,
     fp8_state,
 )
 from astrai.extension.ops.fp8 import (
@@ -86,6 +93,49 @@ def test_quantize_bf16_e5m2_format():
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
 
 
+@skip_no_fp8
+def test_quantize_ring_in_kernel_finalize():
+    """The quantize kernel finalizes the delayed-scaling ring in-kernel: the
+    measured amax lands in hist[idx], the window reduces to the next step's
+    scale on device, and the counter re-arms for the next launch."""
+    torch.manual_seed(21)
+    dev = torch.device("cuda")
+    ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0))
+    x0 = torch.randn(256, 256, device=dev, dtype=torch.bfloat16)
+    w = torch.randn(256, 256, device=dev, dtype=torch.bfloat16)
+    sw = torch.tensor([1.0], device=dev)
+    ring.seed(x0, "e4m3")
+    hist0 = ring.hist.clone()
+
+    # Step over three fresh tensors: each launch folds its amax into
+    # hist[idx] and publishes max(hist)/448 as the next scale.
+    idx = 0
+    for _ in range(3):
+        x = torch.randn(256, 256, device=dev, dtype=torch.bfloat16) * (2.0 + 4.0 * _)
+        _ = linear_forward_fp8(
+            x,
+            w,
+            None,
+            ring.scale,
+            sw,
+            "e4m3",
+            None,
+            ring.state,
+            idx,
+            0,
+        )
+        torch.cuda.synchronize()
+        expected_hist = hist0.clone()
+        expected_hist[idx] = x.abs().amax().float()
+        torch.testing.assert_close(ring.hist, expected_hist)
+        expected_scale = (expected_hist.max() / 448.0).reshape(1)
+        torch.testing.assert_close(ring.scale, expected_scale, rtol=1e-6, atol=1e-12)
+        # counter re-armed to int32 zero
+        assert ring.state[-1].view(torch.int32).item() == 0
+        hist0 = expected_hist.clone()
+        idx = (idx + 1) % 4
+
+
 @skip_no_fp8
 def test_fp8_linear_forward_and_backward():
     torch.manual_seed(7)
@@ -96,7 +146,7 @@ def test_fp8_linear_forward_and_backward():
     bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
     scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad)
 
-    out, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w)
+    out, x8, w8, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w)
     grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
         grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3"
     )
@@ -160,7 +210,7 @@ def test_fp8_linear_static_fp8_weight_and_bias():
 
     w8, _ = quantize_bf16(weight, sw, "e4m3")
     b8, _ = quantize_bf16(bias, sb, "e4m3")
-    out, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb)
+    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)
@@ -169,9 +219,10 @@ def test_fp8_linear_static_fp8_weight_and_bias():
     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")
+    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)
 
@@ -181,10 +232,6 @@ def test_fp8_linear_backward_outside_autocast():
     """aten::linear records an fp8 autograd node inside fp8_autocast; the
     backward runs fp8 kernels even after the context exits (loss.backward()
     placement is free), instead of falling back to bf16 mm."""
-    import torch.nn.functional as F
-
-    import astrai.extension.fp8 as f8mod
-
     torch.manual_seed(5)
     x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
     weight = torch.randn(
@@ -196,9 +243,9 @@ def test_fp8_linear_backward_outside_autocast():
     calls = {"bwd": 0}
     orig = f8mod.linear_backward_fp8
 
-    def spy(g, xx, ww, masks, sg, sw, sx, fmt="e5m2"):
+    def spy(*args, **kwargs):
         calls["bwd"] += 1
-        return orig(g, xx, ww, masks, sg, sw, sx, fmt)
+        return orig(*args, **kwargs)
 
     f8mod.linear_backward_fp8 = spy
     try:
@@ -331,14 +378,20 @@ def test_fp8_autocast_context():
 
 
 def test_fp8_tensor_meta_delayed_update():
-    """Meta seeds from data and refreshes the scale from the amax ring."""
+    """Meta seeds from data; hist/scale are packed views of one state buffer."""
     meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0))
     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))
-    meta.w.update(torch.tensor([4.0]), "e4m3")
-    torch.testing.assert_close(meta.w.scale, torch.tensor(4.0 / 448.0).reshape(1))
+    # [hist | scale | counter] packing: views alias the single state buffer.
+    assert meta.w.state.numel() == 4 + 2
+    assert meta.w.hist.data_ptr() == meta.w.state.data_ptr()
+    assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr()
+    # counter slot stays int32-zero (float bits) between launches
+    assert meta.w.state[-1].view(torch.int32).item() == 0
+    meta.w.advance()
+    assert meta.w.idx == 1
 
 
 def test_quantize_bf16_cpu_fallback():
@@ -373,3 +426,63 @@ def test_mm_fp8_fp8_output_cpu():
     assert out8.dtype == torch.float8_e4m3fn
     ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn)
     assert torch.equal(out8, ref)
+
+
+# --------------------------------------------------------------------------
+# torch-autocast parity: context semantics (nesting, thread locality, switch)
+# --------------------------------------------------------------------------
+
+
+def _linear():
+    """Shared helper: a small bf16 linear operand set on CUDA (grad-tracking
+    so aten::linear records an autograd node)."""
+    torch.manual_seed(31)
+    x = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    w = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    return x, w
+
+
+@skip_no_fp8
+def test_nested_disabled_region_redispatches_bf16():
+    """A nested fp8_autocast(enabled=False) region temporarily restores the
+    bf16 aten::linear path (torch's nested-disable semantics), and fp8
+    resumes when it exits."""
+    x, w = _linear()
+    with fp8_autocast(enabled=True):
+        out_fp8 = F.linear(x, w)
+        with fp8_autocast(enabled=False):
+            out_bf16 = F.linear(x, w)
+            assert type(out_bf16.grad_fn).__name__ != "_LinearFp8Backward"
+            assert out_bf16.dtype == torch.bfloat16
+        out_again = F.linear(x, w)
+        assert type(out_again.grad_fn).__name__ == "_LinearFp8Backward"
+
+
+@skip_no_fp8
+def test_global_switch_routes_without_region():
+    """fp8_linear_enable(True) routes aten::linear to fp8 outside any region
+    (the persistent default); disabling restores bf16."""
+    x, w = _linear()
+    state = fp8_state()
+    try:
+        fp8_linear_enable(True)
+        out = F.linear(x, w)
+        assert type(out.grad_fn).__name__ == "_LinearFp8Backward"
+        fp8_linear_enable(False)
+        out = F.linear(x, w)
+        assert type(out.grad_fn).__name__ != "_LinearFp8Backward"
+    finally:
+        state.reset()
+
+
+def test_autocast_state_is_thread_local():
+    """torch parity: the active config is thread-local — another thread does
+    not see an open region (CPU-only check of the flag, no kernels)."""
+    seen = {}
+    with fp8_autocast(enabled=True):
+        assert fp8_linear_enabled()
+        t = threading.Thread(target=lambda: seen.update(enabled=fp8_linear_enabled()))
+        t.start()
+        t.join()
+    assert seen["enabled"] is False
+    assert not fp8_linear_enabled()