perf: interleave prefetch into the mma phase

- our 128x128 fast loop padded 26 NOPs between the 32 QMMAs while all four LDGSTS sat bunched at the loop tail: ptxas had no independent instructions to fill the tensor-pipe issue gaps, the exact structure the decompiled cuBLAS loop (166 i, NOP=0) and CUTLASS MmaMultistage avoid by issuing cp.async in small groups inside the MMA phase (copy_tiles_and_advance per warp-tile batch)
- the steady-state prefetch is now a loop-carried register pair per congruous operand (PrefetchCarry: swizzled stage offset + global source, constructed once from the same (r, c0) mapping as the interior loader), whose chunks ride after the first and last k_seg MMA batches — SASS: 282 -> 110 instructions, 0 branches, 0 UIMAD.WIDE magic-divisions, LDGSTS interleaved inside the QMMA range, 26 -> 19 NOPs, still 128 regs (2 CTAs/SM)
- the wait-count dispatch ladder (16 instructions of ISETP/SEL picking DEPBAR immediates) and the per-k-tile (tile % ring) * stage_bytes recomputation (UIMAD.WIDE by 0x55555555) are gone: the prologue commits unconditionally so the wait_group<kStages-1> immediate is valid for every iteration, and both read and write stage addresses advance as carried pointers with an equality wrap
- cp_async.cuh splits the emitter from its policies: one raw PTX site (cp_async_16_raw) plus wrappers for unconditional/predicated and pointer/offset destinations, and the now-unused wait_group_dispatch ladder is deleted; the dispatch flip: with the stall gone the big CTA wins the whole former dip band (1280^3 kernel-level 98.2->104.4T), so prefer_small_cta keeps only the sub-5/8-wave band and the single-wave s3 variant is retired

Benchmark: L20 (sm_89), kernel-level sweep 128s2ff: 1024^3 104.2->114.0T (cuBLAS 151.2), 1152^3 130.4->145.9T (157.2), 1280^3 98.2->104.7T (163.8), 1536^3 141.4->150.4T (182.3), 2048^3 178.1->189.6T (202.8), 4096^3 197.6->208.8T (223.7), 8192^3 207.6->219.1T (227.5). CUDA-graph e2e: 1024^3 108.2->118.4T, 1152^3 137.7->150.7T, 1536^3 143.4->153.7T, 2048^3 179.4->191.9T, 8192^3 199.0->208.9T, 1280^3 108.3->106.6T (old dip-band rule re-measured 106.1T — within noise). Four-layout C++ suite and 596 pytests pass.

Dockerfile

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

astrai/extension/fp8.py

@@ -36,10 +36,7 @@ from typing import Dict, List, Optional
 import torch
 from torch.library import Library
 
-from astrai.extension.ops.fp8 import (
+from astrai.extension.ops.fp8 import mm_fp8, quantize
-    linear_backward_fp8,
-    linear_forward_fp8,
-)
 
 # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
 FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -95,11 +92,10 @@ class DynamicScaling(FP8Recipe):
 
 class _ScaleRing:
     """One operand's delayed-scaling state: a float32 buffer
-    ``[hist[n] | scale | counter]`` (views). The quantize kernel's last-finishing
+    ``[hist[n] | scale | counter]`` (views). ``update`` folds the amax
-    block records the measured amax into ``hist[idx]``, reduces the window and
+    returned by the quantize primitive into ``hist[idx]`` and publishes the
-    publishes the next scale entirely on device — the Python-side write/max/write
+    next scale from the window; ``idx`` advances host-side each step. The
-    chain is gone. The counter slot stays int32-zero (float bits) between
+    trailing slot is a legacy counter kept for state-buffer compatibility.
-    launches; ``idx`` advances host-side each step.
     """
 
     __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -114,7 +110,7 @@ class _ScaleRing:
         self.initialized = False
 
     def advance(self) -> None:
-        """Rotate to the next history slot after an in-kernel finalize."""
+        """Rotate to the next history slot after metadata update."""
         self.idx = (self.idx + 1) % self.hist.numel()
 
     def seed(self, t: torch.Tensor, fmt: str) -> None:
@@ -123,12 +119,15 @@ class _ScaleRing:
         self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
         self.initialized = True
 
+    def update(self, amax: torch.Tensor, fmt: str) -> None:
+        self.hist[self.idx].copy_(amax.reshape(()))
+        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
+
 
 class FP8TensorMeta:
-    """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/
+    """Per-weight delayed-scaling state for ``w``, ``x`` and ``g``.
-    ``g``). Fused kernels record amax while quantizing, so the scale used at step
+
-    N reflects amax from steps < N. DynamicScaling never allocates a meta — it
+    DynamicScaling never allocates a meta; it measures the current amax inline.
-    measures the current amax inline.
     """
 
     __slots__ = ("w", "x", "g")
@@ -213,7 +212,11 @@ class FP8State:
         return meta
 
     def reset(self) -> None:
+        """Restore construction defaults (switch, recipe, format) and drop all
+        per-weight metas — a full state reset for tests / reconfiguration."""
         self.default_enabled = False
+        self.default_recipe = DelayedScaling()
+        self.default_format = FP8Format.HYBRID
         self._metas.clear()
 
 
@@ -307,18 +310,9 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor
     return recipe.scale_from_history(amax, fmt)
 
 
-_zero_bias: Dict[Optional[int], torch.Tensor] = {}
+def _is_fp8(dtype: torch.dtype) -> bool:
-
+    """A pre-quantized weight takes the GEMM directly (no re-quantize)."""
-
+    return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
-def _empty_bias(x: torch.Tensor) -> torch.Tensor:
-    """Per-device cached 0-element bf16 bias (the binding only checks numel —
-    never mutated), saving a CUDA allocation per bias-less linear."""
-    key = x.device.index
-    t = _zero_bias.get(key)
-    if t is None:
-        t = torch.empty(0, device=x.device, dtype=torch.bfloat16)
-        _zero_bias[key] = t
-    return t
 
 
 def fp8_linear_forward(
@@ -326,9 +320,10 @@ def fp8_linear_forward(
 ):
     """Scaled fp8 linear forward (called from the aten::linear impl).
 
-    Pure FP8 path for both recipes: quantize x/w with the active scales, run the
+    Composed from the two stateless primitives: quantize x/w with the active
-    pre-quantized GEMM. Delayed scaling finalizes the rings inside the quantize
+    scales, run the pre-quantized GEMM with the bias fused into its epilogue.
-    kernels (amax folded into the window, next scale published on device);
+    Delayed scaling 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).
@@ -337,13 +332,17 @@ def fp8_linear_forward(
     if cfg is None:
         cfg = _current_config()
     fmt = cfg.fp8_format.fwd()
-    margin = cfg.recipe.margin
+    if isinstance(cfg.recipe, DynamicScaling):
-    if bias is None:
-        bias = _empty_bias(x)
-    if isinstance(cfg.recipe, DynamicScaling):  # measure-then-quantize, no state
         sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
         sw = _dynamic_scale(w, cfg.recipe, fmt)
-        out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt)
+        x8, _ = quantize(x, sx.reciprocal(), fmt)
+        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
+        # Bias fuses into the GEMM epilogue (fp32 add before the single bf16
+        # rounding — one rounding fewer than the separate out + bias pass);
+        # None passes through to the kernel's no-bias path.
+        out = mm_fp8(
+            x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
+        ).reshape(*x.shape[:-1], w.size(0))
         return out, sx, sw
 
     meta = state.get_weight_meta(w)
@@ -351,35 +350,20 @@ def fp8_linear_forward(
         meta.w.seed(w, fmt)
     if not meta.x.initialized:
         meta.x.seed(x, fmt)
-    # The quantize kernels finalize each ring in-kernel and overwrite the ring's
+    sx, sw = meta.x.scale.clone(), meta.w.scale.clone()
-    # scale slot, which ALIASES meta.*.scale (a view into the state buffer).
+    x8, amax_x = quantize(x, sx.reciprocal(), fmt)
-    # Snapshot the scales first so the GEMM dequantizes with the SAME scale the
+    if _is_fp8(w.dtype):
-    # operands were quantized with, and so the backward can reuse this step's
+        w8, amax_w = w, None
-    # scale (gradient consistency with the forward). The ring finalize still
-    # publishes the next step's scale into the original slot.
-    if w.dtype is not torch.bfloat16:  # static pre-quantized weight
-        w_arg, sw_arg, w_ring = w, meta.w.scale, None
     else:
-        w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state
+        w8, amax_w = quantize(w, sw.reciprocal(), fmt)
-    sx = meta.x.scale.clone()
+    out = mm_fp8(
-    sw = sw_arg.clone()
+        x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
-    out, _x8, _w8, _ax, _aw = linear_forward_fp8(
+    ).reshape(*x.shape[:-1], w.size(0))
-        x,
+    meta.x.update(amax_x, fmt)
-        w_arg,
+    if amax_w is not None:
-        bias,
+        meta.w.update(amax_w, fmt)
-        sx,
-        sw,
-        fmt,
-        None,
-        meta.x.state,
-        meta.x.idx,
-        margin,
-        w_ring,
-        meta.w.idx,
-        margin,
-    )
     meta.x.advance()
-    if w_ring is not None:
+    if amax_w is not None:
         meta.w.advance()
     return out, sx, sw
 
@@ -410,37 +394,28 @@ class _LinearFp8(torch.autograd.Function):
     def backward(ctx, g):
         x, w, _sx_fwd, _sw_fwd = ctx.saved_tensors
         fmt = ctx.fmt_bwd
-        # Per-recipe scale/ring selection; both branches share one call below.
+        # Flatten leading dims (the forward GEMMs ran on [-1, N] / [-1, K]
+        # views; the kernels only accept 2D operands).
+        g2 = g.reshape(-1, g.size(-1))
         if ctx.is_dynamic:
-            sg = _dynamic_scale(g, ctx.recipe, fmt)
+            sg = _dynamic_scale(g2, ctx.recipe, fmt)
             sw = _dynamic_scale(w, ctx.recipe, fmt)
             sx = _dynamic_scale(x, ctx.recipe, fmt)
-            ring, idx = None, 0
         else:
             meta = ctx.meta
             if not meta.g.initialized:
-                meta.g.seed(g, fmt)
+                meta.g.seed(g2, fmt)
-            # Snapshot the g scale before its ring finalize overwrites the slot
-            # (same aliasing as the forward); reuse the forward's w/x scales so
-            # the backward quantizes with the scale the forward actually used.
             sg = meta.g.scale.clone()
-            ring, idx = meta.g.state, meta.g.idx
             sw, sx = _sw_fwd, _sx_fwd
-        grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8(
+        g8, amax_g = quantize(g2, sg.reciprocal(), fmt)
-            g,
+        x8, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt)
-            x,
+        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
-            w,
+        grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape)  # g8[m,n] @ w8[n,k]
-            list(ctx.needs_input_grad),
+        grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True)  # g8.T @ x8
-            sg,
+        grad_b = g2.sum(0).to(torch.bfloat16)
-            sw,
-            sx,
-            fmt,
-            ring,
-            idx,
-            ctx.recipe.margin,
-        )
         if not ctx.is_dynamic:
-            meta.g.advance()  # the g quantize kernel finalized the ring in-kernel
+            meta.g.update(amax_g, fmt)
+            meta.g.advance()
         return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
 
 

astrai/extension/ops/fp8.py

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

csrc/kernels/common/cp_async.cuh

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

csrc/kernels/fp8/common.h

@@ -48,53 +48,49 @@ using transpose_layout_t = typename transpose_layout<Layout>::type;
 
 // Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
 // WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
-// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile and
+// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
-// the cp.async pipeline depth.
+// warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128 CTA, or 32x32 on the
-template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages>
+// cuBLAS-style 64x64 small CTA that lifts small-shape occupancy) and the
+// cp.async pipeline depth.
+template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
+          int WarpM = 64, int WarpN = 32>
 struct Fp8GemmTraits {
     static constexpr FP8Format kFormat = Fmt;
     static constexpr int kBlockM = BlockM;
     static constexpr int kBlockN = BlockN;
     static constexpr int kK = K;
     static constexpr int kStages = Stages;
+    static constexpr int kWarpM = WarpM;
+    static constexpr int kWarpN = WarpN;
     static constexpr bool kIsE5M2 = (Fmt == FP8Format::E5M2);
     static constexpr __nv_fp8_interpretation_t kNvFormat =
         kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
     static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
 
-    // Derived launch geometry: 64x32 warp tiles give the CTA thread count.
+    // Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. The
-    // The shared-memory budget is layout-aware (crosswise operands add K-
+    // shared-memory budget is layout-aware (crosswise operands add K-major
-    // major staging + a canonical buffer), so it lives in Fp8GemmSmem in
+    // staging + a canonical buffer), so it lives in Fp8GemmSmem in gemm.cuh
-    // gemm.cuh together with the resident-CTA hint for __launch_bounds__.
+    // together with the resident-CTA hint for __launch_bounds__.
-    static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
+    static constexpr int kWarpsM = BlockM / WarpM;
+    static constexpr int kWarpsN = BlockN / WarpN;
+    static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
+    static_assert(kWarpsM * WarpM == BlockM && kWarpsN * WarpN == BlockN,
+                  "warp tiles must exactly tile the CTA");
+    static_assert(WarpM % 16 == 0 && WarpN % 8 == 0,
+                  "warp tile must be a multiple of the m16n8 MMA shape");
 };
 
-// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional
+// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
-// delayed-scaling ring finalization. Separate from FP8Params so each
+// with fused amax.
-// operator owns exactly the fields it touches (the GEMM never reads amax /
-// ring state). Same NSDMI rationale: amax / ring_state gate optional paths
-// via null checks. Still an aggregate, still trivially copyable.
 struct FP8QuantizeParams {
-    // BF16 input and FP8 output buffers; scale_a is the quantization step
+    // Float input and FP8 output buffers; scale is the quantization
-    // (device scalar). amax_a (may be null) is zero-initialized by the
+    // multiplier (device scalar). amax (may be null) is zero-initialized by
-    // binding and receives the raw-domain absolute maximum.
+    // the binding and receives the raw-domain absolute maximum.
-    const void* __restrict__ a_ptr = nullptr;
+    const void* __restrict__ input_ptr = nullptr;
-    void* __restrict__ out_ptr = nullptr;
+    void* __restrict__ output_ptr = nullptr;
-    const float* __restrict__ scale_a = nullptr;
-    float* __restrict__ amax_a = nullptr;
 
-    // Optional delayed-scaling ring finalization. ring_state packs
+    const float* __restrict__ scale = nullptr;
-    // [hist[ring_len] | scale | counter] with ring_len = numel - 2. When
+    float* __restrict__ amax = nullptr;
-    // non-null and amax_a is set, the last-finishing block records the
-    // measured amax into hist[ring_idx], reduces the window and publishes
-    // the next step's scale (max(hist) / fp8_max / 2^ring_margin) — the
-    // fused replacement for the eager hist-write / max / scale-write chain,
-    // at zero extra launches. The counter slot is a persistent zero-armed
-    // int32 (float bits) electing the last block each launch.
-    float* ring_state = nullptr;
-    int ring_len = 0;
-    int ring_idx = 0;
-    int ring_margin = 0;
 
     // Element count (only the elementwise quantize kernel uses it).
     int total = 0;
@@ -103,28 +99,30 @@ struct FP8QuantizeParams {
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
 // through the pre-quantized GEMM kernels. Each kernel touches only the
 // fields it needs; buffers are raw pointers packed by the torch binding.
-// Pointer members default to null (same NSDMI rationale as AttentionParams:
+// Pointer members default to null so optional paths cannot hold garbage.
-// bias / out_scale gate optional paths via null checks, so a partially
-// packed struct must never hold garbage non-null pointers). Still an
-// aggregate, still trivially copyable.
 struct FP8Params {
     // Inputs: a/b are FP8 for the pre-quantized path. Scales are
     // quantization steps (device scalars).
+    // 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.
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
-    const void* __restrict__ bias = nullptr;
+    const void* __restrict__ bias_ptr = nullptr;
-    const float* __restrict__ scale_a = nullptr;
-    const float* __restrict__ scale_b = nullptr;
-    const float* __restrict__ bias_scale = nullptr;
-    // Output: BF16 or FP8 (E4M3). out_scale is the output quantization step
-    // (FP8 output only).
     void* __restrict__ out_ptr = nullptr;
-    const float* __restrict__ out_scale = nullptr;
 
+    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;
 
+    // Batched (bmm) geometry: grid.z slices step the operand/output pointers
+    // by 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

csrc/kernels/fp8/ops.cu

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

csrc/kernels/fp8/quantize.cuh

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

csrc/tests/fp8_test.cu

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

docs/developer/cuda_kernels.md

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

tests/extension/test_fp8_mma.py

@@ -1,9 +1,9 @@
 """FP8 primitives: kernel-level (CUDA) and policy-level (CPU-verifiable) tests.
 
-The kernel-level tests exercise the pure FP8 path (quantize_bf16 + mm_fp8 for
+The kernel-level tests exercise the two stateless primitives (``quantize`` for
-the forward GEMM, quantize + pre-quantized GEMMs for the backward); the
+bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed
-policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks
+operands); the policy-level tests (recipes, autocast context, per-tensor meta,
-of the custom ops) run without a GPU.
+CPU fallbacks of the custom ops) run without a GPU.
 """
 
 import threading
@@ -24,12 +24,7 @@ from astrai.extension.fp8 import (
     fp8_linear_enabled,
     fp8_state,
 )
-from astrai.extension.ops.fp8 import (
+from astrai.extension.ops.fp8 import mm_fp8, quantize
-    linear_backward_fp8,
-    linear_forward_fp8,
-    mm_fp8,
-    quantize_bf16,
-)
 from tests.conftest import skip_no_fp8
 
 
@@ -37,8 +32,11 @@ def _scale(tensor):
     return (tensor.abs().amax().float() / 448.0).clamp_min(1e-12)
 
 
-def _quantize(tensor, scale):
+def _quantize(tensor, scale, fmt="e4m3"):
-    return (tensor.float() / scale).to(torch.float8_e4m3fn).float()
+    """Reference quantize: multiply by the reciprocal (the kernel's exact
+    arithmetic — a plain divide flips fp8 boundary cases by one ulp)."""
+    dtype = torch.float8_e5m2 if fmt == "e5m2" else torch.float8_e4m3fn
+    return (tensor.float() * scale.reciprocal()).to(dtype).float()
 
 
 # --------------------------------------------------------------------------
@@ -57,9 +55,9 @@ def test_fp8_mm_matches_explicit_quantization(m, n, k):
     b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
     scale_a = _scale(a)
     scale_b = _scale(b)
-    a8, _ = quantize_bf16(a, scale_a, "e4m3")
+    a8, _ = quantize(a, scale_a.reciprocal(), "e4m3")
-    b8, _ = quantize_bf16(b, scale_b, "e4m3")
+    b8, _ = quantize(b, scale_b.reciprocal(), "e4m3")
-    out = mm_fp8(a8, b8, scale_a, scale_b)
+    out = mm_fp8(a8, b8, scale_a * scale_b)
     expected = (_quantize(a, scale_a) @ _quantize(b, scale_b) * scale_a * scale_b).to(
         torch.bfloat16
     )
@@ -70,79 +68,167 @@ def test_fp8_mm_matches_explicit_quantization(m, n, k):
 
 
 @skip_no_fp8
-def test_quantize_bf16_returns_amax():
+@pytest.mark.parametrize("in_dtype", [torch.bfloat16, torch.float16, torch.float32])
-    """quantize_bf16 returns (x8, amax); amax tracks the *raw* values and the
+@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"])
-    caller never clears it (zero-initialized inside the kernel entry)."""
+def test_quantize_input_dtypes(in_dtype, fmt):
+    """quantize accepts bf16/fp16/fp32 inputs; bytes and amax match the
+    explicit (value * multiplier) reference."""
     torch.manual_seed(3)
-    x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16)
+    x = torch.randn(64, 128, device="cuda", dtype=torch.float32) * 0.5
+    x = x.to(in_dtype)
     scale = torch.tensor([0.5], device="cuda")
-    x8, amax = quantize_bf16(x, scale, "e4m3")
+    x8, amax = quantize(x, scale, fmt)
-    assert x8.dtype == torch.float8_e4m3fn
+    out_dtype = torch.float8_e5m2 if fmt == "e5m2" else torch.float8_e4m3fn
+    assert x8.dtype == out_dtype
     assert x8.shape == x.shape
     assert amax.shape == (1,)
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
-    ref = (x.float() / 0.5).to(torch.float8_e4m3fn)
+    ref = (x.float() * 0.5).to(out_dtype)
     assert torch.equal(x8, ref)
 
 
 @skip_no_fp8
-def test_quantize_bf16_e5m2_format():
+def test_quantize_e5m2_format():
     x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
-    x8, amax = quantize_bf16(x, torch.tensor([0.1], device="cuda"), "e5m2")
+    x8, amax = quantize(x, torch.tensor([10.0], device="cuda"), "e5m2")
     assert x8.dtype == torch.float8_e5m2
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
 
 
 @skip_no_fp8
-def test_quantize_ring_in_kernel_finalize():
+@pytest.mark.parametrize("trans_a", [False, True])
-    """The quantize kernel finalizes the delayed-scaling ring in-kernel: the
+@pytest.mark.parametrize("trans_b", [False, True])
-    measured amax lands in hist[idx], the window reduces to the next step's
+def test_mm_fp8_transposed_operands(trans_a, trans_b):
-    scale on device, and the counter re-arms for the next launch."""
+    """mm_fp8 handles all four operand layouts via trans_a/trans_b."""
-    torch.manual_seed(21)
+    torch.manual_seed(17)
-    dev = torch.device("cuda")
+    m, n, k = 19, 13, 37
-    ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0))
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)  # A [M][K]
-    x0 = torch.randn(256, 256, device=dev, dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)  # B^T [N][K]
-    w = torch.randn(256, 256, device=dev, dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
-    sw = torch.tensor([1.0], device=dev)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
-    ring.seed(x0, "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
-    hist0 = ring.hist.clone()
+    a_op = a8.t().contiguous() if trans_a else a8
+    b_op = b8 if trans_b else b8.t().contiguous()
 
-    # Step over three fresh tensors: each launch folds its amax into
+    out = mm_fp8(a_op, b_op, sa * sb, trans_a=trans_a, trans_b=trans_b)
-    # hist[idx] and publishes max(hist)/448 as the next scale.
+    assert out.shape == (m, n)
-    idx = 0
+    expected = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
-    for _ in range(3):
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
-        x = torch.randn(256, 256, device=dev, dtype=torch.bfloat16) * (2.0 + 4.0 * _)
+
-        _ = linear_forward_fp8(
+
-            x,
+@skip_no_fp8
-            w,
+@pytest.mark.parametrize("bias_on", [False, True])
-            None,
+def test_mm_fp8_fused_bias(bias_on):
-            ring.scale,
+    """Epilogue-fused bias matches the unfused out + bias reference (single
-            sw,
+    fp32 rounding vs the reference's double rounding keeps it within 1 ulp),
-            "e4m3",
+    including N-tail columns and batched broadcast."""
-            None,
+    torch.manual_seed(31)
-            ring.state,
+    m, n, k = 19, 13, 37  # odd n exercises the guarded bias loads
-            idx,
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-            0,
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-        )
+    sa, sb = _scale(a), _scale(b)
-        torch.cuda.synchronize()
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
-        expected_hist = hist0.clone()
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
-        expected_hist[idx] = x.abs().amax().float()
+    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
-        torch.testing.assert_close(ring.hist, expected_hist)
+
-        expected_scale = (expected_hist.max() / 448.0).reshape(1)
+    out = mm_fp8(a8, b8, sa * sb, trans_b=True, bias=bias if bias_on else None)
-        torch.testing.assert_close(ring.scale, expected_scale, rtol=1e-6, atol=1e-12)
+    base = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
-        # counter re-armed to int32 zero
+    expected = base + bias if bias_on else base
-        assert ring.state[-1].view(torch.int32).item() == 0
+    # bias is O(1) against O(sqrt(k)) accumulators: absolute tolerance rules
-        hist0 = expected_hist.clone()
+    torch.testing.assert_close(out, expected, atol=0.13, rtol=0.01)
-        idx = (idx + 1) % 4
+
+    # Batched broadcast: bias applies to every batch slice (each slice gets
+    # its own reference from its own operand values).
+    ab = torch.randn(3, m, k, device="cuda", dtype=torch.bfloat16)
+    ab8, _ = quantize(ab, sa.reciprocal(), "e4m3")
+    outb = mm_fp8(ab8, b8, sa * sb, trans_b=True, bias=bias)
+    assert outb.shape == (3, m, n)
+    for i in range(3):
+        expected_b = (_quantize(ab[i], sa) @ _quantize(b, sb).t() * sa * sb).to(
+            torch.bfloat16
+        ) + bias
+        torch.testing.assert_close(outb[i], expected_b, atol=0.13, rtol=0.01)
+
+
+@skip_no_fp8
+@pytest.mark.parametrize("trans_a", [False, True])
+@pytest.mark.parametrize("trans_b", [False, True])
+def test_mm_fp8_batched(trans_a, trans_b):
+    """3D operands run as one bmm launch: all four layouts, odd shapes."""
+    torch.manual_seed(23)
+    batch, m, n, k = 4, 19, 13, 37
+    a = torch.randn(batch, m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(batch, n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+    a_op = a8.transpose(-2, -1).contiguous() if trans_a else a8
+    b_op = b8 if trans_b else b8.transpose(-2, -1).contiguous()
+
+    out = mm_fp8(a_op, b_op, sa * sb, trans_a=trans_a, trans_b=trans_b)
+    assert out.shape == (batch, m, n)
+    # flags + transposed buffers reconstruct the original operands: the math
+    # is always A_orig @ B_orig^T regardless of the layout combination.
+    expected = (_quantize(a, sa) @ _quantize(b, sb).transpose(-2, -1) * sa * sb).to(
+        torch.bfloat16
+    )
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
+
+
+@skip_no_fp8
+def test_mm_fp8_batched_broadcast():
+    """A size-1 batch broadcasts across the other operand (matmul rules),
+    and a 2D operand broadcasts across a 3D one."""
+    torch.manual_seed(29)
+    batch, m, n, k = 3, 16, 8, 32
+    a = torch.randn(batch, m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(1, n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+
+    out = mm_fp8(a8, b8, sa * sb, trans_b=True)
+    assert out.shape == (batch, m, n)
+    expected = (_quantize(a, sa) @ _quantize(b, sb).transpose(-2, -1) * sa * sb).to(
+        torch.bfloat16
+    )
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
+
+    # 2D weight broadcast over 3D activations
+    w8 = b8[0]
+    out2 = mm_fp8(a8, w8, sa * sb, trans_b=True)
+    assert out2.shape == (batch, m, n)
+    torch.testing.assert_close(out2, expected, atol=0.125, rtol=0.01)
+
+
+@skip_no_fp8
+def test_mm_fp8_col_major_view_zero_copy():
+    """An inner-transposed view (.t() of a contiguous buffer) folds into the
+    layout tag with no device copy — the only allocation is the output."""
+    torch.manual_seed(31)
+    m, n, k = 64, 64, 64
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+
+    torch.cuda.synchronize()
+    before = torch.cuda.memory_allocated()
+    out = mm_fp8(a8.t(), b8, sa * sb, trans_a=True, trans_b=True)
+    torch.cuda.synchronize()
+    grew = torch.cuda.memory_allocated() - before
+    assert grew == out.numel() * out.element_size()  # no operand copy
+
+    expected = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
 
 
 @skip_no_fp8
 def test_delayed_scaling_forward_uses_snapshot_scale():
-    """Regression: the in-kernel ring finalize overwrites the scale slot, which
+    """The delayed scale for step N is computed from amax(steps < N); the
-    aliases the scale the GEMM must dequantize with. The forward must snapshot
+    forward must snapshot the scale before the ring update, so a changing
-    the delayed scale first, so a changing amax across steps does not leak the
+    amax across steps does not leak the next-step scale into the output."""
-    next-step scale into the output (otherwise out is off by
-    scale_next / scale_current)."""
     torch.manual_seed(11)
     dev = torch.device("cuda")
     state = f8mod.fp8_state()
@@ -153,8 +239,7 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
         m, n, k = 32, 16, 64
         x1 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.5
         # Smaller amax than x1: the delayed scale (amax(x1)/448) still covers
-        # x2 without fp8 saturation, while the next-step scale would differ —
+        # x2 without fp8 saturation, while the next-step scale would differ.
-        # exactly the condition that exposed the overwrite bug.
         x2 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.35
         w = torch.randn(n, k, device=dev, dtype=torch.bfloat16) * 0.5
         bias = torch.zeros(n, device=dev, dtype=torch.bfloat16)
@@ -177,93 +262,57 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
 
 @skip_no_fp8
 def test_fp8_linear_forward_and_backward():
+    """The composed strategy path: forward quantize+GEMM+bias, backward
+    dX/dW GEMMs on transposed operands (E5M2 in hybrid)."""
     torch.manual_seed(7)
     m, n, k = 19, 13, 37
     x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
     weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16)
     bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
-    scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad)
 
-    out, x8, w8, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w)
+    state = f8mod.fp8_state()
-    grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
+    state.reset()
-        grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3"
+    state.default_recipe = DynamicScaling()
-    )
+    try:
+        out, _, _ = f8mod.fp8_linear_forward(x, weight, bias)
 
-    qx = _quantize(x, scale_x)
+        sx, sw = _scale(x), _scale(weight)
-    qw = _quantize(weight, scale_w)
+        qx = _quantize(x, sx)
-    qg = _quantize(grad, scale_g)
+        qw = _quantize(weight, sw)
-    expected_out = (qx @ qw.t() * scale_x * scale_w + bias).to(torch.bfloat16)
+        expected_out = (qx @ qw.t() * sx * sw + bias).to(torch.bfloat16)
-    expected_grad_x = (qg @ qw * scale_g * scale_w).to(torch.bfloat16)
+        torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01)
-    expected_grad_w = (qg.t() @ qx * scale_g * scale_x).to(torch.bfloat16)
 
-    torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01)
+        # backward through the aten::linear integration (hybrid E5M2). The
-    torch.testing.assert_close(grad_x, expected_grad_x, atol=0.125, rtol=0.01)
+        # incoming gradient is 2*out of the *fp8* forward (bf16-rounded), not
-    torch.testing.assert_close(grad_w, expected_grad_w, atol=0.125, rtol=0.01)
+        # 2*exact — derive the reference from the actual output.
-    torch.testing.assert_close(grad_b, grad.sum(0).to(torch.bfloat16))
+        xr = x.detach().clone().requires_grad_()
-    torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1))
+        wr = weight.detach().clone().requires_grad_()
-    torch.testing.assert_close(amax_w, weight.abs().amax().float().reshape(1))
+        br = bias.detach().clone().requires_grad_()
-    torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
+        with fp8_autocast(enabled=True):
+            loss = F.linear(xr, wr, br).float().pow(2).sum()
+        loss.backward()
 
-
+        g = (2 * out.float()).to(torch.bfloat16).float()  # actual grad wrt out
-@skip_no_fp8
+        # the dynamic path measures current-step amax in the bwd fmt (E5M2);
-def test_linear_backward_e5m2_gradients():
+        # amax must be taken in fp32 — a bf16-rounded scale flips E5M2
-    """Hybrid backward: gradient GEMMs run in E5M2 (larger dynamic range)."""
+        # boundary rounding (2-bit mantissa) and the reference drifts.
-    torch.manual_seed(5)
+        e5 = 57344.0
-    m, n, k = 32, 16, 64
+        sg = (g.abs().amax() / e5).clamp_min(1e-12)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 3.0
+        sw5 = (weight.abs().amax().float() / e5).clamp_min(1e-12)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+        sx5 = (x.abs().amax().float() / e5).clamp_min(1e-12)
-    grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) * 10.0
+        expected_grad_x = (
-    sg = _scale(grad) * 0.5
+            _quantize(g, sg, "e5m2") @ _quantize(weight, sw5, "e5m2") * sg * sw5
-    sw = _scale(weight)
+        ).to(torch.bfloat16)
-    sx = _scale(x)
+        expected_grad_w = (
-
+            _quantize(g, sg, "e5m2").t() @ _quantize(x, sx5, "e5m2") * sg * sx5
-    grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
+        ).to(torch.bfloat16)
-        grad, x, weight, [1, 1, 1], sg, sw, sx, "e5m2"
+        torch.testing.assert_close(xr.grad, expected_grad_x, atol=0.5, rtol=0.05)
-    )
+        torch.testing.assert_close(wr.grad, expected_grad_w, atol=0.5, rtol=0.05)
-
+        torch.testing.assert_close(
-    def q5(t, s):
+            br.grad, g.sum(0).to(torch.bfloat16), atol=0.5, rtol=0.05
-        return (t.float() / s).to(torch.float8_e5m2).float()
+        )
-
+    finally:
-    qg = q5(grad, sg)
+        state.reset()
-    qw = q5(weight, sw)
-    qx = q5(x, sx)
-    expected_grad_x = (qg @ qw * sg * sw).to(torch.bfloat16)
-    expected_grad_w = (qg.t() @ qx * sg * sx).to(torch.bfloat16)
-    torch.testing.assert_close(grad_x, expected_grad_x, atol=0.5, rtol=0.05)
-    torch.testing.assert_close(grad_w, expected_grad_w, atol=0.5, rtol=0.05)
-    torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
-
-
-@skip_no_fp8
-def test_fp8_linear_static_fp8_weight_and_bias():
-    """Static fp8 inference: pre-quantized w8/b8 + their scales take the GEMM
-    directly (no weight quantize, amax_w = 0); the bias is fused in the
-    epilogue (bf16 and fp8 bias share the fused path)."""
-    torch.manual_seed(9)
-    m, n, k = 67, 45, 129
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.5
-    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) * 0.5
-    sx, sw, sb = _scale(x), _scale(weight), _scale(bias)
-
-    w8, _ = quantize_bf16(weight, sw, "e4m3")
-    b8, _ = quantize_bf16(bias, sb, "e4m3")
-    out, x8, w8_back, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb)
-
-    qx = _quantize(x, sx)
-    qw = _quantize(weight, sw)
-    qb = _quantize(bias, sb)
-    expected = (qx @ qw.t() * sx * sw + qb * sb).to(torch.bfloat16)
-    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
-    torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1))
-    assert amax_w.item() == 0.0  # nothing measured on the static path
-    assert w8_back is w8  # pre-quantized w handed straight back
-
-    # bf16 bias stays bf16 on the same fused-epilogue path
-    out_bf16bias, *_ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3")
-    expected_b = (qx @ qw.t() * sx * sw + bias.float()).to(torch.bfloat16)
-    torch.testing.assert_close(out_bf16bias, expected_b, atol=0.125, rtol=0.01)
 
 
 @skip_no_fp8
@@ -279,24 +328,24 @@ def test_fp8_linear_backward_outside_autocast():
     bias = torch.randn(96, device="cuda", dtype=torch.bfloat16, requires_grad=True)
     xr, wr, br = (t.detach().clone().requires_grad_() for t in (x, weight, bias))
 
-    calls = {"bwd": 0}
+    calls = {"fwd": 0}
-    orig = f8mod.linear_backward_fp8
+    orig = f8mod.fp8_linear_forward
 
     def spy(*args, **kwargs):
-        calls["bwd"] += 1
+        calls["fwd"] += 1
         return orig(*args, **kwargs)
 
-    f8mod.linear_backward_fp8 = spy
+    f8mod.fp8_linear_forward = spy
     try:
         with fp8_autocast(enabled=True):
             out = F.linear(x, weight, bias)
         assert type(out.grad_fn).__name__ == "_LinearFp8Backward"
         out.float().pow(2).sum().backward()  # outside the autocast region
     finally:
-        f8mod.linear_backward_fp8 = orig
+        f8mod.fp8_linear_forward = orig
         f8mod.fp8_state().reset()
 
-    assert calls["bwd"] == 1  # fp8 kernels, not the bf16 fallback
+    assert calls["fwd"] == 1  # fp8 kernels, not the bf16 fallback
     ref = F.linear(xr, wr, br)
     ref.float().pow(2).sum().backward()
 
@@ -319,15 +368,15 @@ def test_mm_fp8_matches_scaled_mm():
     m, n, k = 512, 4096, 4096
     a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
     b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
-    sa = torch.tensor([2.5], device="cuda")
+    sa = _scale(a)
-    sb = torch.tensor([1.5], device="cuda")
+    sb = _scale(b)
-    a8, _ = quantize_bf16(a, sa, "e4m3")
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
-    b8, _ = quantize_bf16(b, sb, "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
-    out = mm_fp8(a8, b8, sa, sb)
+    out = mm_fp8(a8, b8, sa * sb)
     assert out.dtype == torch.bfloat16
     assert out.shape == (m, n)
 
-    ref = (a8.float().double() @ b8.float().double() * 2.5 * 1.5).to(torch.bfloat16)
+    ref = (a8.float().double() @ b8.float().double() * sa * sb).to(torch.bfloat16)
     torch.testing.assert_close(out, ref, atol=6.0, rtol=0.05)
 
     try:
@@ -342,30 +391,6 @@ def test_mm_fp8_matches_scaled_mm():
     )
 
 
-@skip_no_fp8
-def test_mm_fp8_fp8_output():
-    """mm_fp8 with out_dtype='e4m3' produces an FP8 output (layer-to-layer)."""
-    torch.manual_seed(12)
-    m, n, k = 256, 128, 64
-    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
-    sa = torch.tensor([2.0], device="cuda")
-    sb = torch.tensor([1.0], device="cuda")
-    os_ = torch.tensor([0.5], device="cuda")
-    a8, _ = quantize_bf16(a, sa, "e4m3")
-    b8, _ = quantize_bf16(b, sb, "e4m3")
-    out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_)
-    assert out8.dtype == torch.float8_e4m3fn
-    assert out8.shape == (m, n)
-
-    ref = (a8.float().double() @ b8.float().double() * 2.0 * 1.0 * 0.5).to(
-        torch.bfloat16
-    )
-    torch.testing.assert_close(
-        out8.float().to(torch.bfloat16), ref, atol=6.0, rtol=0.05
-    )
-
-
 # --------------------------------------------------------------------------
 # Policy-level (CPU-verifiable)
 # --------------------------------------------------------------------------
@@ -423,23 +448,26 @@ def test_fp8_tensor_meta_delayed_update():
     meta.w.seed(w, "e4m3")
     assert meta.w.initialized
     torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1))
-    # [hist | scale | counter] packing: views alias the single state buffer.
+    # [hist | scale] packing: views alias the single state buffer.
     assert meta.w.state.numel() == 4 + 2
     assert meta.w.hist.data_ptr() == meta.w.state.data_ptr()
     assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr()
-    # counter slot stays int32-zero (float bits) between launches
-    assert meta.w.state[-1].view(torch.int32).item() == 0
     meta.w.advance()
     assert meta.w.idx == 1
 
+    # update folds a fresh amax into the window and publishes the next scale
+    amax = torch.tensor([8.0])
+    meta.w.update(amax, "e4m3")
+    torch.testing.assert_close(meta.w.scale, torch.tensor([8.0 / 448.0]))
 
-def test_quantize_bf16_cpu_fallback():
+
+def test_quantize_cpu_fallback():
     """CPU fallback of the quantize primitive (scale semantics + amax)."""
     x = torch.randn(16, 32, dtype=torch.bfloat16)
-    scale = torch.tensor([0.5])
+    scale = torch.tensor([0.5])  # quantize multiplier
-    x8, amax = quantize_bf16(x, scale, "e4m3")
+    x8, amax = quantize(x, scale, "e4m3")
     assert x8.dtype == torch.float8_e4m3fn
-    ref = (x.float() / 0.5).to(torch.float8_e4m3fn)
+    ref = (x.float() * 0.5).to(torch.float8_e4m3fn)
     assert torch.equal(x8, ref)
     torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
 
@@ -447,26 +475,12 @@ def test_quantize_bf16_cpu_fallback():
 def test_mm_fp8_cpu_fallback():
     a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
     b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
-    sa = torch.tensor([2.0])
+    scale = torch.tensor([1.0])
-    sb = torch.tensor([0.5])
+    out = mm_fp8(a8, b8, scale)
-    out = mm_fp8(a8, b8, sa, sb)
+    ref = (a8.float() @ b8.float() * 1.0).to(torch.bfloat16)
-    ref = (a8.float() @ b8.float() * 2.0 * 0.5).to(torch.bfloat16)
     torch.testing.assert_close(out, ref)
 
 
-def test_mm_fp8_fp8_output_cpu():
-    """CPU fallback with an FP8 output (out_dtype='e4m3' + out_scale)."""
-    a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
-    b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn)
-    sa = torch.tensor([2.0])
-    sb = torch.tensor([0.5])
-    os_ = torch.tensor([0.25])
-    out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_)
-    assert out8.dtype == torch.float8_e4m3fn
-    ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn)
-    assert torch.equal(out8, ref)
-
-
 # --------------------------------------------------------------------------
 # torch-autocast parity: context semantics (nesting, thread locality, switch)
 # --------------------------------------------------------------------------
@@ -488,7 +502,7 @@ def test_nested_disabled_region_redispatches_bf16():
     resumes when it exits."""
     x, w = _linear()
     with fp8_autocast(enabled=True):
-        out_fp8 = F.linear(x, w)
+        F.linear(x, w)
         with fp8_autocast(enabled=False):
             out_bf16 = F.linear(x, w)
             assert type(out_bf16.grad_fn).__name__ != "_LinearFp8Backward"