refactor: unify operator selection behind generic dispatch

- add astrai/extension/dispatch.py: per-family decision tables over composable Specs with explicit-strict / implicit-loose resolution, ASTR_OPS env overrides, profile presets, and explain traces
- make the axis schema family-owned: register_family takes an axes extractor that snapshots whatever decision axes that family needs from the call, and the core only supplies the axis() predicate vocabulary plus a tensor_axes helper
- drop the central CallContext dataclass; resolve and explain take the raw call arguments, so unregistered handles are probed through supports_call on the same args
- migrate attention and rotary onto family-owned axes with behavior-preserving specs and spec-vs-supports_call mirror tests
- replace the non-ASCII member-of glyph in spec descriptions with plain ASCII " in "

astrai/extension/__init__.py

@@ -27,6 +27,22 @@ from astrai.extension.backend import (
     attn_backend,
     get_backend,
 )
+from astrai.extension.dispatch import (
+    Axes,
+    ExplicitSelectionError,
+    ImplRecord,
+    Resolution,
+    Spec,
+    axis,
+    explain,
+    explain_plan,
+    op_backend,
+    register_env_alias,
+    register_family,
+    resolve,
+    resolve_plan,
+    tensor_axes,
+)
 from astrai.extension.loader import KERNEL_NAMES, is_available
 from astrai.extension.ops import (
     TensorLayout,
@@ -52,4 +68,18 @@ __all__ = [
     "is_available",
     "KERNEL_NAMES",
     "apply_rotary_emb",
+    "Axes",
+    "ExplicitSelectionError",
+    "ImplRecord",
+    "Resolution",
+    "Spec",
+    "axis",
+    "explain",
+    "explain_plan",
+    "op_backend",
+    "register_env_alias",
+    "register_family",
+    "resolve",
+    "resolve_plan",
+    "tensor_axes",
 ]

astrai/extension/backend/attention.py

@@ -21,10 +21,15 @@ Usage — mirroring ``torch.nn.attention.sdpa_kernel``:
         ...
 
 Thread-safe via ``contextvars`` — each scheduler thread gets its own
-active backend. Backend resolution follows a strict precedence:
+active backend. Backend resolution is a thin facade over the generic
+operator dispatcher (``astrai.extension.dispatch``): the three backends
+are registered as the "attention" family and the decision table lives in
+``_attention_records``.  Resolution follows a strict precedence:
 
 1. explicit ``attn_backend(...)`` context (wins over everything),
-2. the process-wide ``ASTR_BACKEND`` environment override,
+2. the process-wide ``ASTR_BACKEND`` environment override
+   (or an ``ASTR_OPS`` ``attention=`` entry, which wins over the legacy
+   variable),
 3. an implicit default picked from the available backends
    (cuda > flash > torch).
 
@@ -40,12 +45,9 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 (blhd). The backend returns ``[batch, seq_len, n_heads * head_dim]``.
 """
 
-import contextvars
 import enum
 import functools
 import logging
-import os
-import threading
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
@@ -54,6 +56,22 @@ import torch
 import torch.nn.functional as F
 from torch import Tensor
 
+from astrai.extension.dispatch import (
+    Axes,
+    ImplRecord,
+    Spec,
+    axis,
+    env_selection,
+    get_override,
+    register_env_alias,
+    register_family,
+    reset_override,
+    set_override,
+    tensor_axes,
+)
+from astrai.extension.dispatch import (
+    resolve as _dispatch_resolve,
+)
 from astrai.extension.loader import is_available
 from astrai.extension.ops.attention import (
     attn_paged_decode,
@@ -72,15 +90,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-_default_backend_lock = threading.Lock()
-_env_backend_name: Optional[str] = None
-_env_backend: Optional["AttentionBackend"] = None
-_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
-    contextvars.ContextVar("attn_backend", default=None)
-)
-
-# Backends are stateless — one canonical instance per class, created lazily
-# and reused everywhere (resolution, fallback, context managers).
 _singletons: Dict[type, "AttentionBackend"] = {}
 
 
@@ -157,28 +166,6 @@ def _resolve_default_backend() -> "AttentionBackend":
     return _priority_backends()[0]
 
 
-def _environment_backend() -> Optional["AttentionBackend"]:
-    """Resolve the process-wide ``ASTR_BACKEND`` override, if configured."""
-    global _env_backend, _env_backend_name
-    name = os.environ.get("ASTR_BACKEND", "").strip().lower()
-    if not name:
-        return None
-    if name != _env_backend_name:
-        with _default_backend_lock:
-            if name != _env_backend_name:
-                try:
-                    _env_backend = _resolve_backend(name)
-                except (ValueError, RuntimeError):
-                    _env_backend = None
-                    logger.warning(
-                        "ASTR_BACKEND=%r is not a registered attention backend; "
-                        "falling back to default resolution",
-                        name,
-                    )
-                _env_backend_name = name
-    return _env_backend
-
-
 def _resolve_backend(
     backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
 ) -> "AttentionBackend":
@@ -204,18 +191,44 @@ def _resolve_backend(
     return _resolve_default_backend()
 
 
+_ENV_WARNED: set = set()
+
+
+def _environment_backend() -> Optional["AttentionBackend"]:
+    """Resolve the process-wide env override (``ASTR_OPS`` or the legacy
+    ``ASTR_BACKEND``) to a backend instance, if it names a registered one.
+
+    Invalid names warn once and are ignored, falling back to default
+    resolution — the override is soft, never fatal.
+    """
+    name = env_selection("attention")
+    if name is None:
+        return None
+    try:
+        return _resolve_backend(name)
+    except (ValueError, RuntimeError):
+        message = (
+            f"ASTR_BACKEND/ASTR_OPS value {name!r} is not a registered "
+            f"attention backend; falling back to default resolution"
+        )
+        if message not in _ENV_WARNED:
+            _ENV_WARNED.add(message)
+            logger.warning(message)
+        return None
+
+
 def get_backend(
     use_default: bool = True,
 ) -> Optional["AttentionBackend"]:
     """Resolve the active backend: explicit context > env > default.
 
     An ``attn_backend(...)`` context is the caller's explicit choice and
-    always wins.  ``ASTR_BACKEND`` is a process-wide override consulted
-    only when no context is set.  Pass ``use_default=False`` at request
-    submission to retain only an environment override or the caller's
-    :func:`attn_backend` value.
+    always wins.  ``ASTR_BACKEND`` (or ``ASTR_OPS``) is a process-wide
+    override consulted only when no context is set.  Pass
+    ``use_default=False`` at request submission to retain only an
+    environment override or the caller's :func:`attn_backend` value.
     """
-    context_backend = _current_backend.get()
+    context_backend = get_override("attention")
     if context_backend is not None:
         return context_backend
     env_backend = _environment_backend()
@@ -241,11 +254,11 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
             ...
     """
     instance = _resolve_backend(backend)
-    token = _current_backend.set(instance)
+    token = set_override("attention", instance)
     try:
         yield instance
     finally:
-        _current_backend.reset(token)
+        reset_override(token)
 
 
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
@@ -260,6 +273,24 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
     )
 
 
+def _axes(
+    q: Tensor,
+    kv_cache: Optional["KVCache"],
+    attn_mask: Optional[Tensor],
+    is_causal: bool,
+    fwd: Optional[str],
+) -> Axes:
+    """Snapshot the axes the attention decision table depends on."""
+    return tensor_axes(
+        q,
+        fwd=fwd,
+        ndim=q.dim(),
+        head_dim=q.size(-1) if q.dim() >= 1 else None,
+        has_cache=kv_cache is not None,
+        has_mask=attn_mask is not None,
+    )
+
+
 def attention(
     q: Tensor,
     k: Tensor,
@@ -300,37 +331,13 @@ def attention(
     Returns:
         [batch, q_len, n_heads * head_dim]
     """
-    if backend is not None:
-        selected = _resolve_backend(backend)
-        explicit = True
-    else:
-        context_backend = _current_backend.get()
-        explicit = context_backend is not None
-        # Resolve through the same chain as inference: explicit context >
-        # ASTR_BACKEND env > default.  Training calls (fwd=None, no cache)
-        # land on the CUDA backend and fall back by capability below —
-        # flash when it can handle the call, else torch SDPA.
-        selected = get_backend()
-        assert selected is not None
-
-    if not selected.supports_call(q, kv_cache, attn_mask, is_causal, fwd):
-        if explicit:
-            raise RuntimeError(
-                f"Explicitly-set backend {type(selected).__name__} cannot "
-                f"handle this attention call (shape={q.shape}, "
-                f"dtype={q.dtype}, kv_cache={'none' if kv_cache is None else 'present'}, "
-                f"attn_mask={'none' if attn_mask is None else 'present'}). "
-                f"Remove the attn_backend() context or switch to a compatible backend."
+    explicit = _resolve_backend(backend) if backend is not None else None
+    resolution = _dispatch_resolve(
+        "attention", q, kv_cache, attn_mask, is_causal, fwd, explicit=explicit
     )
-        selected = next(
-            (
-                candidate
-                for candidate in _priority_backends()
-                if candidate.supports_call(q, kv_cache, attn_mask, is_causal, fwd)
-            ),
-            _instance(TorchNativeBackend),
+    return resolution.record.obj.forward(
+        q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd
     )
-    return selected.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
 
 
 class AttentionBackend(ABC):
@@ -362,11 +369,11 @@ class AttentionBackend(ABC):
     """
 
     def __enter__(self) -> "AttentionBackend":
-        self._token = _current_backend.set(self)
+        self._token = set_override("attention", self)
         return self
 
     def __exit__(self, *exc) -> None:
-        _current_backend.reset(self._token)
+        reset_override(self._token)
 
     @classmethod
     @abstractmethod
@@ -694,12 +701,13 @@ class CudaBackend(AttentionBackend):
 class FlashAttnBackend(AttentionBackend):
     """FlashAttention backend via the optional ``flash-attn`` package.
 
-    Decode (q_len=1, contiguous cache): uses ``flash_attn_with_kvcache``,
-    which reads K/V directly from the flat pool via cache_batch_idx +
-    cache_seqlens — no materialized KV gather.
+    Decode (q_len=1, contiguous cache): writes K/V to the pool, gathers
+    flat K/V via the ``req_to_token`` page table, and calls
+    ``flash_attn_varlen_func`` over the ragged batch
+    (``qo_indptr``/``kv_indptr``).
 
-    Prefill / non-contiguous decode: falls back to KV gather +
-    ``flash_attn_func``.
+    Prefill: packed 3-D calls share the ``flash_attn_varlen_func`` path;
+    dense 4-D calls go through ``flash_attn_func`` (mask-free only).
     """
 
     @classmethod
@@ -816,3 +824,78 @@ class FlashAttnBackend(AttentionBackend):
             causal=True,
         )
         return out
+
+
+# Family registration over the generic dispatcher: the "attention" decision
+# table.  The Specs mirror each backend's ``supports_call`` exactly (a unit
+# test asserts they never drift).  The provider is re-evaluated per
+# resolution, so monkeypatching ``flash_attn_available`` (plus clearing
+# ``_priority_backends``) is honored, as before.
+
+_CLASS_TO_NAME: Dict[type, str] = {
+    CudaBackend: ATTN_BACKEND.CUDA.value,
+    FlashAttnBackend: ATTN_BACKEND.FLASH.value,
+    TorchNativeBackend: ATTN_BACKEND.TORCH_NATIVE.value,
+}
+
+_SPEC_CUDA = (
+    axis("fwd").in_("prefill", "decode")
+    & axis("has_cache").truthy()
+    & axis("ndim").eq(3)
+    & axis("dtype").in_(torch.bfloat16)
+    & axis("head_dim").in_(*CudaBackend.HEAD_DIMS)
+    & Spec.of(
+        lambda ax: is_available(f"attn_paged_{ax.get('fwd')}"), "paged kernels loaded"
+    )
+)
+
+_SPEC_FLASH = (
+    axis("dtype").in_(torch.float16, torch.bfloat16)
+    & Spec.of(lambda ax: flash_attn_available(), "flash-attn available")
+    & (
+        (
+            axis("fwd").not_none()
+            & axis("ndim").eq(3)
+            & Spec.of(
+                lambda ax: (
+                    _flash_attn is not None
+                    and hasattr(_flash_attn, "flash_attn_varlen_func")
+                ),
+                "varlen api present",
+            )
+        )
+        | (axis("fwd").none() & axis("has_mask").falsy())
+    )
+)
+
+
+def _attention_records() -> list:
+    specs = {
+        CudaBackend: _SPEC_CUDA,
+        FlashAttnBackend: _SPEC_FLASH,
+        TorchNativeBackend: Spec.always(),
+    }
+    return [
+        ImplRecord(
+            family="attention",
+            name=_CLASS_TO_NAME[type(backend)],
+            obj=backend,
+            spec=specs[type(backend)],
+            priority=position,
+        )
+        for position, backend in enumerate(_priority_backends())
+    ]
+
+
+def _reference_record() -> ImplRecord:
+    return ImplRecord(
+        family="attention",
+        name=ATTN_BACKEND.TORCH_NATIVE.value,
+        obj=_instance(TorchNativeBackend),
+        spec=Spec.always(),
+        priority=999,
+    )
+
+
+register_family("attention", _axes, _attention_records, _reference_record)
+register_env_alias("attention", "ASTR_BACKEND")

astrai/extension/backend/rotary.py

@@ -1,7 +1,9 @@
-"""Rotary embedding with auto-dispatch to CUDA kernel.
+"""Rotary embedding dispatch (family "rotary").
 
-Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
-CUDA kernel when available, falls back to torch complex multiply otherwise.
+Registered rows: the fused CUDA kernel (bf16 CUDA, inference-only) and the
+torch complex-multiply fallback (autograd-safe).  Selection runs through
+the generic dispatcher, so ``op_backend(rotary=...)`` and
+``ASTR_OPS=rotary=torch`` work exactly like for attention.
 
 Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
 freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
@@ -10,16 +12,22 @@ freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
 import torch
 from torch import Tensor
 
+from astrai.extension.dispatch import (
+    ImplRecord,
+    Spec,
+    axis,
+    register_family,
+    resolve,
+    tensor_axes,
+)
 from astrai.extension.loader import is_available
 from astrai.extension.ops.rotary import rotary_emb as _cuda_rotary
 
-_cache = {"available": None}
-
-
-def _cuda_available() -> bool:
-    if _cache["available"] is None:
-        _cache["available"] = is_available("rotary_emb")
-    return _cache["available"]
+_SPEC_CUDA = (
+    axis("device_cuda").truthy()
+    & axis("dtype").in_(torch.bfloat16)
+    & axis("grad_enabled").eq(False)
+)
 
 
 def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
@@ -33,6 +41,34 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
     return x_out.to(dtype)
 
 
+def _rotary_records() -> list:
+    return [
+        ImplRecord(
+            family="rotary",
+            name="cuda",
+            obj=_cuda_rotary,
+            spec=_SPEC_CUDA,
+            available=lambda: is_available("rotary_emb"),
+            priority=0,
+        ),
+        ImplRecord(
+            family="rotary",
+            name="torch",
+            obj=_torch_apply,
+            spec=Spec.always(),
+            priority=99,
+        ),
+    ]
+
+
+register_family(
+    "rotary",
+    lambda x, freqs_cis: tensor_axes(x),
+    _rotary_records,
+    lambda: _rotary_records()[-1],
+)
+
+
 def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     """Apply rotary embedding to x.
 
@@ -43,11 +79,4 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     Returns:
         [batch, seq_len, n_heads, head_dim] (bf16)
     """
-    if (
-        _cuda_available()
-        and not torch.is_grad_enabled()
-        and x.is_cuda
-        and x.dtype == torch.bfloat16
-    ):
-        return _cuda_rotary(x, freqs_cis)
-    return _torch_apply(x, freqs_cis)
+    return resolve("rotary", x, freqs_cis).record.obj(x, freqs_cis)

astrai/extension/dispatch.py

@@ -0,0 +1,400 @@
+"""Operator dispatch: one selection mechanism for all op families.
+
+A family registers three things with the core: an ``axes`` extractor whose
+signature mirrors the op call and snapshots whatever decision axes *that
+family* needs, an ordered list of ``ImplRecord`` rows (name, impl object,
+capability ``Spec``, machine-level ``available``), and a fallback record.
+The core defines no axes itself — each ``Spec`` predicates over the axes
+dict produced by the family's own extractor.  Resolution: explicit/context
+selection (strict — raises when incapable) > ``ASTR_OPS`` env entry (soft —
+falls through) > first capable row > family fallback.  The rows are the
+family's decision table, printable via ``explain``.
+
+Records flagged ``faithful=False`` change numerics (e.g. fp8) and are only
+reachable through an explicit selection, never the implicit chain.
+"""
+
+import contextvars
+import logging
+import os
+import threading
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+Axes = Mapping[str, Any]
+Call = Tuple[Tuple, Dict[str, Any]]
+
+
+def _fmt(value: Any) -> str:
+    return str(value)
+
+
+class Spec:
+    """Composable, self-describing predicate over a family's axes dict."""
+
+    __slots__ = ("_fn", "_desc")
+
+    def __init__(self, fn: Callable[[Axes], bool], desc: str):
+        self._fn = fn
+        self._desc = desc
+
+    def matches(self, ax: Axes) -> bool:
+        return bool(self._fn(ax))
+
+    @property
+    def description(self) -> str:
+        return self._desc
+
+    def __and__(self, other: "Spec") -> "Spec":
+        return Spec(
+            lambda ax: self._fn(ax) and other._fn(ax),
+            f"({self._desc} and {other._desc})",
+        )
+
+    def __or__(self, other: "Spec") -> "Spec":
+        return Spec(
+            lambda ax: self._fn(ax) or other._fn(ax),
+            f"({self._desc} or {other._desc})",
+        )
+
+    def __invert__(self) -> "Spec":
+        return Spec(lambda ax: not self._fn(ax), f"not({self._desc})")
+
+    @classmethod
+    def always(cls) -> "Spec":
+        return cls(lambda ax: True, "always")
+
+    @classmethod
+    def of(cls, fn: Callable[[Axes], bool], desc: str) -> "Spec":
+        return cls(fn, desc)
+
+
+class Axis:
+    """Named-axis predicate builder: ``axis("dtype").in_(torch.bfloat16)``.
+
+    Axis names belong to each family; the core never defines or inspects
+    them beyond the predicate the builder closes over.
+    """
+
+    __slots__ = ("_name",)
+
+    def __init__(self, name: str):
+        self._name = name
+
+    def in_(self, *values: Any) -> Spec:
+        rendered = ", ".join(_fmt(v) for v in values)
+        return Spec(
+            lambda ax: ax.get(self._name) in values,
+            f"{self._name} in {{{rendered}}}",
+        )
+
+    def eq(self, value: Any) -> Spec:
+        return Spec(
+            lambda ax: ax.get(self._name) == value, f"{self._name}=={_fmt(value)}"
+        )
+
+    def is_(self, value: Any) -> Spec:
+        return Spec(
+            lambda ax: ax.get(self._name) is value, f"{self._name} is {_fmt(value)}"
+        )
+
+    def none(self) -> Spec:
+        return Spec(lambda ax: ax.get(self._name) is None, f"{self._name} is None")
+
+    def not_none(self) -> Spec:
+        return Spec(
+            lambda ax: ax.get(self._name) is not None, f"{self._name} is not None"
+        )
+
+    def truthy(self) -> Spec:
+        return Spec(lambda ax: bool(ax.get(self._name)), self._name)
+
+    def falsy(self) -> Spec:
+        return Spec(lambda ax: not ax.get(self._name), f"!{self._name}")
+
+
+def axis(name: str) -> Axis:
+    """Entry point for named-axis predicates; see ``Axis``."""
+    return Axis(name)
+
+
+def tensor_axes(x: torch.Tensor, **extra: Any) -> Dict[str, Any]:
+    """Tensor-derived axes shared by most families; opt-in, extendable."""
+    return {
+        "dtype": x.dtype,
+        "device_cuda": x.is_cuda,
+        "grad_enabled": torch.is_grad_enabled(),
+        **extra,
+    }
+
+
+@dataclass(frozen=True)
+class ImplRecord:
+    """One decision-table row: an implementation plus its capability."""
+
+    family: str
+    name: str
+    obj: Any
+    spec: Spec
+    available: Callable[[], bool] = lambda: True
+    priority: int = 100
+    faithful: bool = True
+
+
+@dataclass
+class OpFamily:
+    name: str
+    axes: Callable[..., Axes]
+    provider: Callable[[], List[ImplRecord]]
+    fallback: Callable[[], ImplRecord]
+
+
+_FAMILIES: Dict[str, OpFamily] = {}
+_ENV_ALIASES: Dict[str, str] = {}
+
+_current_overrides: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
+    "astrai_op_overrides", default={}
+)
+
+_env_lock = threading.Lock()
+_env_cache: Dict[tuple, Optional[Dict[str, str]]] = {}
+_warned: set = set()
+
+
+def register_family(
+    name: str,
+    axes: Callable[..., Axes],
+    provider: Callable[[], List[ImplRecord]],
+    fallback: Callable[[], ImplRecord],
+) -> None:
+    """Register (or replace) a family; ``provider`` is re-evaluated per
+    resolution so availability changes (tests, late imports) are honored.
+    ``axes`` mirrors the op call signature and snapshots that family's
+    decision axes; unregistered handles are probed through the same args.
+    """
+    _FAMILIES[name] = OpFamily(name, axes, provider, fallback)
+
+
+def register_env_alias(family: str, varname: str) -> None:
+    """Legacy single-value env var for a family (e.g. attention →
+    ASTR_BACKEND); an ASTR_OPS entry wins when both are set."""
+    _ENV_ALIASES[family] = varname
+
+
+def _family(name: str) -> OpFamily:
+    fam = _FAMILIES.get(name)
+    if fam is None:
+        raise KeyError(f"no operator family registered under {name!r}")
+    return fam
+
+
+def _warn_once(message: str) -> None:
+    if message not in _warned:
+        _warned.add(message)
+        logger.warning(message)
+
+
+def set_override(family: str, handle: Any) -> contextvars.Token:
+    overrides = dict(_current_overrides.get())
+    overrides[family] = handle
+    return _current_overrides.set(overrides)
+
+
+def reset_override(token: contextvars.Token) -> None:
+    _current_overrides.reset(token)
+
+
+def get_override(family: str) -> Optional[Any]:
+    return _current_overrides.get().get(family)
+
+
+@contextmanager
+def op_backend(**handles: Any):
+    """Select implementations per family for the enclosed scope::
+
+        with op_backend(attention="torch_native", rotary="torch"):
+            engine.generate(...)
+
+    String handles are validated eagerly against the family's currently
+    available implementations; object handles pass through unchecked.
+    """
+    for family, handle in handles.items():
+        if isinstance(handle, str):
+            fam = _FAMILIES.get(family)
+            if fam is None:
+                raise ValueError(f"unknown operator family {family!r}")
+            if _record_for_handle(fam, handle) is None:
+                raise ValueError(f"Unknown {family} implementation: {handle!r}")
+    tokens = [set_override(f, h) for f, h in handles.items()]
+    try:
+        yield
+    finally:
+        for token in reversed(tokens):
+            reset_override(token)
+
+
+def env_overrides() -> Dict[str, str]:
+    """Merged ASTR_OPS + legacy-alias selections (family or "profile").
+
+    Cached per distinct env content; unknown families / malformed entries
+    warn once and are dropped (soft override, never fatal).
+    """
+    with _env_lock:
+        merged: Dict[str, str] = {}
+        raw = os.environ.get("ASTR_OPS", "").strip()
+        if raw:
+            key = ("ASTR_OPS", raw)
+            if key not in _env_cache:
+                parsed: Dict[str, str] = {}
+                for item in raw.split(","):
+                    key_part, sep, value = item.strip().partition("=")
+                    key_part, value = key_part.strip(), value.strip()
+                    if not sep or not key_part or not value:
+                        _warn_once(f"ASTR_OPS: ignoring malformed entry {item!r}")
+                        continue
+                    parsed[key_part] = value
+                _env_cache[key] = parsed or None
+            merged.update(_env_cache[key] or {})
+        for fam, varname in _ENV_ALIASES.items():
+            raw = os.environ.get(varname, "").strip()
+            if raw:
+                key = (varname, raw)
+                if key not in _env_cache:
+                    _env_cache[key] = {fam: raw.lower()}
+                merged.setdefault(fam, _env_cache[key][fam])
+        for fam in [f for f in merged if f not in _FAMILIES and f != "profile"]:
+            _warn_once(f"ASTR_OPS: unknown operator family {fam!r}; dropping it")
+            merged.pop(fam)
+        return merged
+
+
+def env_selection(family: str) -> Optional[str]:
+    return env_overrides().get(family)
+
+
+@dataclass(frozen=True)
+class Resolution:
+    record: ImplRecord
+    origin: str
+
+
+class ExplicitSelectionError(RuntimeError):
+    """An explicitly selected implementation cannot handle the call."""
+
+
+def _record_for_handle(fam: OpFamily, handle: Any) -> Optional[ImplRecord]:
+    records = sorted(fam.provider(), key=lambda r: r.priority)
+    if isinstance(handle, str):
+        return next((r for r in records if r.name == handle), None)
+    return next((r for r in records if r.obj is handle), None)
+
+
+def _adhoc_record(family: str, handle: Any, args: Tuple, kwargs: Dict) -> ImplRecord:
+    """Wrap an unregistered object; capability probes its own method on
+    the original call arguments."""
+    supports = getattr(handle, "supports_call", None)
+    if supports is not None:
+        spec = Spec.of(
+            lambda ax: bool(supports(*args, **kwargs)),
+            f"{type(handle).__name__}.supports_call",
+        )
+    else:
+        spec = Spec.always()
+    return ImplRecord(family, type(handle).__name__, handle, spec)
+
+
+def resolve(
+    family: str, *args: Any, explicit: Optional[Any] = None, **kwargs: Any
+) -> Resolution:
+    """Resolve one family for one call (explicit-strict / implicit-loose).
+
+    ``args``/``kwargs`` mirror the op call: the family's ``axes`` extractor
+    snapshots the decision axes from them, and unregistered handles are
+    probed through their own ``supports_call`` with the same arguments.
+    """
+    fam = _family(family)
+    ax = fam.axes(*args, **kwargs)
+
+    handle: Optional[Any] = None
+    origin = "chain"
+    if explicit is not None:
+        handle, origin = explicit, "explicit"
+    elif get_override(family) is not None:
+        handle, origin = get_override(family), "context"
+    else:
+        env_name = env_selection(family)
+        if env_name is not None:
+            handle, origin = env_name, "env"
+
+    if handle is not None:
+        record = _record_for_handle(fam, handle)
+        if record is None and not isinstance(handle, str):
+            record = _adhoc_record(family, handle, args, kwargs)
+        if record is None:
+            if origin in ("explicit", "context"):
+                raise ValueError(f"Unknown {family} implementation: {handle!r}")
+            _warn_once(f"ASTR_OPS: {family}={handle!r} is not registered; ignoring")
+        else:
+            if record.available() and record.spec.matches(ax):
+                return Resolution(record, origin)
+            if origin in ("explicit", "context"):
+                raise ExplicitSelectionError(
+                    f"Explicitly-set backend {type(record.obj).__name__} cannot "
+                    f"handle this {family} call; required: {record.spec.description}"
+                )
+
+    if handle is None and env_overrides().get("profile") == "reference":
+        return Resolution(fam.fallback(), "profile")
+
+    for record in sorted(fam.provider(), key=lambda r: r.priority):
+        if record.available() and record.faithful and record.spec.matches(ax):
+            return Resolution(record, "chain")
+    return Resolution(fam.fallback(), "fallback")
+
+
+def resolve_plan(calls: Mapping[str, Call]) -> Dict[str, Resolution]:
+    """Resolve several families at once (one decision snapshot)."""
+    return {
+        family: resolve(family, *args, **kwargs)
+        for family, (args, kwargs) in calls.items()
+    }
+
+
+def _describe_axes(ax: Axes) -> str:
+    return " ".join(f"{key}={ax[key]}" for key in sorted(ax))
+
+
+def explain(
+    family: str, *args: Any, explicit: Optional[Any] = None, **kwargs: Any
+) -> str:
+    """Human-readable decision trace for one family call."""
+    fam = _family(family)
+    ax = fam.axes(*args, **kwargs)
+    records = sorted(fam.provider(), key=lambda r: r.priority)
+    lines = [f"[{family}] {_describe_axes(ax)}"]
+    for record in records:
+        if not record.available():
+            lines.append(f"  {record.name}: SKIP unavailable")
+        elif not record.faithful:
+            lines.append(f"  {record.name}: SKIP not faithful (explicit-only)")
+        elif record.spec.matches(ax):
+            lines.append(f"  {record.name}: MATCH ({record.spec.description})")
+        else:
+            lines.append(f"  {record.name}: reject ({record.spec.description})")
+    try:
+        resolution = resolve(family, *args, explicit=explicit, **kwargs)
+        lines.append(f"  => {resolution.record.name} (origin={resolution.origin})")
+    except (ExplicitSelectionError, ValueError) as exc:
+        lines.append(f"  => ERROR: {exc}")
+    return "\n".join(lines)
+
+
+def explain_plan(calls: Mapping[str, Call]) -> str:
+    return "\n".join(
+        explain(family, *args, **kwargs) for family, (args, kwargs) in calls.items()
+    )

astrai/extension/fp8.py

@@ -31,12 +31,12 @@ import functools
 from contextvars import ContextVar, Token
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Dict, List, NamedTuple, Optional
 
 import torch
 from torch.library import Library
 
-from astrai.extension.ops.fp8 import mm_fp8, quantize
+from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
 
 # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
 FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -56,46 +56,35 @@ class FP8Format(str, Enum):
         return "e5m2" if self is FP8Format.HYBRID else self.value
 
 
+@dataclass
 class FP8Recipe:
     """Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
 
-    ``scale_from_history`` receives the operand's amax tensor (a ring window for
-    delayed scaling, the current amax for dynamic scaling) and returns the
-    quantization step. Subclasses set ``history_len`` / ``margin``.
+    ``dynamic=False`` (default) is TE-style delayed scaling: max over the
+    amax history window (amax from *previous* steps; the window trades
+    responsiveness against stability). ``dynamic=True`` is current-amax
+    scaling (torchao DYNAMIC): measure, then quantize — no history, at an
+    extra pass. ``scale_from_history`` receives the operand's amax tensor
+    (a ring window / the current amax) and returns the quantization step.
     """
 
     history_len: int = 16
     margin: int = 0
+    dynamic: bool = False
 
     def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
         peak = amax.max()
         return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
 
 
-@dataclass
-class DelayedScaling(FP8Recipe):
-    """TE-style delayed scaling: max over the amax history window (amax from
-    *previous* steps; the window trades responsiveness against stability)."""
-
-    history_len: int = 16
-    margin: int = 0
-
-
-@dataclass
-class DynamicScaling(FP8Recipe):
-    """Current-amax scaling (torchao DYNAMIC): measure, then quantize. No
-    history — the scale is derived from the same-step amax, at an extra pass."""
-
-    history_len: int = 1
-    margin: int = 0
-
-
 class _ScaleRing:
     """One operand's delayed-scaling state: a float32 buffer
-    ``[hist[n] | scale | counter]`` (views). ``update`` folds the amax
-    returned by the quantize primitive into ``hist[idx]`` and publishes the
-    next scale from the window; ``idx`` advances host-side each step. The
-    trailing slot is a legacy counter kept for state-buffer compatibility.
+    ``[hist[n] | scale | legacy | amax | done]`` (views). The quantize
+    kernel folds its fused amax into ``hist[idx]`` and publishes the next
+    scale from the window in its own last block (``fold_args`` passes the
+    buffer + recipe constants); ``idx`` advances host-side each use. The
+    ``amax``/``done`` tail slots are kernel scratch (self-cleaning across
+    launches); the legacy slot keeps state-buffer compatibility.
     """
 
     __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -103,7 +92,7 @@ class _ScaleRing:
     def __init__(self, device: torch.device, recipe: FP8Recipe):
         self.recipe = recipe
         n = recipe.history_len
-        self.state = torch.zeros(n + 2, device=device, dtype=torch.float32)
+        self.state = torch.zeros(n + 4, device=device, dtype=torch.float32)
         self.hist = self.state[:n]
         self.scale = self.state[n : n + 1]
         self.idx = 0
@@ -119,23 +108,25 @@ class _ScaleRing:
         self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
         self.initialized = True
 
-    def update(self, amax: torch.Tensor, fmt: str) -> None:
-        self.hist[self.idx].copy_(amax.reshape(()))
-        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
+    def fold_args(self, fmt: str) -> dict:
+        """Keyword arguments for quantize()'s in-kernel history fold."""
+        return {
+            "ring_state": self.state,
+            "hist_idx": self.idx,
+            "fp8_max": FP8_MAX[fmt],
+            "pow2_margin": float(2**self.recipe.margin),
+        }
 
 
-class FP8TensorMeta:
-    """Per-weight delayed-scaling state for ``w``, ``x`` and ``g``.
+class FP8TensorMeta(NamedTuple):
+    """Per-weight delayed-scaling rings for ``w``, ``x`` and ``g``.
 
-    DynamicScaling never allocates a meta; it measures the current amax inline.
+    Dynamic scaling never allocates a meta; it measures the current amax inline.
     """
 
-    __slots__ = ("w", "x", "g")
-
-    def __init__(self, device: torch.device, recipe: FP8Recipe):
-        self.w = _ScaleRing(device, recipe)
-        self.x = _ScaleRing(device, recipe)
-        self.g = _ScaleRing(device, recipe)
+    w: _ScaleRing
+    x: _ScaleRing
+    g: _ScaleRing
 
 
 @dataclass(frozen=True)
@@ -159,55 +150,29 @@ _active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
 class FP8State:
     """Global fp8 training state: per-tensor metas + out-of-region defaults.
 
-    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set
-    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.
+    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar``
+    set by ``fp8_autocast`` (see ``_active``/``_current_config``); these plain
+    attributes are the persistent defaults applied outside any region —
+    ``fp8_linear_enable`` writes ``default_enabled``. 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.default_enabled = False
-        self.default_recipe: FP8Recipe = DelayedScaling()
+        self.default_recipe: FP8Recipe = FP8Recipe()
         self.default_format: FP8Format = FP8Format.HYBRID
         self._metas: Dict[tuple, FP8TensorMeta] = {}
 
-    # Active-config views (region config if open, else the defaults).
-    @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:
+    def get_weight_meta(self, w: torch.Tensor, recipe: FP8Recipe) -> FP8TensorMeta:
         key = (w.data_ptr(), w.shape, w.dtype)
         meta = self._metas.get(key)
         if meta is None:
-            meta = FP8TensorMeta(w.device, self.recipe)
+            meta = FP8TensorMeta(
+                _ScaleRing(w.device, recipe),
+                _ScaleRing(w.device, recipe),
+                _ScaleRing(w.device, recipe),
+            )
             self._metas[key] = meta
         return meta
 
@@ -215,7 +180,7 @@ class FP8State:
         """Restore construction defaults (switch, recipe, format) and drop all
         per-weight metas — a full state reset for tests / reconfiguration."""
         self.default_enabled = False
-        self.default_recipe = DelayedScaling()
+        self.default_recipe = FP8Recipe()
         self.default_format = FP8Format.HYBRID
         self._metas.clear()
 
@@ -278,7 +243,7 @@ class fp8_autocast:
         margin: int = 0,
     ):
         if recipe is None:
-            recipe = DelayedScaling(history_len=update_interval, margin=margin)
+            recipe = FP8Recipe(history_len=update_interval, margin=margin)
         self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
         self._tokens: List[Token] = []
 
@@ -322,17 +287,17 @@ def fp8_linear_forward(
 
     Composed from the two stateless primitives: quantize x/w with the active
     scales, run the pre-quantized GEMM with the bias fused into its epilogue.
-    Delayed scaling folds
-    the returned amax into the history ring and publishes the next scale;
-    dynamic scaling measures the current amax itself. Training quantizes the
-    weight every step (the optimizer bumps its version, so there is no cast
-    cache, matching ``cached_cast``-less behavior).
+    Delayed scaling lets the quantize kernel fold the fused amax into the
+    history ring and publish the next scale in its own last block; 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).
     """
     state = fp8_state()
     if cfg is None:
         cfg = _current_config()
     fmt = cfg.fp8_format.fwd()
-    if isinstance(cfg.recipe, DynamicScaling):
+    if cfg.recipe.dynamic:
         sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
         sw = _dynamic_scale(w, cfg.recipe, fmt)
         x8, _ = quantize(x, sx.reciprocal(), fmt)
@@ -345,25 +310,25 @@ def fp8_linear_forward(
         ).reshape(*x.shape[:-1], w.size(0))
         return out, sx, sw
 
-    meta = state.get_weight_meta(w)
+    meta = state.get_weight_meta(w, cfg.recipe)
     if not meta.w.initialized:
         meta.w.seed(w, fmt)
     if not meta.x.initialized:
         meta.x.seed(x, fmt)
     sx, sw = meta.x.scale.clone(), meta.w.scale.clone()
-    x8, amax_x = quantize(x, sx.reciprocal(), fmt)
+    # The clones feed this call's kernels (stream-ordered before the in-kernel
+    # fold overwrites the ring scale slots); the fp8 quantize kernel folds the
+    # amax into the history window and publishes the next scale itself.
+    x8, _ = quantize(x, sx.reciprocal(), fmt, **meta.x.fold_args(fmt))
     if _is_fp8(w.dtype):
-        w8, amax_w = w, None
+        w8 = w
     else:
-        w8, amax_w = quantize(w, sw.reciprocal(), fmt)
+        w8, _ = quantize(w, sw.reciprocal(), fmt, **meta.w.fold_args(fmt))
     out = mm_fp8(
         x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
     ).reshape(*x.shape[:-1], w.size(0))
-    meta.x.update(amax_x, fmt)
-    if amax_w is not None:
-        meta.w.update(amax_w, fmt)
     meta.x.advance()
-    if amax_w is not None:
+    if not _is_fp8(w.dtype):
         meta.w.advance()
     return out, sx, sw
 
@@ -385,8 +350,8 @@ class _LinearFp8(torch.autograd.Function):
         ctx.save_for_backward(x, w, sx, sw)
         ctx.fmt_bwd = cfg.fp8_format.bwd()
         ctx.recipe = cfg.recipe
-        ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
-        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w)
+        ctx.is_dynamic = cfg.recipe.dynamic
+        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w, cfg.recipe)
         return out
 
     @staticmethod
@@ -407,16 +372,32 @@ class _LinearFp8(torch.autograd.Function):
                 meta.g.seed(g2, fmt)
             sg = meta.g.scale.clone()
             sw, sx = _sw_fwd, _sx_fwd
-        g8, amax_g = quantize(g2, sg.reciprocal(), fmt)
-        x8, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt)
-        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
-        grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape)  # g8[m,n] @ w8[n,k]
-        grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True)  # g8.T @ x8
-        grad_b = g2.sum(0).to(torch.bfloat16)
+        # Backward GEMMs route through the NT fast path via transposed
+        # quantize outputs: g8 [m,n] with w8T [k,n] (trans_b=True) gives
+        # grad_x, g8T [n,m] with x8T [k,m] gives grad_w — no NN-swap or TT
+        # crosswise kernel in the training path. g is consumed in both
+        # orientations, so quantize_dual's single pass feeds both.
+        # The g quantize folds the gradient amax into its ring in-kernel;
+        # the x8T/w8T orientation copies discard amax (those rings were
+        # folded at forward time).
+        g8, g8T, _ = quantize_dual(g2, sg.reciprocal(), fmt, **meta.g.fold_args(fmt))
+        x8T, _ = quantize(
+            x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, transposed=True
+        )
+        if _is_fp8(w.dtype):
+            # Pre-quantized weight has no transposed copy: keep the swap
+            # path for grad_x (grad_w is unaffected).
+            grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape)
+        else:
+            w8T, _ = quantize(w, sw.reciprocal(), fmt, transposed=True)
+            grad_x = mm_fp8(g8, w8T, sg * sw, trans_b=True).reshape(x.shape)
+        grad_w = mm_fp8(g8T, x8T, sg * sx, trans_b=True)  # g8.T @ x8
+        # bias-free linears must not pay the column-sum
+        # reduce: g2.sum(0) is another full read of the gradient.
+        grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None
         if not ctx.is_dynamic:
-            meta.g.update(amax_g, fmt)
             meta.g.advance()
-        return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
+        return grad_x, grad_w, grad_b
 
 
 # ---------------------------------------------------------------------------

astrai/extension/ops/fp8.py

@@ -1,14 +1,20 @@
 """FP8 CUDA kernel interface adapter (the only module touching the pybind).
 
-Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
+Attention-style thin wrappers: one Python entry per binding, called directly
+— no torch.library dispatch layer. Optional arguments (``ring_state``,
+``bias``) keep native Optional semantics at the pybind boundary, and
+in-place buffer updates (the delayed-scaling ring fold, like attention's
+KV-cache appends) happen on-stream without mutation declarations. CUDA-only:
+non-CUDA or unsupported inputs raise from the binding's TORCH_CHECKs.
 
-- ``quantize(x, scale, fmt) -> (x8, amax)`` — BF16/FP16/FP32 → FP8 with fused amax
+- ``quantize(x, scale, fmt, transposed=False) -> (x8|x8T, amax)`` — BF16/FP16/FP32
+  → FP8 with fused amax (``transposed`` picks the orientation; arity is fixed)
+- ``quantize_dual(x, scale, fmt) -> (x8, x8T, amax)`` — both orientations, one read
 - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
 
-Scale semantics: scales are *quantization steps* — the value divided out when
-quantizing (``x8 = x / scale``). Every primitive computes its own inverse
-internally; callers never pass ``scale_inv``. ``amax`` values are *returned*,
-never passed as output arguments. ``fmt`` is ``"e4m3"`` or ``"e5m2"``.
+``scale`` is the quantization multiplier (device scalar); ``fmt`` is
+``"e4m3"`` or ``"e5m2"``. ``amax`` values are *returned*, never passed as
+output arguments.
 
 Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
 this module is stateless.
@@ -17,7 +23,6 @@ this module is stateless.
 from typing import Optional, Tuple
 
 import torch
-from torch.library import custom_op
 
 from astrai.extension.loader import get_module
 
@@ -32,120 +37,63 @@ def _fmt_int(fmt: str) -> int:
         raise ValueError(f"unsupported fp8 format {fmt!r} (expected 'e4m3' or 'e5m2')")
 
 
-def _fmt_name(fmt: int) -> str:
-    if fmt == 0:
-        return "e4m3"
-    if fmt == 1:
-        return "e5m2"
-    raise ValueError(f"unsupported quantization type {fmt!r}")
-
-
-def _fmt_dtype(fmt: str) -> torch.dtype:
-    return torch.float8_e5m2 if _fmt_int(fmt) else torch.float8_e4m3fn
-
-
-@custom_op("custom::fp8_quantize", mutates_args=())
-def fp8_quantize(
-    x: torch.Tensor, scale: torch.Tensor, fmt: int
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; ``scale`` is a multiplier."""
-
-
-@fp8_quantize.register_fake
-def _fp8_quantize_fake(x, scale, fmt):
-    dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
-    return (
-        torch.empty(x.shape, device=x.device, dtype=dtype),
-        torch.empty(1, device=x.device, dtype=torch.float32),
-    )
-
-
-_QUANT_INPUT_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
-
-
-@fp8_quantize.register_kernel("cuda")
-def _fp8_quantize_cuda(x, scale, fmt):
-    if x.dtype not in _QUANT_INPUT_DTYPES:
-        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
-    return get_module("fp8_ops").quantize(x, scale, int(fmt))
-
-
-@fp8_quantize.register_kernel("cpu")
-def _fp8_quantize_cpu(x, scale, fmt):
-    x8 = (x.float() * scale).to(_fmt_dtype(_fmt_name(fmt)))
-    amax = x.abs().amax().float().reshape(1).clamp_min(1e-12)
-    return x8, amax
-
-
-@custom_op("custom::fp8_gemm", mutates_args=())
-def fp8_gemm(
-    a: torch.Tensor,
-    b: torch.Tensor,
-    scale: torch.Tensor,
-    trans_a: int = 0,
-    trans_b: int = 0,
-    bias: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """FP8 GEMM: ``a @ b * scale (+ bias)`` with FP32 accumulation.
-
-    2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules).
-    ``bias`` (bf16, length n) fuses into the epilogue in fp32 before the
-    single bf16 rounding. The result is always BF16; FP8 output is a
-    separate quantize operation.
-    """
-
-
-@fp8_gemm.register_fake
-def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0, bias=None):
-    dtype = torch.bfloat16
-    rows = a.size(2) if trans_a else a.size(1)
-    cols = b.size(1) if trans_b else b.size(2)
-    batches = [t.size(0) for t in (a, b) if t.dim() == 3]
-    shape = (max(batches), rows, cols) if batches else (rows, cols)
-    return torch.empty(shape, device=a.device, dtype=dtype)
-
-
-@fp8_gemm.register_kernel("cuda")
-def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0, bias=None):
-    if a.dtype != b.dtype or a.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
-        raise TypeError(
-            f"fp8 GEMM requires matching fp8 inputs, got {a.dtype}/{b.dtype}"
-        )
-    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
-
-
-@fp8_gemm.register_kernel("cpu")
-def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
-    aa = a.float().transpose(-2, -1) if trans_a else a.float()
-    bb = b.float().transpose(-2, -1) if trans_b else b.float()
-    acc = aa @ bb * scale
-    if bias is not None and bias.numel() > 0:
-        acc = acc + bias.float()
-    return acc.to(torch.bfloat16)
-
-
 def quantize(
-    x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    fmt: str = "e4m3",
+    transposed: bool = False,
+    ring_state: Optional[torch.Tensor] = None,
+    hist_idx: int = 0,
+    fp8_max: float = 448.0,
+    pow2_margin: float = 1.0,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; returns
-    ``(x8, amax)``.
+    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax.
 
     ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
     E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
+    ``transposed=True`` swaps ``x8`` for ``x8T``, the ``[cols][rows]``
+    row-major transpose of the quantized input — the K-contiguous operand
+    orientation NT GEMMs want — at the same 2-tuple arity.
+
+    ``ring_state`` (a 1D float32 CUDA buffer laid out
+    ``[hist n | scale | legacy | amax | done]``) switches on the in-kernel
+    delayed-scaling fold: the kernel's last block folds the amax into
+    ``hist[hist_idx]`` and publishes the next scale as
+    ``max(hist) / fp8_max / pow2_margin`` — the returned ``amax`` is then the
+    self-cleaned persistent slot (reads zero). None keeps the classic
+    fresh-amax return.
     """
-    # Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a
-    # 512-wide GEMM): real CUDA tensors of a supported dtype go straight to
-    # the extension. Fake/subclass tensors and non-CUDA inputs keep the
-    # custom_op route so torch.compile / meta / fake-tensor tracing and the
-    # CPU fallback behave exactly as before.
-    if (
-        type(x) is torch.Tensor
-        and x.is_cuda
-        and x.dtype in _QUANT_INPUT_DTYPES
-        and fmt in _FMT_TO_INT
-    ):
-        return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt])
-    return fp8_quantize(x, scale, _fmt_int(fmt))
+    return get_module("fp8_ops").quantize(
+        x,
+        scale,
+        _fmt_int(fmt),
+        transposed,
+        ring_state,
+        hist_idx,
+        fp8_max,
+        pow2_margin,
+    )
+
+
+def quantize_dual(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    fmt: str = "e4m3",
+    ring_state: Optional[torch.Tensor] = None,
+    hist_idx: int = 0,
+    fp8_max: float = 448.0,
+    pow2_margin: float = 1.0,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Dual-orientation quantize: one read of ``x`` produces both the
+    row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
+    consumed by GEMMs in both orientations (backward ``g``).
+
+    ``ring_state`` switches on the in-kernel delayed-scaling fold exactly as
+    in :func:`quantize`.
+    """
+    return get_module("fp8_ops").quantize_dual(
+        x, scale, _fmt_int(fmt), ring_state, hist_idx, fp8_max, pow2_margin
+    )
 
 
 def mm_fp8(
@@ -165,15 +113,4 @@ def mm_fp8(
     kernel epilogue in fp32 — no separate elementwise pass. The result is
     BF16; FP8 output is a separate quantize operation.
     """
-    # Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep
-    # validation identical on the direct route (bias may be None — the
-    # binding resolves it to the no-bias path).
-    if (
-        type(a) is torch.Tensor
-        and a.is_cuda
-        and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
-    ):
-        return get_module("fp8_ops").mm_fp8(
-            a, b, scale, int(trans_a), int(trans_b), bias
-        )
-    return fp8_gemm(a, b, scale, trans_a, trans_b, bias)
+    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)

astrai/inference/runtime/executor.py

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

astrai/inference/workspace.py

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

astrai/logging.py

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

astrai/parallel/setup.py

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

astrai/trainer/train_callback.py

@@ -54,8 +54,11 @@ class TrainCallback(Protocol):
     def on_batch_end(self, context: TrainContext):
         """Called at the end of each batch."""
 
-    def on_optimizer_step(self, context: TrainContext):
-        """Called on every optimizer step (sync step only)."""
+    def before_optimizer_step(self, context: TrainContext):
+        """Called immediately before every optimizer step (sync step only)."""
+
+    def after_optimizer_step(self, context: TrainContext):
+        """Called after the optimizer and scheduler step (sync step only)."""
 
     def on_error(self, context: TrainContext):
         """Called when an error occurs during training."""
@@ -82,7 +85,7 @@ class GradientClippingCallback(TrainCallback):
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
 
-    def on_optimizer_step(self, context: TrainContext):
+    def before_optimizer_step(self, context: TrainContext):
         context.grad_norm = context.executor.clip_grad_norm(
             context.model, self.max_grad_norm
         )
@@ -116,6 +119,8 @@ class GradientCheckpointingCallback(TrainCallback):
             del module._original_forward
 
     def on_train_begin(self, context: TrainContext):
+        if not self.modules:
+            return
         context.model.apply(self._enable)
         logger.info("Gradient checkpointing enabled")
 
@@ -168,7 +173,7 @@ class CheckpointCallback(TrainCallback):
                 )
                 context.checkpoint.save(save_path)
 
-    def on_batch_end(self, context: TrainContext):
+    def after_optimizer_step(self, context: TrainContext):
         if context.optimizer_step - self.last_ckpt_step >= self.interval:
             self._save_checkpoint(context)
 
@@ -214,7 +219,7 @@ class ProgressBarCallback(TrainCallback):
         )
 
     @only_on_rank(0)
-    def on_optimizer_step(self, context: TrainContext):
+    def before_optimizer_step(self, context: TrainContext):
         postfix = {
             "step": f"{context.optimizer_step:d}",
             "loss": f"{context.loss:.4f}",
@@ -341,7 +346,7 @@ class MetricCallback(TrainCallback):
             for log in self.log_cache:
                 f.write(json.dumps(log) + "\n")
 
-    def on_optimizer_step(self, context):
+    def before_optimizer_step(self, context):
         context.grad_snr_tracker.update(context.model)
 
         if (

astrai/trainer/trainer.py

@@ -93,7 +93,7 @@ class Trainer:
                         self._call_callbacks("on_batch_end", context)
 
                         if executor.sync_gradients:
-                            self._call_callbacks("on_optimizer_step", context)
+                            self._call_callbacks("before_optimizer_step", context)
                             context.optimizer.step()
                             context.strategy.on_optimizer_step()
                             context.optimizer.zero_grad()
@@ -101,6 +101,8 @@ class Trainer:
                             if context.scheduler:
                                 context.scheduler.step()
 
+                            self._call_callbacks("after_optimizer_step", context)
+
                 self._call_callbacks("on_epoch_end", context)
 
             if context.stop_requested:

csrc/CMakeLists.txt

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

csrc/kernels/attention/common.h

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

csrc/kernels/attention/decode_split_kv.cuh

@@ -2,8 +2,8 @@
 #include <cuda_bf16.h>
 #include <float.h>
 #include "common.h"
+#include "common/reduce.cuh"
 #include "layout_policies.cuh"
-#include "../common/reduce.cuh"
 
 namespace astrai {
 namespace attention {

csrc/kernels/attention/dispatchers.cuh

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

csrc/kernels/attention/layout_policies.cuh

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

csrc/kernels/attention/mma_utils.cuh

@@ -3,8 +3,8 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
-#include "../common/cp_async.cuh"
-#include "../common/mma.cuh"
+#include "common/cp_async.cuh"
+#include "common/mma.cuh"
 
 // Predicated cp.async (4-operand form) requires CUDA 11.2+.
 // bf16 mma.sync requires sm_80+ (guarded at build time by ASTRAI_NO_MMA).

csrc/kernels/attention/prefill_split_q.cuh

@@ -2,8 +2,8 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "common.h"
+#include "common/reduce.cuh"
 #include "layout_policies.cuh"
-#include "../common/reduce.cuh"
 
 namespace astrai {
 namespace attention {

csrc/kernels/attention/prefill_split_q_mma.cuh

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

csrc/kernels/fp8/common.h

@@ -11,47 +11,23 @@
 namespace astrai {
 namespace fp8 {
 
-// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
-// E5M2 (gradient / large dynamic range, max 57344).
+// Compile-time FP8 format: E4M3 (forward, max 448) or E5M2 (gradients,
+// max 57344).
 enum class FP8Format : int {
     E4M3 = 0,
     E5M2 = 1,
 };
 
-// Operand memory layouts as types (CUTLASS-style tags). The tag names the
-// storage order of the raw buffer relative to the operand's canonical GEMM
-// matrix — A is [M][K], B is [K][N]:
-//   A RowMajor = [M][K] storage (K-contiguous rows; the default)
-//   A ColMajor = [K][M] storage (M-contiguous; A^T)
-//   B RowMajor = [K][N] storage (N-contiguous; the plain a @ b operand)
-//   B ColMajor = [N][K] storage (K-contiguous; the nn.Linear weight layout)
-// Empty tags: selection happens by type at compile time (see load_operand_tile).
+// Operand storage tags (CUTLASS-style) relative to the canonical matrices
+// A [M][K] / B [K][N]: A RowMajor = [M][K] (default), A ColMajor = [K][M],
+// B RowMajor = [K][N], B ColMajor = [N][K] (the nn.Linear weight). Selection
+// is by type at compile time (see gemm.cuh's stage loads).
 struct RowMajor {};
 struct ColMajor {};
 
-// Transpose of a layout tag: the same buffer with the rows and contract dims
-// swapped. B's tag is relative to the canonical [K][N] GEMM matrix, so the
-// stage-load (which views any operand as [rows][contract]) sees the transposed
-// tag — this trait makes that inversion explicit.
-template <typename Layout>
-struct transpose_layout;
-template <>
-struct transpose_layout<RowMajor> {
-    using type = ColMajor;
-};
-template <>
-struct transpose_layout<ColMajor> {
-    using type = RowMajor;
-};
-template <typename Layout>
-using transpose_layout_t = typename transpose_layout<Layout>::type;
-
-// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
-// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
-// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
-// warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128 CTA, or 32x32 on the
-// cuBLAS-style 64x64 small CTA that lifts small-shape occupancy) and the
-// cp.async pipeline depth.
+// Compile-time tile configuration, mirroring KernelTraits in the attention
+// kernels: CTA tile, warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128
+// CTA, 32x32 on the 64x64 small CTA) and cp.async pipeline depth.
 template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
           int WarpM = 64, int WarpN = 32>
 struct Fp8GemmTraits {
@@ -67,10 +43,8 @@ struct Fp8GemmTraits {
         kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
     static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
 
-    // Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. 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__.
+    // Derived geometry: warp tiles tile the CTA. The smem budget is
+    // layout-aware, so it lives in Fp8GemmSmem (gemm.cuh).
     static constexpr int kWarpsM = BlockM / WarpM;
     static constexpr int kWarpsN = BlockN / WarpN;
     static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
@@ -80,54 +54,75 @@ struct Fp8GemmTraits {
                   "warp tile must be a multiple of the m16n8 MMA shape");
 };
 
-// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
-// with fused amax.
+// Quantize output orientation: RowMajor = x8 only; Transposed = the
+// [cols][rows] x8T only; Dual = both from a single read. Transposed/Dual
+// produce K-contiguous operands so crosswise consumers (backward
+// grad_x / grad_w) route through the NT fast path.
+enum class QuantLayout : int {
+    RowMajor = 0,
+    Transposed = 1,
+    Dual = 2,
+};
+
+// Quantize-kernel parameter POD: float input -> FP8 with fused amax.
 struct FP8QuantizeParams {
-    // Float input and FP8 output buffers; scale is the quantization
-    // multiplier (device scalar). amax (may be null) is zero-initialized by
-    // the binding and receives the raw-domain absolute maximum.
     const void* __restrict__ input_ptr = nullptr;
     void* __restrict__ output_ptr = nullptr;
+    void* __restrict__ output_transposed_ptr = nullptr;  // [cols][rows]
+    QuantLayout out_layout = QuantLayout::RowMajor;
 
-    const float* __restrict__ scale = nullptr;
-    float* __restrict__ amax = nullptr;
+    const float* __restrict__ scale = nullptr;  // device multiplier
+    float* __restrict__ amax = nullptr;         // raw-domain max out
 
-    // Element count (only the elementwise quantize kernel uses it).
+    // Optional delayed-scaling ring fold: when fold_ring is set, the kernel's
+    // last-finishing block folds the final amax into hist[hist_idx], reduces
+    // the window and publishes the next scale — replacing the host-side
+    // update chain. amax then points at a persistent self-cleaning slot
+    // (zeroed by the same last block) inside the caller's ring state.
+    bool fold_ring = false;
+    float* __restrict__ hist = nullptr;  // [hist_len] amax history window
+    float* __restrict__ scale_out = nullptr;
+    unsigned int* __restrict__ done = nullptr;  // block-completion counter
+    int hist_len = 0;
+    int hist_idx = 0;
+    float fp8_max = 448.0f;   // scale = max(hist) / fp8_max / pow2_margin
+    float pow2_margin = 1.0f;
+
+    // Element count (elementwise kernel); the tiled kernel views the same
+    // buffer as [rows][cols] row-major.
     int total = 0;
+    int rows = 0;
+    int cols = 0;
 };
 
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
-// through the pre-quantized GEMM kernels. Each kernel touches only the
-// fields it needs; buffers are raw pointers packed by the torch binding.
-// Pointer members default to null so optional paths cannot hold garbage.
+// through the kernels; each kernel touches only the fields it needs.
 struct FP8Params {
-    // Inputs: a/b are FP8 for the pre-quantized path. Scales are
-    // quantization steps (device scalars).
-    // Optional bf16 bias broadcast over output rows (fused into the epilogue
-    // before the bf16 rounding, so it adds in fp32 — one rounding fewer than
-    // the separate out + bias elementwise kernel it replaces). Null disables.
+    // FP8 operands + output; scales are quantization steps (device
+    // scalars). Optional bf16 bias fuses into the epilogue (fp32 add before
+    // the single bf16 rounding); null disables.
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
     const void* __restrict__ bias_ptr = nullptr;
     void* __restrict__ out_ptr = nullptr;
 
     const float* __restrict__ scale = nullptr;
-    // Shapes. `int` covers every realistic LLM shape; the kernels promote
-    // to int64 for all pointer arithmetic.
-    int m, n, k;
+    // NN-swap mode (canonicalize_gemm): the kernel computes the transposed
+    // problem and the epilogue scatters D[row][col] to out[col * p.m + row]
+    // in the caller's [M][N] buffer. Zero in the plain orientation.
+    int out_transposed = 0;
+    int m, n, k;  // int covers LLM shapes; kernels promote to int64
 
-    // Batched (bmm) geometry: grid.z slices step the operand/output pointers
-    // by these element strides (0 broadcasts the operand across batches).
+    // Batched (bmm) geometry: grid.z steps these element strides (0
+    // broadcasts the operand across batches).
     int batch = 1;
     int64_t a_batch_stride = 0;
     int64_t b_batch_stride = 0;
     int64_t out_batch_stride = 0;
 
-    // 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 dims (row strides) of A and B; the binding packs
+    // them so the kernel reads each buffer naturally or transposed per the
+    // LayoutA/LayoutB tags.
     int a_ld, b_ld;
 };
 

csrc/kernels/fp8/gemm/epilogue.cuh

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

csrc/kernels/fp8/gemm/load.cuh

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

csrc/kernels/fp8/gemm/mainloop.cuh

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

csrc/kernels/fp8/gemm/policy.cuh

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

csrc/kernels/fp8/gemm/scheduler.cuh

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

csrc/kernels/fp8/ops.cu

@@ -1,4 +1,4 @@
-// CUDA bindings for the two stateless FP8 primitives.
+// CUDA bindings for the stateless FP8 quantize/GEMM primitives.
 
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -6,10 +6,9 @@
 
 #include <cstdint>
 #include <mutex>
-#include <tuple>
 #include <unordered_map>
 
-#include "../common/device.cuh"
+#include "common/device.cuh"
 #include "gemm.cuh"
 #include "quantize.cuh"
 
@@ -46,58 +45,13 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input) {
                 "scale must be a CUDA float32 scalar on the input device");
 }
 
-void pack_quantize(FP8QuantizeParams& p, const void* input, void* output,
-                   const torch::Tensor& scale, torch::Tensor& amax,
-                   int64_t total) {
-    p.input_ptr = input;
-    p.output_ptr = output;
-    p.scale = scale.data_ptr<float>();
-    p.amax = amax.data_ptr<float>();
-    p.total = static_cast<int>(total);
-}
-
-void pack_gemm(FP8Params& p, const void* a, const void* b, void* output,
-               const torch::Tensor& scale, int64_t m, int64_t n, int64_t k,
-               int64_t a_ld, int64_t b_ld) {
-    p.a_ptr = a;
-    p.b_ptr = b;
-    p.out_ptr = output;
-    p.scale = scale.data_ptr<float>();
-    p.m = static_cast<int>(m);
-    p.n = static_cast<int>(n);
-    p.k = static_cast<int>(k);
-    p.a_ld = static_cast<int>(a_ld);
-    p.b_ld = static_cast<int>(b_ld);
-}
-
-template <FP8Format Fmt, int Variant>
-void launch_variant(const FP8Params& p, cudaStream_t stream) {
-    using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>;
-    using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>;
-    launch_fp8_gemm<Fmt, LayoutA, LayoutB>(p, stream);
-}
-
-template <FP8Format Fmt>
-void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a,
-                   bool trans_b) {
-    const int variant = (static_cast<int>(trans_a) << 1) |
-                        static_cast<int>(trans_b);
-    switch (variant) {
-        case 0: launch_variant<Fmt, 0>(p, stream); break;
-        case 1: launch_variant<Fmt, 1>(p, stream); break;
-        case 2: launch_variant<Fmt, 2>(p, stream); break;
-        case 3: launch_variant<Fmt, 3>(p, stream); break;
-    }
-}
-
 // Inner-layout resolution for one GEMM operand. The user flag names the
-// math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
-// the storage may independently be a col-major view (.t() of a contiguous
+// math (0 = last two dims are [rows][contract], 1 = transposed); the
+// storage may independently be a col-major view (.t() of a contiguous
 // buffer), which folds into the returned dispatch flag at zero copy — the
 // kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the
-// user flag only; the fold never swaps them (see the layout table in
-// gemm.cuh). Tensors whose inner dims are neither natural layout fall back
-// to .contiguous().
+// user flag only. Tensors whose inner dims are neither natural layout fall
+// back to .contiguous().
 bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
                      int64_t& batch_stride, torch::Tensor& storage) {
     torch::Tensor t = t_in;
@@ -115,11 +69,42 @@ bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
     return flag ^ col_major;
 }
 
-}  // namespace
+// Dtype dispatch over the unified quantize launcher.
+template <bool Tiled, FP8Format Fmt>
+void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p,
+                      cudaStream_t stream) {
+    switch (x.scalar_type()) {
+    case torch::kHalf:
+        launch_fp8_quantize<Fmt, __half, Tiled>(p, stream);
+        break;
+    case torch::kFloat32:
+        launch_fp8_quantize<Fmt, float, Tiled>(p, stream);
+        break;
+    default:
+        launch_fp8_quantize<Fmt, __nv_bfloat16, Tiled>(p, stream);
+    }
+}
 
-std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
-                                                  torch::Tensor scale,
-                                                  int64_t fmt) {
+template <bool Tiled>
+void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p,
+                         bool e5m2, cudaStream_t stream) {
+    if (e5m2)
+        launch_for_dtype<Tiled, FP8Format::E5M2>(x, p, stream);
+    else
+        launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
+}
+
+// Shared binding body for the two quantize entry points: RowMajor /
+// Transposed (single output) serve quantize(), Dual (both orientations from
+// one read) serves quantize_dual(). A ring tensor switches
+// on the in-kernel delayed-scaling fold: state layout
+// [hist n | scale | legacy | amax | done-as-int], and the returned amax is
+// the (self-cleaned) persistent slot. Without it, amax is reduced into a
+// fresh buffer armed by a driver memset — cheaper than the zeros() fill
+// kernel.
+py::object quantize_impl(torch::Tensor x, torch::Tensor scale, int64_t fmt,
+                         QuantLayout layout, py::object ring, int64_t hist_idx,
+                         double fp8_max, double pow2_margin) {
     TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
     TORCH_CHECK(x.scalar_type() == torch::kBFloat16 ||
                     x.scalar_type() == torch::kHalf ||
@@ -128,43 +113,103 @@ std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
     TORCH_CHECK(fmt == static_cast<int64_t>(FP8Format::E4M3) ||
                     fmt == static_cast<int64_t>(FP8Format::E5M2),
                 "unsupported quantization type: expected E4M3 (0) or E5M2 (1)");
+    TORCH_CHECK(layout == QuantLayout::RowMajor || x.dim() >= 2,
+                "transposed quantize layouts need a 2D+ tensor");
     check_scale(scale, x);
     check_fp8_device(x);
     const at::cuda::OptionalCUDAGuard guard(x.device());
     auto stream = at::cuda::getCurrentCUDAStream();
     auto input = x.contiguous();
-    auto output = torch::empty_like(
-        input, input.options().dtype(fmt ? torch::kFloat8_e5m2
-                                          : torch::kFloat8_e4m3fn));
-    auto amax = torch::zeros({1}, input.options().dtype(torch::kFloat32));
-    FP8QuantizeParams p;
-    pack_quantize(p, input.data_ptr(), output.data_ptr(), scale, amax,
-                  input.numel());
-    const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
-    if (x.scalar_type() == torch::kHalf) {
-        if (e5m2)
-            launch_fp8_quantize<FP8Format::E5M2, __half>(p, stream.stream());
-        else
-            launch_fp8_quantize<FP8Format::E4M3, __half>(p, stream.stream());
-    } else if (x.scalar_type() == torch::kFloat32) {
-        if (e5m2)
-            launch_fp8_quantize<FP8Format::E5M2, float>(p, stream.stream());
-        else
-            launch_fp8_quantize<FP8Format::E4M3, float>(p, stream.stream());
+    auto out_opts = input.options().dtype(
+        fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn);
+    torch::Tensor amax;
+    float *ring_hist = nullptr, *ring_scale_out = nullptr;
+    unsigned int* ring_done = nullptr;
+    int ring_len = 0;
+    if (!ring.is_none()) {
+        auto st = ring.cast<torch::Tensor>();
+        TORCH_CHECK(st.is_cuda() && st.dim() == 1 &&
+                        st.scalar_type() == torch::kFloat32,
+                    "ring state must be a 1D float32 CUDA tensor");
+        const int64_t n = st.numel() - 4;
+        TORCH_CHECK(n > 0 && hist_idx >= 0 && hist_idx < n,
+                    "ring state too small or hist_idx out of range");
+        float* base = st.data_ptr<float>();
+        amax = st.narrow(0, n + 2, 1);
+        ring_hist = base;
+        ring_scale_out = base + n;
+        ring_done = reinterpret_cast<unsigned int*>(base + n + 3);
+        ring_len = static_cast<int>(n);
     } else {
-        if (e5m2)
-            launch_fp8_quantize<FP8Format::E5M2, __nv_bfloat16>(
-                p, stream.stream());
-        else
-            launch_fp8_quantize<FP8Format::E4M3, __nv_bfloat16>(
-                p, stream.stream());
+        amax = torch::empty({1}, input.options().dtype(torch::kFloat32));
+        cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream());
     }
+
+    FP8QuantizeParams p;
+    p.input_ptr = input.data_ptr();
+    p.scale = scale.data_ptr<float>();
+    p.amax = amax.data_ptr<float>();
+    if (ring_hist) {
+        p.fold_ring = true;
+        p.hist = ring_hist;
+        p.scale_out = ring_scale_out;
+        p.done = ring_done;
+        p.hist_len = ring_len;
+        p.hist_idx = static_cast<int>(hist_idx);
+        p.fp8_max = static_cast<float>(fp8_max);
+        p.pow2_margin = static_cast<float>(pow2_margin);
+    }
+    p.total = static_cast<int>(input.numel());
+    p.out_layout = layout;
+    p.rows = static_cast<int>(input.size(-2));
+    p.cols = static_cast<int>(input.size(-1));
+    torch::Tensor output, output_t;
+    if (layout != QuantLayout::Transposed) {
+        output = torch::empty_like(input, out_opts);
+        p.output_ptr = output.data_ptr();
+    }
+    if (layout != QuantLayout::RowMajor) {
+        output_t = torch::empty({input.size(-1), input.size(-2)}, out_opts);
+        p.output_transposed_ptr = output_t.data_ptr();
+    }
+    const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
+    if (layout == QuantLayout::RowMajor)
+        launch_quantize_for<false>(input, p, e5m2, stream.stream());
+    else
+        launch_quantize_for<true>(input, p, e5m2, stream.stream());
     C10_CUDA_CHECK(cudaGetLastError());
-    return {output, amax};
+    if (layout == QuantLayout::Dual)
+        return py::make_tuple(output, output_t, amax);
+    return py::make_tuple(
+        layout == QuantLayout::Transposed ? output_t : output, amax);
+}
+
+}  // namespace
+
+// Single-orientation quantize binding: row-major x8, or its [cols][rows]
+// transpose when transposed is set — the K-contiguous operand orientation
+// NT GEMMs want. Returns (x8|x8T, amax).
+py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
+                    bool transposed, py::object ring, int64_t hist_idx,
+                    double fp8_max, double pow2_margin) {
+    const QuantLayout layout =
+        transposed ? QuantLayout::Transposed : QuantLayout::RowMajor;
+    return quantize_impl(x, scale, fmt, layout, ring, hist_idx, fp8_max,
+                         pow2_margin);
+}
+
+// Dual-orientation quantize binding: one read of x produces both the
+// row-major x8 and its transpose (plus amax), for tensors consumed by GEMMs
+// in both orientations (backward g). Returns (x8, x8T, amax).
+py::object quantize_dual(torch::Tensor x, torch::Tensor scale, int64_t fmt,
+                         py::object ring, int64_t hist_idx, double fp8_max,
+                         double pow2_margin) {
+    return quantize_impl(x, scale, fmt, QuantLayout::Dual, ring, hist_idx,
+                         fp8_max, pow2_margin);
 }
 
 torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
-                     int64_t trans_a, int64_t trans_b, torch::Tensor bias) {
+                     bool trans_a, bool trans_b, py::object bias) {
     TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
     TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn ||
                     a.scalar_type() == torch::kFloat8_e5m2,
@@ -174,6 +219,18 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
                     (b.dim() == 2 || b.dim() == 3),
                 "a and b must be 2D or 3D (batched)");
     TORCH_CHECK(a.device() == b.device(), "a and b must share device");
+    // Python None and an omitted argument both mean "no bias" — an undefined
+    // tensor below. (py::isinstance<torch::Tensor> is false for real tensors
+    // here — torch's caster registers no pybind type info — so validate by
+    // attempting the cast itself.)
+    torch::Tensor bias_t;
+    if (!bias.is_none()) {
+        try {
+            bias_t = bias.cast<torch::Tensor>();
+        } catch (const py::cast_error&) {
+            TORCH_CHECK(false, "bias must be a torch.Tensor or None");
+        }
+    }
     check_scale(scale, a);
     check_fp8_device(a);
     const at::cuda::OptionalCUDAGuard guard(a.device());
@@ -190,10 +247,8 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
 
     torch::Tensor a_st, b_st;
     int64_t a_ld, b_ld, a_bstride, b_bstride;
-    const bool tag_a =
-        resolve_operand(a, trans_a != 0, a_ld, a_bstride, a_st);
-    const bool tag_b =
-        resolve_operand(b, trans_b != 0, b_ld, b_bstride, b_st);
+    const bool tag_a = resolve_operand(a, trans_a, a_ld, a_bstride, a_st);
+    const bool tag_b = resolve_operand(b, trans_b, b_ld, b_bstride, b_st);
     // GEMM dims from the user flags; storage layout never swaps them.
     const int64_t m = trans_a ? a.size(-1) : a.size(-2);
     const int64_t k = trans_a ? a.size(-2) : a.size(-1);
@@ -206,53 +261,47 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
             ? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16))
             : torch::empty({m, n}, a.options().dtype(torch::kBFloat16));
     FP8Params p;
-    pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
-              m, n, k, a_ld, b_ld);
+    p.a_ptr = a_st.data_ptr();
+    p.b_ptr = b_st.data_ptr();
+    p.out_ptr = output.data_ptr();
+    p.scale = scale.data_ptr<float>();
+    p.m = static_cast<int>(m);
+    p.n = static_cast<int>(n);
+    p.k = static_cast<int>(k);
+    p.a_ld = static_cast<int>(a_ld);
+    p.b_ld = static_cast<int>(b_ld);
     // Fused epilogue bias (bf16, broadcast over rows and batches). An
     // undefined or 0-element tensor keeps the plain scaled output.
-    if (bias.defined() && bias.numel() > 0) {
-        TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kBFloat16,
+    if (bias_t.defined() && bias_t.numel() > 0) {
+        TORCH_CHECK(bias_t.is_cuda() && bias_t.scalar_type() == torch::kBFloat16,
                     "fp8 gemm bias must be a CUDA bf16 tensor");
-        TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n,
+        TORCH_CHECK(bias_t.dim() == 1 && bias_t.size(0) == n,
                     "fp8 gemm bias must be 1D of length n=", n);
-        TORCH_CHECK(bias.is_contiguous(), "fp8 gemm bias must be contiguous");
-        p.bias_ptr = bias.data_ptr();
+        TORCH_CHECK(bias_t.is_contiguous(), "fp8 gemm bias must be contiguous");
+        p.bias_ptr = bias_t.data_ptr();
     }
     p.batch = static_cast<int>(batch);
     p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride;
     p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride;
     p.out_batch_stride = m * n;
     if (a.scalar_type() == torch::kFloat8_e4m3fn)
-        dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
+        gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
     else
-        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
+        gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
     C10_CUDA_CHECK(cudaGetLastError());
     return output;
 }
 
-// mm_fp8 binding: Python None and an omitted argument both mean "no bias"
-// (resolved to an undefined tensor here, so every Python layer can pass its
-// bias argument through untouched instead of normalizing it host-side).
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
-          py::arg("fmt"));
-    m.def(
-        "mm_fp8",
-        [](torch::Tensor a, torch::Tensor b, torch::Tensor scale,
-           int64_t trans_a, int64_t trans_b, py::object bias) {
-            torch::Tensor t;
-            if (!bias.is_none()) {
-                // (py::isinstance<torch::Tensor> is false for real tensors
-                // here — torch's caster registers no pybind type info — so
-                // validate by attempting the cast itself.)
-                try {
-                    t = bias.cast<torch::Tensor>();
-                } catch (const py::cast_error&) {
-                    TORCH_CHECK(false, "bias must be a torch.Tensor or None");
-                }
-            }
-            return mm_fp8(a, b, scale, trans_a, trans_b, t);
-        },
-        py::arg("a"), py::arg("b"), py::arg("scale"), py::arg("trans_a") = 0,
-        py::arg("trans_b") = 0, py::arg("bias") = py::none());
+          py::arg("fmt"), py::arg("transposed") = false,
+          py::arg("ring") = py::none(), py::arg("hist_idx") = 0,
+          py::arg("fp8_max") = 448.0, py::arg("pow2_margin") = 1.0);
+    m.def("quantize_dual", &quantize_dual, py::arg("x"), py::arg("scale"),
+          py::arg("fmt"), py::arg("ring") = py::none(),
+          py::arg("hist_idx") = 0, py::arg("fp8_max") = 448.0,
+          py::arg("pow2_margin") = 1.0);
+    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
+          py::arg("trans_a") = false, py::arg("trans_b") = false,
+          py::arg("bias") = py::none());
 }

csrc/kernels/fp8/quantize.cuh

@@ -1,10 +1,7 @@
 #pragma once
-// FP8 quantize device code — pure CUDA, no torch. Any float input element
-// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over
-// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the
-// FP8QuantizeParams POD, formats and input types ride on template parameters,
-// and the launcher is a plain function usable from both the torch binding and
-// pure C tests.
+// FP8 quantize device code — pure CUDA, no torch: kernels take the
+// FP8QuantizeParams POD, format and input type ride on template parameters,
+// and the launcher is shared by the torch binding and the C tests.
 
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
@@ -13,13 +10,13 @@
 #include <cstdint>
 
 #include "common.h"
-#include "../common/reduce.cuh"
+#include "common/reduce.cuh"
 
 namespace astrai {
 namespace fp8 {
 
-// Input element type traits: one element -> float, and the vectorized
-// unpack of one 16-byte load into kVecElems floats.
+// Input element type traits: one element -> float, the unpack of one
+// 16-byte load into kVecElems floats, and a native 2-element pair load.
 template <typename InT>
 struct quant_in_traits;
 
@@ -31,14 +28,22 @@ struct quant_in_traits<__nv_bfloat16> {
     }
     static __device__ __forceinline__ void load_vec(const uint4& raw,
                                                     float* f) {
-        const unsigned w[4] = {raw.x, raw.y, raw.z, raw.w};
+        const __nv_bfloat162* b2 =
+            reinterpret_cast<const __nv_bfloat162*>(&raw);
 #pragma unroll
         for (int j = 0; j < 4; ++j) {
-            f[2 * j] =
-                __bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu));
-            f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16));
+            const float2 p = __bfloat1622float2(b2[j]);
+            f[2 * j] = p.x;
+            f[2 * j + 1] = p.y;
         }
     }
+    static __device__ __forceinline__ void load_pair(const __nv_bfloat16* p,
+                                                     float* f) {
+        const float2 v = __bfloat1622float2(
+            *reinterpret_cast<const __nv_bfloat162*>(p));
+        f[0] = v.x;
+        f[1] = v.y;
+    }
 };
 
 template <>
@@ -57,6 +62,13 @@ struct quant_in_traits<__half> {
             f[2 * j + 1] = p.y;
         }
     }
+    static __device__ __forceinline__ void load_pair(const __half* p,
+                                                     float* f) {
+        const float2 v =
+            __half22float2(*reinterpret_cast<const __half2*>(p));
+        f[0] = v.x;
+        f[1] = v.y;
+    }
 };
 
 template <>
@@ -65,15 +77,27 @@ struct quant_in_traits<float> {
     static __device__ __forceinline__ float to_float(float v) { return v; }
     static __device__ __forceinline__ void load_vec(const uint4& raw,
                                                     float* f) {
-        f[0] = __uint_as_float(raw.x);
-        f[1] = __uint_as_float(raw.y);
-        f[2] = __uint_as_float(raw.z);
-        f[3] = __uint_as_float(raw.w);
+        const unsigned* w = reinterpret_cast<const unsigned*>(&raw);
+#pragma unroll
+        for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]);
+    }
+    static __device__ __forceinline__ void load_pair(const float* p,
+                                                     float* f) {
+        f[0] = p[0];
+        f[1] = p[1];
     }
 };
 
-// Convert one float pair to one packed fp8 pair. The stored bytes see
-// value * mult (round-nearest-even + satfinite).
+// One float -> one fp8 byte (round-nearest-even + satfinite).
+template <FP8Format Fmt>
+__device__ __forceinline__ uint8_t cvt_fp8(float v) {
+    if constexpr (Fmt == FP8Format::E5M2)
+        return __nv_fp8_e5m2(v).__x;
+    else
+        return __nv_fp8_e4m3(v).__x;
+}
+
+// One float pair -> one packed fp8x2 word (round-nearest-even + satfinite).
 template <FP8Format Fmt>
 __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
     constexpr __nv_fp8_interpretation_t kFmt =
@@ -82,23 +106,52 @@ __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
         make_float2(a, b), __NV_SATFINITE, kFmt));
 }
 
-// Quantize kernel: float input -> FP8 (E4M3 or E5M2), fused amax over raw
-// values.
+// Block-wide amax reduce -> one atomic per block: warp-reduce, park one
+// value per warp, thread 0 folds. kWarps must cover the block's warp count.
+// With p.fold_ring, the last-finishing block additionally folds the final
+// amax into the history window and publishes the next scale (atomicAdd
+// ticket + fences), re-zeroing the amax slot and the counter for the next
+// launch — the host-side delayed-scaling update chain disappears.
+template <int kWarps>
+__device__ __forceinline__ void publish_amax(const FP8QuantizeParams& p,
+                                             float v) {
+    v = warp_reduce_max(v);
+    __shared__ float slots[kWarps];
+    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
+    if ((tid & 31) == 0) slots[tid >> 5] = v;
+    __syncthreads();
+    if (tid == 0) {
+#pragma unroll
+        for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
+        atomic_max_float(p.amax, v);
+        if (!p.fold_ring) return;
+        __threadfence();
+        const unsigned int ticket = atomicAdd(p.done, 1u);
+        __threadfence();
+        if (ticket != gridDim.x - 1u) return;
+        p.hist[p.hist_idx] = *p.amax;
+        float peak = p.hist[0];
+        for (int i = 1; i < p.hist_len; ++i) peak = fmaxf(peak, p.hist[i]);
+        *p.scale_out = fmaxf(peak / p.fp8_max / p.pow2_margin, 1e-12f);
+        *p.amax = 0.0f;
+        *p.done = 0u;
+    }
+}
+
+// Elementwise quantize kernel (QuantLayout::RowMajor): vectorized 16B loads
+// -> fp8 stores, fused amax over raw values.
 template <FP8Format Fmt, typename InT>
 __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
     const float mult = *p.scale;
     const auto* x = static_cast<const InT*>(p.input_ptr);
-    void* x8 = p.output_ptr;
-    float* amax = p.amax;
+    uint8_t* x8 = static_cast<uint8_t*>(p.output_ptr);
     float local_amax = 0.0f;
     const int64_t stride = (int64_t)blockDim.x * gridDim.x;
 
-    // Vectorized body: one 16B load -> kVecElems fp8 bytes per step (8
-    // elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B
-    // aligned and the binding passes freshly allocated contiguous buffers,
-    // so element 0 keeps the uint4 access natural; a misaligned base
-    // (contiguous view with an odd storage offset) falls back to the scalar
-    // loop below via total_vec = 0.
+    // One 16B load -> kVecElems bytes per step. Torch allocations are >=16B
+    // aligned, so element 0 keeps the uint4 access natural; a misaligned
+    // base (odd storage offset view) falls to the scalar tail via
+    // total_vec = 0.
     constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
     const bool aligned =
         ((reinterpret_cast<uintptr_t>(x) |
@@ -125,8 +178,7 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
             packed[j] = (lo & 0xffffu) | (hi << 16);
         }
         if constexpr (kVecElems == 8)
-            reinterpret_cast<uint2*>(x8)[i] =
-                make_uint2(packed[0], packed[1]);
+            reinterpret_cast<uint2*>(x8)[i] = make_uint2(packed[0], packed[1]);
         else
             reinterpret_cast<unsigned*>(x8)[i] = packed[0];
     }
@@ -136,38 +188,124 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
          i < p.total; i += stride) {
         const float v = quant_in_traits<InT>::to_float(x[i]);
         local_amax = fmaxf(local_amax, fabsf(v));
-        if constexpr (Fmt == FP8Format::E5M2) {
-            reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] =
-                __nv_fp8_e5m2(v * mult);
-        } else {
-            reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] =
-                __nv_fp8_e4m3(v * mult);
-        }
-    }
-    if (amax) {
-        local_amax = warp_reduce_max(local_amax);
-        __shared__ float slots[32];
-        if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax;
-        __syncthreads();
-        if (threadIdx.x == 0) {
-            float v = 0.0f;
-            for (int w = 0; w < (blockDim.x >> 5); ++w)
-                v = fmaxf(v, slots[w]);
-            atomic_max_float(amax, v);
-        }
+        x8[i] = cvt_fp8<Fmt>(v * mult);
     }
+    if (p.amax) publish_amax<8>(p, local_amax);
 }
 
+// Tiled transpose quantize (QuantLayout::Transposed/Dual): reads the
+// [rows][cols] input
+// once and writes the fp8 bytes transposed ([cols][rows], so the contract
+// dim lands K-contiguous for NT GEMM operands) and, in mode 2, the row-major
+// copy too. 64x32 tiles, one native pair load per row (a full 128B warp
+// read); rows whose pair is unaligned or ragged (odd widths, misaligned
+// bases) fall back to element loads in place. Staging goes through a byte
+// tile whose pitch keeps the store stride coprime with the 32 banks.
+// (+25-35% over the former 32x32 scalar kernel on sub-4M tensors; ~5%
+// slower once DRAM-saturated — accepted for the single-kernel shape.)
 template <FP8Format Fmt, typename InT>
+__global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
+    constexpr int kTileC = 64, kTileR = 32;
+    // 34B pitch: staging stride is 17 words (coprime with the 32 banks) so
+    // the pair-byte stores stay conflict-free, and the byte-wise consume
+    // reads still span distinct words.
+    __shared__ uint8_t tile[kTileC][kTileR + 2];
+    const float mult = *p.scale;
+    const auto* x = static_cast<const InT*>(p.input_ptr);
+    const int r0 = blockIdx.y * kTileR;
+    const int c0 = blockIdx.x * kTileC;
+    const int r = r0 + threadIdx.y * 4;
+    const int c = c0 + threadIdx.x * 2;  // cols even => the pair is in-bounds
+
+    uint8_t q[4][2];
+    float local_amax = 0.0f;
+    // Vectorize the pair when both elements are in-bounds and the native
+    // 2-element load is aligned; odd widths, misaligned bases and ragged
+    // edges fall back to element loads row by row.
+    constexpr int kPairAlign = 2 * (int)sizeof(InT);
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        q[j][0] = 0;
+        q[j][1] = 0;
+        if (r + j < p.rows && c < p.cols) {
+            const InT* a = x + (int64_t)(r + j) * p.cols + c;
+            if (c + 1 < p.cols &&
+                (reinterpret_cast<uintptr_t>(a) & (kPairAlign - 1)) == 0) {
+                float f[2];
+                quant_in_traits<InT>::load_pair(a, f);
+#pragma unroll
+                for (int k = 0; k < 2; ++k) {
+                    local_amax = fmaxf(local_amax, fabsf(f[k]));
+                    q[j][k] = cvt_fp8<Fmt>(f[k] * mult);
+                }
+            } else {
+                const float v0 = quant_in_traits<InT>::to_float(a[0]);
+                local_amax = fmaxf(local_amax, fabsf(v0));
+                q[j][0] = cvt_fp8<Fmt>(v0 * mult);
+                if (c + 1 < p.cols) {
+                    const float v1 = quant_in_traits<InT>::to_float(a[1]);
+                    local_amax = fmaxf(local_amax, fabsf(v1));
+                    q[j][1] = cvt_fp8<Fmt>(v1 * mult);
+                }
+            }
+        }
+    }
+    if (p.out_layout == QuantLayout::Dual) {
+        uint8_t* out = static_cast<uint8_t*>(p.output_ptr);
+#pragma unroll
+        for (int j = 0; j < 4; ++j)
+            if (r + j < p.rows && c < p.cols) {
+                uint8_t* o = out + (int64_t)(r + j) * p.cols + c;
+                const int64_t off = (int64_t)(r + j) * p.cols + c;
+                if (c + 1 < p.cols && (off & 1) == 0)
+                    *reinterpret_cast<unsigned short*>(o) =
+                        (unsigned short)(q[j][0] | (q[j][1] << 8));
+                else {
+                    o[0] = q[j][0];
+                    if (c + 1 < p.cols) o[1] = q[j][1];
+                }
+            }
+    }
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+#pragma unroll
+        for (int k = 0; k < 2; ++k)
+            tile[threadIdx.x * 2 + k][threadIdx.y * 4 + j] = q[j][k];
+    __syncthreads();
+    // Transposed scatter: output element (c, r) lives at c * rows + r;
+    // threadIdx.x tracks r so each warp writes one contiguous run. tile is
+    // [col][row]; warp y walks 8 columns, threads read down one column.
+    uint8_t* out_t = static_cast<uint8_t*>(p.output_transposed_ptr);
+#pragma unroll
+    for (int i = 0; i < 8; ++i) {
+        const int oc = c0 + threadIdx.y * 8 + i;
+        if (oc < p.cols && r0 + threadIdx.x < p.rows)
+            out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
+                tile[threadIdx.y * 8 + i][threadIdx.x];
+    }
+    if (p.amax) publish_amax<8>(p, local_amax);
+}
+
+// Unified quantize launcher: Tiled selects the transpose kernel
+// (QuantLayout::Transposed/Dual) over the vectorized elementwise one. The
+// transpose kernel vectorizes
+// pair loads in-kernel and falls back to scalar loads at unaligned/ragged
+// rows, so the host side picks only the grid.
+template <FP8Format Fmt, typename InT, bool Tiled = false>
 void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
+    if constexpr (Tiled) {
+        const dim3 grid((p.cols + 63) / 64, (p.rows + 31) / 32);
+        if (grid.x == 0 || grid.y == 0) return;
+        fp8_quantize_tiled_kernel<Fmt, InT><<<grid, dim3(32, 8), 0, stream>>>(p);
+    } else {
         constexpr int kThreads = 256;
-    // One block per 256 vectors; at least one block so the scalar tail of a
-    // tiny / misaligned tensor is still covered.
         constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
-    int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
-    if (blocks < 1) blocks = 1;
+        // Grid-stride loops: any grid >= 1 is correct; one block per 256
+        // vectors plus the tail block covers tiny and misaligned tensors.
+        const int64_t blocks = 1 + p.total / (kVecElems * kThreads);
         fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
     }
+}
 
 }  // namespace fp8
 }  // namespace astrai

csrc/kernels/rotary_emb.cu

@@ -11,34 +11,41 @@ __global__ void rotary_emb_kernel(
     int n_heads,
     int head_dim
 ) {
-    const int half_dim = head_dim >> 1;
-    const int total = n_tokens * n_heads * half_dim;
+    // Each head tiles into exact 2-pair chunks: one 8B x access and one 16B
+    // cos/sin access per chunk (head_dim % 4 == 0 is enforced on the host).
+    const int chunks = head_dim >> 2;
+    const int total = n_tokens * n_heads * chunks;
 
-    for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
-         idx < total;
-         idx += gridDim.x * blockDim.x) {
+    for (int c = blockIdx.x * blockDim.x + threadIdx.x;
+         c < total;
+         c += gridDim.x * blockDim.x) {
+        const int chunk = c % chunks;
+        const int tmp = c / chunks;
+        const int head = tmp % n_heads;
+        const int token = tmp / n_heads;
 
-        int pair = idx % half_dim;
-        int tmp = idx / half_dim;
-        int head = tmp % n_heads;
-        tmp /= n_heads;
-        int token = tmp;
+        const int x_off = (tmp * head_dim) + (chunk << 2);
+        const int f_off = ((token * chunks) + chunk) << 2;
 
-        int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
-        int cs_offset = (token * half_dim + pair) * 2;
+        const float4 f = *reinterpret_cast<const float4*>(freqs_cis + f_off);
+        const uint2 xr = *reinterpret_cast<const uint2*>(x + x_off);
+        __nv_bfloat162 p0 = *reinterpret_cast<const __nv_bfloat162*>(&xr.x);
+        __nv_bfloat162 p1 = *reinterpret_cast<const __nv_bfloat162*>(&xr.y);
 
-        __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
-        float x_even = __bfloat162float(__low2bfloat16(x_pair));
-        float x_odd  = __bfloat162float(__high2bfloat16(x_pair));
+        const float e0 = __bfloat162float(__low2bfloat16(p0));
+        const float o0 = __bfloat162float(__high2bfloat16(p0));
+        const float e1 = __bfloat162float(__low2bfloat16(p1));
+        const float o1 = __bfloat162float(__high2bfloat16(p1));
 
-        float c = freqs_cis[cs_offset];
-        float s = freqs_cis[cs_offset + 1];
+        __nv_bfloat162 r0 = __floats2bfloat162_rn(
+            e0 * f.x - o0 * f.y, e0 * f.y + o0 * f.x);
+        __nv_bfloat162 r1 = __floats2bfloat162_rn(
+            e1 * f.z - o1 * f.w, e1 * f.w + o1 * f.z);
 
-        float out_even = x_even * c - x_odd * s;
-        float out_odd  = x_even * s + x_odd * c;
-
-        __nv_bfloat162 out_pair = __floats2bfloat162_rn(out_even, out_odd);
-        *reinterpret_cast<__nv_bfloat162*>(out + x_offset) = out_pair;
+        uint2 oraw;
+        *reinterpret_cast<__nv_bfloat162*>(&oraw.x) = r0;
+        *reinterpret_cast<__nv_bfloat162*>(&oraw.y) = r1;
+        *reinterpret_cast<uint2*>(out + x_off) = oraw;
     }
 }
 
@@ -64,7 +71,7 @@ torch::Tensor rotary_emb(
     int n_heads = x.size(x.dim() - 2);
     int head_dim = x.size(x.dim() - 1);
 
-    TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
+    TORCH_CHECK(head_dim % 4 == 0, "head_dim must be a multiple of 4");
     TORCH_CHECK(freqs_cis.numel() == (int64_t)n_tokens * head_dim,
                 "freqs_cis token or rotary dimension mismatch");
     TORCH_CHECK(freqs_cis.size(-2) == head_dim / 2, "freqs_cis dim/2 mismatch");
@@ -72,10 +79,9 @@ torch::Tensor rotary_emb(
 
     auto out = torch::empty_like(x);
 
-    int half_dim = head_dim / 2;
-    int total = n_tokens * n_heads * half_dim;
+    int work = n_tokens * n_heads * (head_dim / 4);
     int block = 256;
-    int grid = std::min((total + block - 1) / block, 1024);
+    int grid = std::min((work + block - 1) / block, 2048);
 
     rotary_emb_kernel<<<grid, block, 0, stream>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),

csrc/tests/attn_paged_test.cu

@@ -1,5 +1,5 @@
 // Compile:
-//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
+//   nvcc -I csrc/kernels -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
 //        --extra-device-vectorization -Xcompiler -fopenmp \
 //        csrc/tests/attn_paged_test.cu \
 //        -o /tmp/test_paged && /tmp/test_paged
@@ -7,7 +7,7 @@
 #include <cstring>
 #include <vector>
 #include "test_utils.cuh"
-#include "../kernels/attention/dispatchers.cuh"
+#include "attention/dispatchers.cuh"
 
 using namespace astrai::attention;
 

csrc/tests/attn_test.cu

@@ -1,13 +1,13 @@
 /*
 Pure-C test — uses shared dispatcher.  Combines the decode (split-KV) and
 prefill (split-Q) correctness checks + benchmarks into one binary.
-nvcc -I csrc -arch=sm_89 -O3 \
+nvcc -I csrc/kernels -arch=sm_89 -O3 \
     --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
     -Xcompiler -fopenmp csrc/tests/attn_test.cu -o test && ./test
 */
 
 #include "test_utils.cuh"
-#include "../kernels/attention/dispatchers.cuh"
+#include "attention/dispatchers.cuh"
 
 using namespace astrai::attention;
 

csrc/tests/fp8_test.cu

@@ -6,7 +6,7 @@ Part 1 exercises one bf16 -> fp8 -> mma.sync m16n8k32 instruction pair
 Part 2 checks launch_fp8_gemm across all four operand layouts, both K
 tiles, and ragged shapes against an fp32 CPU reference.
 
-nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
+nvcc -I csrc/kernels -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
     && /tmp/fp8_test
 */
 
@@ -21,8 +21,8 @@ nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test
 #include <type_traits>
 #include <vector>
 
-#include "../kernels/common/mma.cuh"
-#include "../kernels/fp8/gemm.cuh"
+#include "common/mma.cuh"
+#include "fp8/gemm.cuh"
 
 using namespace astrai::fp8;
 
@@ -170,9 +170,37 @@ static bool test_single_mma() {
 // Part 2: GEMM correctness — layouts x K-tiles vs fp32 CPU reference
 // ---------------------------------------------------------------------------
 
+// Naive fp32 reference on the GPU: same layout interpretation as the CPU
+// loop it replaces (O(m*n) to check instead of O(m*n*k) to compute).
+__global__ static void
+naive_gemm_ref(const __nv_fp8_e4m3* a, const __nv_fp8_e4m3* b, float* out,
+               int m, int n, int k, int a_ld, int b_ld, int a_rm, int b_rm) {
+    const int i = blockIdx.y * 32 + threadIdx.y;
+    const int j = blockIdx.x * 32 + threadIdx.x;
+    if (i >= m || j >= n) return;
+    float acc = 0.f;
+    for (int kk = 0; kk < k; ++kk) {
+        float av = a_rm ? (float)a[i * a_ld + kk] : (float)a[kk * a_ld + i];
+        float bv = b_rm ? (float)b[kk * b_ld + j] : (float)b[j * b_ld + kk];
+        acc += av * bv;
+    }
+    out[i * n + j] = acc;
+}
+
+// Big-CTA policies for the direct-layout cases: kK/Stages vary per case;
+// the fast interior loop follows the dual-congruous rule, grouped raster 8
+// matches the production dispatch.
+template <typename LA, typename LB>
+constexpr bool kCaseFast =
+    !std::is_same_v<LA, ColMajor> && !std::is_same_v<LB, RowMajor>;
+template <typename LA, typename LB, int kK, int Stages>
+using CasePolicy =
+    Fp8GemmPolicy<FP8Format::E4M3, 128, 128, LA, LB, 64, 32, kK, Stages, 8,
+                  false, kCaseFast<LA, LB>>;
+
 template <typename LA, typename LB, int kK, int Stages>
 static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
-                          int k, int a_ld, int b_ld) {
+                          int k, int a_ld, int b_ld, int dispatch = 0) {
     __nv_fp8_e4m3 *da, *db;
     __nv_bfloat16* dout;
     float* dscale;
@@ -205,7 +233,27 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     p.k = k;
     p.a_ld = a_ld;
     p.b_ld = b_ld;
-    launch_fp8_gemm<FP8Format::E4M3, LA, LB, kK, Stages>(p, 0);
+    float* d_ref;
+    cudaMalloc(&d_ref, (size_t)m * n * 4);
+    naive_gemm_ref<<<dim3((n + 31) / 32, (m + 31) / 32), dim3(32, 32)>>>(
+        da, db, d_ref, m, n, k, a_ld, b_ld,
+        !std::is_same_v<LA, ColMajor>, !std::is_same_v<LB, ColMajor>);
+    std::vector<float> href((size_t)m * n);
+    cudaMemcpy(href.data(), d_ref, href.size() * 4, cudaMemcpyDeviceToHost);
+    cudaFree(d_ref);
+
+    if (dispatch == 1)
+        // Production route, NN: the dual-N-contiguous problem has no
+        // dedicated instantiation — canonicalize_gemm swaps to the
+        // transposed <ColMajor, ColMajor> kernel with its out-transposed
+        // epilogue (see gemm.cuh).
+        gemm<FP8Format::E4M3>(p, 0, false, false);
+    else if (dispatch == 2)
+        // Production route, NT: exercises plan_gemm's small/narrow/big
+        // selection for this shape.
+        gemm<FP8Format::E4M3>(p, 0, false, true);
+    else
+        launch_policy<CasePolicy<LA, LB, kK, Stages>>(p, 0);
     cudaError_t e = cudaDeviceSynchronize();
     if (e != cudaSuccess) {
         printf("  CUDA err: %s\n", cudaGetErrorString(e));
@@ -218,20 +266,7 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     bool ok = true;
     for (int i = 0; i < m && ok; ++i) {
         for (int j = 0; j < n && ok; ++j) {
-            float ref = 0;
-            for (int kk = 0; kk < k; ++kk) {
-                // A reference reads the actual uploaded buffer: LA ColMajor
-                // means the buffer is [K][M] (ha_t), else [M][K].
-                float av = std::is_same_v<LA, ColMajor>
-                               ? (float)__nv_fp8_e4m3(ha[kk * m + i])
-                               : (float)__nv_fp8_e4m3(ha[i * k + kk]);
-                float bv;
-                if (std::is_same_v<LB, ColMajor>)
-                    bv = (float)__nv_fp8_e4m3(hb[j * k + kk]);
-                else
-                    bv = (float)__nv_fp8_e4m3(hb[kk * n + j]);
-                ref += av * bv;
-            }
+            const float ref = href[(size_t)i * n + j];
             float got =
                 __bfloat162float(__ushort_as_bfloat16(hb16[i * n + j]));
             float err = fabsf(got - ref);
@@ -254,6 +289,7 @@ static bool test_gemm() {
     } cfgs[] = {
         {128, 128, 128}, {256, 128, 256}, {128, 256, 64},
         {100, 130, 96},  {64, 64, 160},   {300, 200, 320},
+        {2048, 256, 512}, {1024, 1024, 512},
     };
     bool all = true;
     for (auto& c : cfgs) {
@@ -274,12 +310,12 @@ static bool test_gemm() {
         printf(" NT K64:");
         all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(ha, hb_colmajor, c.m,
                                                         c.n, c.k, c.k, c.k);
-        printf(" NN K32:");
-        all &= run_gemm_case<RowMajor, RowMajor, 32, 3>(ha, hb_rowmajor, c.m,
-                                                        c.n, c.k, c.k, c.n);
-        printf(" NN K64:");
-        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(ha, hb_rowmajor, c.m,
-                                                        c.n, c.k, c.k, c.n);
+        printf(" NN swap:");
+        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(
+            ha, hb_rowmajor, c.m, c.n, c.k, c.k, c.n, /*dispatch=*/1);
+        printf(" NT disp:");
+        all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(
+            ha, hb_colmajor, c.m, c.n, c.k, c.k, c.k, /*dispatch=*/2);
         printf(" TN K32:");
         all &= run_gemm_case<ColMajor, ColMajor, 32, 3>(ha_t, hb_colmajor, c.m,
                                                         c.n, c.k, c.m, c.k);

docker-compose.yml

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

docs/developer/architecture.md

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

docs/developer/cuda_kernels.md

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

docs/developer/docker-serving.md

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

docs/developer/docker-training.md

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

docs/developer/internals.md

@@ -118,12 +118,13 @@ on_train_begin
         on_batch_end
 
         if executor.sync_gradients:
-          on_optimizer_step
+          before_optimizer_step
           optimizer.step()
           strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
+          after_optimizer_step
     on_epoch_end
 on_train_end
 ```
@@ -139,8 +140,9 @@ Strategy metrics are detached and converted to Python `float` values before the
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback` |
+| `before_optimizer_step` | Every accumulation window, before `optimizer.step()` | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
+| `on_batch_end` | Every batch | — |
+| `after_optimizer_step` | Every accumulation window, after `optimizer.step()` and `scheduler.step()` | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
 | `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
@@ -185,7 +187,7 @@ The extension package separates mechanism from policy:
 Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/backend/attention.py`):
 
 - **`CudaBackend`** (default when supported): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool).
-- **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
+- **`FlashAttnBackend`**: optional flash-attn dispatch; inference paths gather flat K/V from the pool via `req_to_token` and call `flash_attn_varlen_func` over the ragged batch (fp16/bf16 only); dense mask-free training calls use `flash_attn_func`.
 - **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
 - The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call.
 - Resolution precedence is: explicit `attn_backend(...)` context > `ASTR_BACKEND` env > default. An explicit `attn_backend(...)` selection is strict (incompatible calls raise); `ASTR_BACKEND` is a default-level override that falls back to a compatible backend when incapable. Training calls (`fwd=None`, no KV cache) resolve by capability: the CUDA cache kernels cannot run without a cache, so they fall back to flash (mask-free/causal calls only) and finally to torch SDPA.

docs/get-started.md

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

docs/guides/training.md

@@ -70,12 +70,13 @@ on_train_begin
         on_batch_end
 
         if executor.sync_gradients:
-          on_optimizer_step
+          before_optimizer_step
           optimizer.step()
           strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
+          after_optimizer_step
     on_epoch_end
 on_train_end
 ```
@@ -87,8 +88,9 @@ on_train_end
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback` |
+| `before_optimizer_step` | Every accumulation window, before `optimizer.step()` | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
+| `on_batch_end` | Every batch | — |
+| `after_optimizer_step` | Every accumulation window, after `optimizer.step()` and `scheduler.step()` | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
 | `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |

scripts/docker/serve_runtime.py

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

scripts/docker/train_runtime.py

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

scripts/serve.sh

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

scripts/train.sh

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

setup.py

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

tests/extension/test_dispatch.py

@@ -0,0 +1,283 @@
+"""Tests for the generic operator dispatcher (Spec / decision tables)."""
+
+import importlib
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import astrai.extension.dispatch as dispatch
+from astrai.extension import (
+    ATTN_BACKEND,
+    ExplicitSelectionError,
+    ImplRecord,
+    Spec,
+    axis,
+    explain,
+    op_backend,
+    resolve,
+    resolve_plan,
+)
+from astrai.extension.backend import apply_rotary_emb
+
+attn_mod = importlib.import_module("astrai.extension.backend.attention")
+rotary_mod = importlib.import_module("astrai.extension.backend.rotary")
+
+
+@pytest.fixture
+def toy_family():
+    """A toy family: alpha (restricted), beta, and an unfaithful fast row."""
+    calls = []
+
+    def records():
+        return [
+            ImplRecord(
+                "toy",
+                "alpha",
+                "alpha-obj",
+                axis("dtype").in_(torch.bfloat16),
+                priority=0,
+            ),
+            ImplRecord(
+                "toy",
+                "beta",
+                "beta-obj",
+                Spec.always(),
+                priority=10,
+            ),
+            ImplRecord(
+                "toy",
+                "fp8",
+                "fp8-obj",
+                Spec.always(),
+                priority=1,
+                faithful=False,
+            ),
+        ]
+
+    dispatch.register_family(
+        "toy",
+        lambda **kw: kw,
+        records,
+        lambda: ImplRecord("toy", "beta", "beta-obj", Spec.always()),
+    )
+    yield calls
+    dispatch._FAMILIES.pop("toy", None)
+
+
+def test_spec_composition_and_description():
+    spec = axis("dtype").in_(torch.bfloat16) & axis("grad_enabled").eq(False)
+    assert spec.matches({"dtype": torch.bfloat16, "grad_enabled": False})
+    assert not spec.matches({"dtype": torch.bfloat16, "grad_enabled": True})
+    assert "dtype" in spec.description and "grad_enabled" in spec.description
+
+    either = axis("fwd").none() | axis("has_cache").truthy()
+    assert either.matches({"fwd": None})
+    assert either.matches({"fwd": "decode", "has_cache": True})
+    assert not either.matches({"fwd": "decode"})
+
+    assert (~axis("fwd").none()).matches({"fwd": "decode"})
+
+
+def test_chain_returns_first_capable(toy_family):
+    assert resolve("toy", dtype=torch.bfloat16).record.obj == "alpha-obj"
+    assert resolve("toy", dtype=torch.float32).record.obj == "beta-obj"
+
+
+def test_unfaithful_rows_are_chain_invisible(toy_family):
+    resolution = resolve("toy", dtype=torch.float32)
+    assert resolution.record.obj == "beta-obj"
+    with op_backend(toy="fp8"):
+        assert resolve("toy", dtype=torch.float32).record.obj == "fp8-obj"
+
+
+def test_explicit_selection_is_strict(toy_family):
+    with pytest.raises(ExplicitSelectionError):
+        resolve("toy", dtype=torch.float32, explicit="alpha")
+    assert resolve("toy", dtype=torch.bfloat16, explicit="alpha").origin == "explicit"
+
+
+def test_context_selection_is_strict(toy_family):
+    with op_backend(toy="alpha"):
+        with pytest.raises(ExplicitSelectionError):
+            resolve("toy", dtype=torch.float32)
+        assert resolve("toy", dtype=torch.bfloat16).origin == "context"
+
+
+def test_unknown_explicit_name_raises(toy_family):
+    with pytest.raises(ValueError, match="Unknown toy implementation"):
+        resolve("toy", explicit="nope")
+    with pytest.raises(ValueError, match="Unknown toy implementation"):
+        with op_backend(toy="nope"):
+            pass
+
+
+class _Probe:
+    def __init__(self, capable):
+        self.capable = capable
+        self.probed = 0
+
+    def supports_call(self, *args, **kwargs):
+        self.probed += 1
+        return self.capable
+
+
+def test_adhoc_instance_probed_via_supports_call(toy_family):
+    probe = _Probe(capable=True)
+    resolution = resolve("toy", dtype=torch.float32, explicit=probe)
+    assert resolution.record.obj is probe and probe.probed == 1
+
+    incapable = _Probe(capable=False)
+    with pytest.raises(ExplicitSelectionError):
+        resolve("toy", dtype=torch.float32, explicit=incapable)
+
+
+def test_nested_op_backend_scopes(toy_family):
+    with op_backend(toy="alpha"):
+        with op_backend(toy="beta"):
+            assert resolve("toy", dtype=torch.float32).origin == "context"
+        assert resolve("toy", dtype=torch.bfloat16).origin == "context"
+
+
+def test_env_entry_is_soft(toy_family, monkeypatch):
+    monkeypatch.setenv("ASTR_OPS", "toy=alpha")
+    resolution = resolve("toy", dtype=torch.float32)
+    assert resolution.record.obj == "beta-obj" and resolution.origin == "chain"
+    assert resolve("toy", dtype=torch.bfloat16).origin == "env"
+
+
+def test_env_unknown_impl_ignored(toy_family, monkeypatch):
+    monkeypatch.setenv("ASTR_OPS", "toy=missing")
+    assert resolve("toy", dtype=torch.float32).record.obj == "beta-obj"
+
+
+def test_env_profile_reference(toy_family, monkeypatch):
+    monkeypatch.setenv("ASTR_OPS", "profile=reference")
+    resolution = resolve("toy", dtype=torch.bfloat16)
+    assert resolution.origin == "profile"
+    monkeypatch.setenv("ASTR_OPS", "toy=alpha,profile=reference")
+    assert resolve("toy", dtype=torch.bfloat16).origin == "env"
+
+
+def test_legacy_env_alias(monkeypatch):
+    monkeypatch.setenv("ASTR_BACKEND", "torch_native")
+    assert dispatch.env_selection("attention") == "torch_native"
+    monkeypatch.setenv("ASTR_OPS", "attention=cuda")
+    assert dispatch.env_selection("attention") == "cuda"
+
+
+def test_context_beats_env(toy_family, monkeypatch):
+    monkeypatch.setenv("ASTR_OPS", "toy=alpha")
+    with op_backend(toy="beta"):
+        assert resolve("toy", dtype=torch.float32).origin == "context"
+
+
+def test_resolve_plan_snapshots_families(toy_family):
+    plan = resolve_plan(
+        {
+            "toy": ((), {"dtype": torch.bfloat16}),
+            "rotary": (
+                (
+                    SimpleNamespace(dtype=torch.bfloat16, is_cuda=True),
+                    SimpleNamespace(),
+                ),
+                {},
+            ),
+        }
+    )
+    assert plan["toy"].record.obj == "alpha-obj"
+    assert plan["rotary"].record.name in ("cuda", "torch")
+
+
+def test_explain_shows_rejection_reasons(toy_family):
+    text = explain("toy", dtype=torch.float32)
+    assert "alpha: reject" in text and "beta: MATCH" in text
+    assert "=> beta" in text
+
+
+def test_explain_reports_strict_error(toy_family):
+    text = explain("toy", dtype=torch.float32, explicit="alpha")
+    assert "ERROR" in text
+
+
+_DUMMY_CACHE = object()
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("head_dim", [64, 96])
+@pytest.mark.parametrize(
+    "fwd,has_cache,ndim,has_mask",
+    [
+        ("decode", True, 3, False),
+        ("prefill", True, 3, False),
+        (None, False, 4, False),
+        (None, False, 4, True),
+    ],
+)
+def test_attention_specs_mirror_supports_call(
+    dtype, head_dim, fwd, has_cache, ndim, has_mask
+):
+    shape = {3: (1, 2, head_dim), 4: (1, 2, 4, head_dim)}[ndim]
+    q = torch.zeros(shape, dtype=dtype)
+    mask = torch.zeros(1, 1, 2, 2, dtype=torch.bool) if has_mask else None
+    cache = _DUMMY_CACHE if has_cache else None
+    ax = attn_mod._axes(q, cache, mask, False, fwd)
+
+    cuda = attn_mod._instance(attn_mod.CudaBackend)
+    assert attn_mod._SPEC_CUDA.matches(ax) == cuda.supports_call(
+        q, cache, mask, False, fwd
+    )
+
+    flash = attn_mod._instance(attn_mod.FlashAttnBackend)
+    assert attn_mod._SPEC_FLASH.matches(ax) == flash.supports_call(
+        q, cache, mask, False, fwd
+    )
+
+
+def test_attention_resolution_matches_legacy_semantics():
+    q = torch.zeros(1, 2, 4, 8, dtype=torch.float32)
+    resolution = resolve("attention", q, None, None, True, "prefill")
+    assert resolution.record.name == ATTN_BACKEND.TORCH_NATIVE.value
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="rotary CUDA path needs a GPU"
+)
+class TestRotaryDispatch:
+    def _input(self):
+        torch.manual_seed(0)
+        x = torch.randn(1, 5, 3, 16, device="cuda", dtype=torch.bfloat16)
+        freqs = torch.randn(1, 5, 8, 2, device="cuda", dtype=torch.float32)
+        return x, freqs
+
+    def test_cuda_row_selected_under_inference_mode(self):
+        from astrai.extension.loader import is_available
+
+        x, freqs = self._input()
+        with torch.inference_mode():
+            resolution = resolve("rotary", x, freqs)
+        expected = "cuda" if is_available("rotary_emb") else "torch"
+        assert resolution.record.name == expected
+
+    def test_grad_falls_back_to_torch(self):
+        x, freqs = self._input()
+        assert resolve("rotary", x, freqs).record.name == "torch"
+
+    def test_context_switch_to_torch(self):
+        x, freqs = self._input()
+        with torch.inference_mode():
+            with op_backend(rotary="torch"):
+                out = apply_rotary_emb(x, freqs)
+        assert out.shape == x.shape and out.dtype == torch.bfloat16
+
+    def test_cuda_matches_torch_numerics(self):
+        from astrai.extension.loader import is_available
+
+        if not is_available("rotary_emb"):
+            pytest.skip("rotary kernel not built")
+        x, freqs = self._input()
+        with torch.inference_mode():
+            fast = apply_rotary_emb(x, freqs)
+        slow = rotary_mod._torch_apply
+        ref = slow(x, freqs)
+        assert torch.allclose(fast.float(), ref.float(), atol=2e-2, rtol=1e-2)

tests/extension/test_fp8_mma.py

@@ -2,8 +2,9 @@
 
 The kernel-level tests exercise the two stateless primitives (``quantize`` for
 bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed
-operands); the policy-level tests (recipes, autocast context, per-tensor meta,
-CPU fallbacks of the custom ops) run without a GPU.
+operands); the policy-level tests (recipes, autocast context, per-tensor
+meta) run without a GPU. The primitives themselves are CUDA-only
+(attention-style direct wrappers — no torch.library dispatch layer).
 """
 
 import threading
@@ -14,9 +15,8 @@ import torch.nn.functional as F
 
 import astrai.extension.fp8 as f8mod
 from astrai.extension.fp8 import (
-    DelayedScaling,
-    DynamicScaling,
     FP8Format,
+    FP8Recipe,
     FP8TensorMeta,
     _ScaleRing,
     fp8_autocast,
@@ -24,7 +24,7 @@ from astrai.extension.fp8 import (
     fp8_linear_enabled,
     fp8_state,
 )
-from astrai.extension.ops.fp8 import mm_fp8, quantize
+from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
 from tests.conftest import skip_no_fp8
 
 
@@ -233,7 +233,7 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
     dev = torch.device("cuda")
     state = f8mod.fp8_state()
     state.reset()
-    state.default_recipe = DelayedScaling(history_len=1, margin=0)
+    state.default_recipe = FP8Recipe(history_len=1, margin=0)
     state.default_format = FP8Format.E4M3
     try:
         m, n, k = 32, 16, 64
@@ -272,7 +272,7 @@ def test_fp8_linear_forward_and_backward():
 
     state = f8mod.fp8_state()
     state.reset()
-    state.default_recipe = DynamicScaling()
+    state.default_recipe = FP8Recipe(dynamic=True)
     try:
         out, _, _ = f8mod.fp8_linear_forward(x, weight, bias)
 
@@ -399,13 +399,13 @@ def test_mm_fp8_matches_scaled_mm():
 def test_recipe_scale_from_history():
     """Delayed: max over the window + margin; dynamic: current amax."""
     hist = torch.tensor([1.0, 2.0, 0.5])
-    d = DelayedScaling(history_len=3, margin=0)
+    d = FP8Recipe(history_len=3, margin=0)
     assert torch.allclose(d.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0))
-    d_m = DelayedScaling(history_len=3, margin=2)
+    d_m = FP8Recipe(history_len=3, margin=2)
     assert torch.allclose(
         d_m.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0 / 4.0)
     )
-    dyn = DynamicScaling()
+    dyn = FP8Recipe(dynamic=True)
     amax = torch.tensor([0.25])
     assert torch.allclose(
         dyn.scale_from_history(amax, "e4m3"), torch.tensor(0.25 / 448.0)
@@ -423,62 +423,107 @@ def test_fp8_format_enum():
 
 
 def test_fp8_autocast_context():
-    """fp8_autocast sets and restores recipe + format on the global state."""
+    """fp8_autocast pushes and restores the thread-local active config."""
     state = fp8_state()
-    prev = (state.enabled, state.recipe, state.fp8_format)
+    state.reset()
     try:
         with fp8_autocast(enabled=True, fp8_format="hybrid", update_interval=8):
-            assert state.enabled
-            assert isinstance(state.recipe, DelayedScaling)
-            assert state.recipe.history_len == 8
-            assert state.fp8_format is FP8Format.HYBRID
-            with fp8_autocast(enabled=True, recipe=DynamicScaling(), fp8_format="e4m3"):
-                assert isinstance(state.recipe, DynamicScaling)
-                assert state.fp8_format is FP8Format.E4M3
-            assert state.fp8_format is FP8Format.HYBRID  # restored on exit
-        assert not state.enabled
+            cfg = f8mod._active_config.get()
+            assert cfg is not None and cfg.enabled
+            assert not cfg.recipe.dynamic
+            assert cfg.recipe.history_len == 8
+            assert cfg.fp8_format is FP8Format.HYBRID
+            with fp8_autocast(
+                enabled=True, recipe=FP8Recipe(dynamic=True), fp8_format="e4m3"
+            ):
+                inner = f8mod._active_config.get()
+                assert inner.recipe.dynamic
+                assert inner.fp8_format is FP8Format.E4M3
+            assert f8mod._active_config.get() is cfg  # restored on exit
+        assert f8mod._active_config.get() is None
+        assert not fp8_linear_enabled()
     finally:
-        state.enabled, state.recipe, state.fp8_format = prev
+        state.reset()
 
 
 def test_fp8_tensor_meta_delayed_update():
     """Meta seeds from data; hist/scale are packed views of one state buffer."""
-    meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0))
+    recipe = FP8Recipe(history_len=4, margin=0)
+    meta = FP8TensorMeta(
+        _ScaleRing(torch.device("cpu"), recipe),
+        _ScaleRing(torch.device("cpu"), recipe),
+        _ScaleRing(torch.device("cpu"), recipe),
+    )
     w = torch.randn(8, 8)
     meta.w.seed(w, "e4m3")
     assert meta.w.initialized
     torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1))
-    # [hist | scale] packing: views alias the single state buffer.
-    assert meta.w.state.numel() == 4 + 2
+    # [hist | scale | legacy | amax | done] packing: views alias one buffer.
+    assert meta.w.state.numel() == 4 + 4
     assert meta.w.hist.data_ptr() == meta.w.state.data_ptr()
     assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr()
     meta.w.advance()
     assert meta.w.idx == 1
 
-    # update folds a fresh amax into the window and publishes the next scale
-    amax = torch.tensor([8.0])
-    meta.w.update(amax, "e4m3")
-    torch.testing.assert_close(meta.w.scale, torch.tensor([8.0 / 448.0]))
+    # fold_args hands the kernel the buffer, the slot and the recipe constants
+    args = meta.w.fold_args("e4m3")
+    assert args["ring_state"] is meta.w.state and args["hist_idx"] == 1
+    assert args["fp8_max"] == 448.0 and args["pow2_margin"] == 1.0
 
 
-def test_quantize_cpu_fallback():
-    """CPU fallback of the quantize primitive (scale semantics + amax)."""
-    x = torch.randn(16, 32, dtype=torch.bfloat16)
-    scale = torch.tensor([0.5])  # quantize multiplier
-    x8, amax = quantize(x, scale, "e4m3")
-    assert x8.dtype == torch.float8_e4m3fn
-    ref = (x.float() * 0.5).to(torch.float8_e4m3fn)
-    assert torch.equal(x8, ref)
+@skip_no_fp8
+@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"])
+def test_quantize_dual_and_transposed_orientations(fmt):
+    """quantize_dual yields both orientations from one read; quantize's
+    transposed switch keeps the 2-tuple arity with the [cols][rows] layout."""
+    torch.manual_seed(11)
+    x = torch.randn(37, 67, device="cuda", dtype=torch.bfloat16) * 3
+    mult = _scale(x).reciprocal()
+    x8, amax = quantize(x, mult, fmt)
+    x8T, _ = quantize(x, mult, fmt, transposed=True)
+    d8, d8T, _ = quantize_dual(x, mult, fmt)
+    assert x8T.shape == (67, 37)
+    assert torch.equal(x8.view(torch.uint8), d8.view(torch.uint8))
+    assert torch.equal(x8T.view(torch.uint8), d8T.view(torch.uint8))
+    assert torch.equal(x8T.t().contiguous().view(torch.uint8), x8.view(torch.uint8))
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
 
 
-def test_mm_fp8_cpu_fallback():
-    a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
-    b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
-    scale = torch.tensor([1.0])
-    out = mm_fp8(a8, b8, scale)
-    ref = (a8.float() @ b8.float() * 1.0).to(torch.bfloat16)
-    torch.testing.assert_close(out, ref)
+@skip_no_fp8
+@pytest.mark.parametrize("fmt,fmax", [("e4m3", 448.0), ("e5m2", 57344.0)])
+@pytest.mark.parametrize("margin", [0, 1])
+def test_quantize_ring_fold_matches_host_update(fmt, fmax, margin):
+    """The in-kernel delayed-scaling fold matches a host-side reference."""
+    dev = torch.device("cuda")
+    n, idx = 4, 2
+    torch.manual_seed(3)
+    x = torch.randn(128, 96, dtype=torch.bfloat16, device=dev) * 3
+    mult = torch.tensor([0.01], device=dev)
+    pow2m = float(2**margin)
+
+    # Reference: legacy quantize + the host fold it used to return amax for.
+    x8_ref, amax = quantize(x, mult, fmt)
+    hist = torch.full((n,), 1.0, device=dev)
+    hist[idx] = amax.to(torch.float32)
+    scale = (hist.max() / fmax / pow2m).clamp_min(1e-12).reshape(1)
+
+    # Fused: same window, fold inside the quantize kernel's last block.
+    ring = torch.zeros(n + 4, device=dev)
+    ring[:n].fill_(1.0)
+    x8, _ = quantize(
+        x,
+        mult,
+        fmt,
+        ring_state=ring,
+        hist_idx=idx,
+        fp8_max=fmax,
+        pow2_margin=pow2m,
+    )
+    assert torch.equal(x8.view(torch.uint8), x8_ref.view(torch.uint8))
+    torch.testing.assert_close(ring[:n], hist, rtol=0, atol=0)
+    torch.testing.assert_close(ring[n : n + 1], scale, rtol=0, atol=0)
+    assert float(ring[n + 2]) == 0.0  # amax slot self-cleaned
+    assert int(ring[n + 3].view(torch.int32)) == 0  # done counter reset
 
 
 # --------------------------------------------------------------------------

tests/inference/test_scheduler.py

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

tests/test_serve_runtime.py

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

tests/trainer/test_callbacks.py

@@ -1,8 +1,12 @@
+from pathlib import Path
+
 import torch
 
 from astrai.model.components.decoder_block import DecoderBlock
+from astrai.serialization import Checkpoint
 from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
 from astrai.trainer.trainer import Trainer
+from tests.helpers import RandomTokenDataset
 
 
 def test_gradient_checkpointing_enable_disable(test_model):
@@ -135,3 +139,34 @@ def test_callback_integration(
     assert "on_train_begin" in callback_calls
     assert "on_batch_end" in callback_calls
     assert "on_epoch_end" in callback_calls
+
+
+def test_checkpoint_captures_completed_optimizer_step(
+    base_test_env, train_config_factory, device
+):
+    """Checkpoint state must include the update represented by its step number."""
+    model = base_test_env["model"]
+    initial_state = {
+        name: tensor.detach().cpu().clone()
+        for name, tensor in model.state_dict().items()
+    }
+    train_config = train_config_factory(
+        model_fn=lambda: model,
+        dataset=RandomTokenDataset(length=2),
+        test_dir=base_test_env["test_dir"],
+        device=device,
+        batch_per_device=2,
+        ckpt_interval=1,
+    )
+
+    Trainer(train_config).train()
+
+    checkpoint = Checkpoint.load(
+        str(Path(base_test_env["test_dir"]) / "epoch_0_step_1")
+    )
+    assert any(
+        not torch.equal(checkpoint.state_dict[name].cpu(), initial_tensor)
+        for name, initial_tensor in initial_state.items()
+    )
+    assert checkpoint.extra["optimizer"]["state"]
+    assert checkpoint.extra["scheduler"]["last_epoch"] == 1