ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: fp8 rings, lean autocast, gemm staging

Browse Source 
- Finalize scale rings inside the quantize kernels: a last-block epilogue (threadfence + counter elect) folds amax into hist, reduces the window and publishes the next scale on device, zero extra launches; _ScaleRing packs [hist | scale | counter] into one CUDA buffer.
- Split FP8QuantizeParams out of FP8Params so each operator owns its fields; linear_forward/backward_fp8 take optional ring arguments.
- Drop the inference weight-quantization cache; the optimizer bumps the weight version every step, so a cache would miss anyway.
- Zero amax scratch via empty + cudaMemsetAsync instead of torch::zeros, cutting a ~50us fill_ dispatch per quantize.
- Stage crosswise-B operands K-major with cp.async (contract >= 8192) and PRMT-transpose per k_seg region in smem, interleaved with the MMAs; the sync LDG + byte-scatter path it replaces was long-scoreboard bound (ncu 4.6 vs 0.4 stalls/issue).
- Load crosswise-A direct with an in-register PRMT transpose; its operands are typically L2-resident and the staging round trip measured as a net loss.
- Enable grouped rasterization for the congruous NT forward (shared B stripe keeps the weight operand hot in L2) and make the smem budget layout-aware (Fp8GemmSmem) while holding two CTAs per SM.
- Annotate ops/fp8.py return types; drop weight-cache and decorator tests, hoist their imports to module level.

e2e 12L/dim1024/B4xT512 fused AdamW: fp8 137.8ms/step vs bf16 210.3ms, 1.53x. Kernel vs cuBLASLt _scaled_mm: fwd 1.03-1.09x, dX 1.33-1.47x, dW 1.30-1.39x (from 1.10/1.42-1.49/1.52-1.56x), before the pre-transposed copies cuBLASLt needs for dX/dW. fp8 train step vs bf16: 1.34x at 2048 tokens (was 1.25x), 1.08x at 512.
This commit is contained in:
ViperEkura
2026-08-25 14:24:11 +08:00
parent 4dc5e923e0
commit 5e76fbd1bf
6 changed files with 823 additions and 397 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/fp8.py
+273 -181
View File
@@ -1,35 +1,37 @@
"""FP8 training: scaling recipes, per-tensor state, and aten::linear dispatch. """FP8 training: scaling recipes, per-tensor state, and aten::linear dispatch.
Layered (see also ``ops/fp8.py`` for the CUDA interface adapter): Layered (see ``ops/fp8.py`` for the CUDA interface adapter):
1. ``ops.fp8`` — the only module touching the pybind.
1. Kernel interface: ``ops.fp8`` — the only module touching the pybind. 2. This module (strategy layer): scaling *recipes* (TE-style delayed scaling
2. Training state (this module): scaling *recipes* (TE-style delayed scaling or dynamic current-amax scaling), per-tensor scales + amax history, and the
or dynamic current-amax scaling), per-tensor scales + amax history, and ``fp8_autocast`` context manager (like ``torch.autocast``).
the ``fp8_autocast`` context (like ``torch.autocast``). 3. aten::linear integration: registers the CUDA + AutogradCUDA impls.
3. aten::linear integration (this module): registers the CUDA impl and the
dtype guard.
Usage:: Usage::
from astrai.extension.fp8 import fp8_autocast from astrai.extension.fp8 import fp8_autocast
with fp8_autocast(enabled=True, fp8_format="hybrid"): with fp8_autocast(enabled=True, fp8_format="hybrid"):
logits = model(input_ids) logits = model(input_ids)
loss.backward() # fp8 backward runs wherever it is called: the loss.backward() # fp8 backward runs anywhere; fwd captured state on the node
# forward captures the fmt/recipe/meta on the autograd node
Importing this module registers the aten::linear CUDA and AutogradCUDA Format defaults follow the ecosystem consensus: E4M3 forward / E5M2 backward
implementations. ("hybrid"); every operand's scale is a quantization step derived from its amax
history by the active recipe.
Format defaults follow the ecosystem consensus: E4M3 for the forward pass, The context mirrors ``torch.autocast`` (``autocast_mode.py``): the active
E5M2 for the backward (gradient) pass ("hybrid"); every operand's scale is a ``(enabled, recipe, fp8_format)`` triple is thread-local (a ``contextvars``
quantization step derived from its amax history by the active recipe. ``ContextVar``, absent outside any region), and the manager is class-based and
reentrant with nested ``enabled=False`` disabling dispatch inside it. The module
targets *training*: every step quantizes x/w/g fresh (no weight-cast cache — the
optimizer bumps the weight version each step, so a torch-style cached_cast would
miss anyway), and the per-operand scales come from the delayed/dynamic recipe.
""" """
from contextlib import contextmanager import functools
from contextvars import ContextVar, Token
from dataclasses import dataclass from dataclasses import dataclass
from enum import Enum from enum import Enum
from typing import Optional from typing import Dict, List, Optional
import torch import torch
from torch.library import Library from torch.library import Library
@@ -58,62 +60,46 @@ class FP8Format(str, Enum):
class FP8Recipe: class FP8Recipe:
"""Scale-from-amax policy; the scale computation is the injection point. """Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
``scale_from_history`` receives the amax tensor for this operand (a ring ``scale_from_history`` receives the operand's amax tensor (a ring window for
window for delayed scaling, the current amax for dynamic scaling) and delayed scaling, the current amax for dynamic scaling) and returns the
returns the quantization step: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``. quantization step. Subclasses set ``history_len`` / ``margin``.
""" """
history_len: int = 16 history_len: int = 16
margin: int = 0 margin: int = 0
def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor: def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
raise NotImplementedError peak = amax.max()
return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
@dataclass @dataclass
class DelayedScaling(FP8Recipe): class DelayedScaling(FP8Recipe):
"""TE-style delayed scaling: max over the amax history window. """TE-style delayed scaling: max over the amax history window (amax from
*previous* steps; the window trades responsiveness against stability)."""
The scale is computed from amax measured in *previous* steps (delayed one
step); the window length trades responsiveness against stability.
"""
history_len: int = 16 history_len: int = 16
margin: int = 0 margin: int = 0
def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
peak = amax.max()
return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
@dataclass @dataclass
class DynamicScaling(FP8Recipe): class DynamicScaling(FP8Recipe):
"""Current-amax scaling (torchao DYNAMIC): measure, then quantize. """Current-amax scaling (torchao DYNAMIC): measure, then quantize. No
history — the scale is derived from the same-step amax, at an extra pass."""
No history — the scale is derived from the amax of the tensor being
quantized in the same step, at the cost of an extra reduction pass.
"""
history_len: int = 1 history_len: int = 1
margin: int = 0 margin: int = 0
def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
peak = amax.max()
return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
class _ScaleRing: class _ScaleRing:
"""One operand's delayed-scaling state, packed for in-kernel finalization. """One operand's delayed-scaling state: a float32 buffer
``[hist[n] | scale | counter]`` (views). The quantize kernel's last-finishing
``state`` is a single float32 CUDA buffer ``[hist[n] | scale | counter]`` block records the measured amax into ``hist[idx]``, reduces the window and
(``hist`` / ``scale`` are views). The quantize kernel's last-finishing publishes the next scale entirely on device — the Python-side write/max/write
block records the freshly measured amax into ``hist[idx]``, reduces the chain is gone. The counter slot stays int32-zero (float bits) between
window and publishes the next step's scale entirely on device — the launches; ``idx`` advances host-side each step.
Python-side hist-write / max / scale-write chain is gone. The counter
slot stays int32-zero (float bits) between launches. ``idx`` advances
host-side each step; ``margin`` is fixed by the recipe.
""" """
__slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized") __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -121,7 +107,6 @@ class _ScaleRing:
def __init__(self, device: torch.device, recipe: FP8Recipe): def __init__(self, device: torch.device, recipe: FP8Recipe):
self.recipe = recipe self.recipe = recipe
n = recipe.history_len n = recipe.history_len
# [hist | scale | counter]; the counter slot must start at int 0.
self.state = torch.zeros(n + 2, device=device, dtype=torch.float32) self.state = torch.zeros(n + 2, device=device, dtype=torch.float32)
self.hist = self.state[:n] self.hist = self.state[:n]
self.scale = self.state[n : n + 1] self.scale = self.state[n : n + 1]
@@ -140,12 +125,10 @@ class _ScaleRing:
class FP8TensorMeta: class FP8TensorMeta:
"""Per-weight delayed-scaling state: one ring per operand role. """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/
``g``). Fused kernels record amax while quantizing, so the scale used at step
Holds the ``w`` / ``x`` / ``g`` rings; fused kernels record the amax N reflects amax from steps < N. DynamicScaling never allocates a meta — it
while quantizing, so the scale used at step N reflects amax from steps measures the current amax inline.
< N. DynamicScaling never allocates a meta — it measures the current
amax inline (``_dynamic_scale``), so it needs no history storage.
""" """
__slots__ = ("w", "x", "g") __slots__ = ("w", "x", "g")
@@ -156,14 +139,70 @@ class FP8TensorMeta:
self.g = _ScaleRing(device, recipe) self.g = _ScaleRing(device, recipe)
@dataclass(frozen=True)
class _ActiveConfig:
"""The immutable (enabled, recipe, format) triple of one open region."""
enabled: bool
recipe: FP8Recipe
fp8_format: FP8Format
# Thread-local active configuration (torch's autocast TLS analog): set by
# fp8_autocast on __enter__, absent outside any region. Autograd engine
# threads run backwards with their own empty context — fine, since backward
# only reads state captured on ctx at forward time.
_active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
"astrai_fp8_active_config", default=None
)
class FP8State: class FP8State:
"""Global fp8 training state: active recipe + per-tensor metas.""" """Global fp8 training state: per-tensor metas + out-of-region defaults.
The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set
by ``fp8_autocast``. The properties below read that active config when a
region is open and the global defaults otherwise; the setters (and
``fp8_linear_enable``) write the global defaults — the persistent switch
applying outside any region. The metas registry is shared across threads
(GIL-protected); fp8 backward runs on autograd engine threads and only
touches metas captured on ``ctx`` at forward time.
"""
def __init__(self): def __init__(self):
self.enabled = False self.default_enabled = False
self.recipe: FP8Recipe = DelayedScaling() self.default_recipe: FP8Recipe = DelayedScaling()
self.fp8_format: FP8Format = FP8Format.HYBRID self.default_format: FP8Format = FP8Format.HYBRID
self._metas: dict[tuple, FP8TensorMeta] = {} self._metas: Dict[tuple, FP8TensorMeta] = {}
# Active-config views (region config if open, else the defaults).
@property
def enabled(self) -> bool:
cfg = _active_config.get()
return cfg.enabled if cfg is not None else self.default_enabled
@property
def recipe(self) -> FP8Recipe:
cfg = _active_config.get()
return cfg.recipe if cfg is not None else self.default_recipe
@property
def fp8_format(self) -> FP8Format:
cfg = _active_config.get()
return cfg.fp8_format if cfg is not None else self.default_format
# Persistent (out-of-region) defaults.
@enabled.setter
def enabled(self, value: bool) -> None:
self.default_enabled = bool(value)
@recipe.setter
def recipe(self, value: FP8Recipe) -> None:
self.default_recipe = value
@fp8_format.setter
def fp8_format(self, value: FP8Format) -> None:
self.default_format = FP8Format(value)
def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta: def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
key = (w.data_ptr(), w.shape, w.dtype) key = (w.data_ptr(), w.shape, w.dtype)
@@ -174,13 +213,11 @@ class FP8State:
return meta return meta
def reset(self) -> None: def reset(self) -> None:
self.enabled = False self.default_enabled = False
self._metas.clear() self._metas.clear()
# Global singleton: autograd backward runs on the engine worker threads, so # Process-wide singleton; per-thread/per-region state lives in _active_config.
# thread-local state would lose the fp8 flag during loss.backward(). The GIL
# protects Python-side mutation; the CUDA kernels take their own mutex.
_state = FP8State() _state = FP8State()
@@ -188,43 +225,76 @@ def fp8_state() -> FP8State:
return _state return _state
@contextmanager def _active() -> Optional[_ActiveConfig]:
def fp8_autocast( """The active config when fp8 dispatch is on, else ``None`` (fast guard).
enabled: bool = True,
update_interval: int = 16, A region config wins (honoring nested ``enabled=False`` regions); with no
recipe: Optional[FP8Recipe] = None, region open this falls back to the persistent global switch
fp8_format: str = "hybrid", (``fp8_linear_enable``), so that flag still routes aten::linear to fp8.
margin: int = 0, """
): cfg = _active_config.get()
if cfg is not None:
return cfg if cfg.enabled else None
if _state.default_enabled:
return _ActiveConfig(True, _state.default_recipe, _state.default_format)
return None
def _current_config() -> _ActiveConfig:
"""Like ``_active()`` but always returns a config (disabled regions and
out-of-region direct calls resolve to the global defaults)."""
cfg = _active_config.get()
if cfg is not None:
return cfg
return _ActiveConfig(
_state.default_enabled, _state.default_recipe, _state.default_format
)
class fp8_autocast:
"""Autocast-style context: fp8 linear dispatch on this thread. """Autocast-style context: fp8 linear dispatch on this thread.
Usage:: Mirrors ``torch.autocast`` — a class-based, reentrant, nestable context
over thread-local state::
with fp8_autocast(enabled=True, fp8_format="hybrid"): with fp8_autocast(enabled=True, fp8_format="hybrid"):
logits = model(input_ids) # aten::linear -> fp8 path logits = model(input_ids) # aten::linear -> fp8 path
loss.backward() # fp8 backward; state was captured at forward time loss.backward() # fp8 backward; state was captured at forward time
Args: Nesting follows torch: each ``__enter__`` pushes the new active config, each
enabled: toggle fp8 dispatch for aten::linear. ``__exit__`` restores the previous one, and a nested ``enabled=False`` region
update_interval: legacy alias for the delayed-scaling history window simply disables dispatch inside it. The instance doubles as a decorator.
(used only when ``recipe`` is not given).
recipe: scaling policy; defaults to ``DelayedScaling(update_interval)``.
fp8_format: ``"e4m3"`` / ``"e5m2"`` / ``"hybrid"`` (default) — hybrid
means E4M3 forward, E5M2 backward.
margin: scale headroom (``scale = (amax / FP8_MAX) / 2^margin``) used
with the default delayed recipe.
""" """
state = fp8_state()
prev = (state.enabled, state.recipe, state.fp8_format) def __init__(
if recipe is None: self,
recipe = DelayedScaling(history_len=update_interval, margin=margin) enabled: bool = True,
state.enabled = enabled update_interval: int = 16,
state.recipe = recipe recipe: Optional[FP8Recipe] = None,
state.fp8_format = FP8Format(fp8_format) fp8_format: str = "hybrid",
try: margin: int = 0,
yield ):
finally: if recipe is None:
state.enabled, state.recipe, state.fp8_format = prev recipe = DelayedScaling(history_len=update_interval, margin=margin)
self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
self._tokens: List[Token] = []
def __enter__(self) -> "fp8_autocast":
self._tokens.append(_active_config.set(self._config))
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
token = self._tokens.pop()
_active_config.reset(token)
return False
def __call__(self, func):
@functools.wraps(func)
def decorate(*args, **kwargs):
with self:
return func(*args, **kwargs)
return decorate
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -237,72 +307,96 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor
return recipe.scale_from_history(amax, fmt) return recipe.scale_from_history(amax, fmt)
def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): _zero_bias: Dict[Optional[int], torch.Tensor] = {}
def _empty_bias(x: torch.Tensor) -> torch.Tensor:
"""Per-device cached 0-element bf16 bias (the binding only checks numel —
never mutated), saving a CUDA allocation per bias-less linear."""
key = x.device.index
t = _zero_bias.get(key)
if t is None:
t = torch.empty(0, device=x.device, dtype=torch.bfloat16)
_zero_bias[key] = t
return t
def fp8_linear_forward(
x: torch.Tensor, w: torch.Tensor, bias=None, cfg: Optional[_ActiveConfig] = None
):
"""Scaled fp8 linear forward (called from the aten::linear impl). """Scaled fp8 linear forward (called from the aten::linear impl).
Pure FP8 path for both recipes: quantize x/w with the active scales, run Pure FP8 path for both recipes: quantize x/w with the active scales, run the
the pre-quantized GEMM. With delayed scaling the rings finalize inside pre-quantized GEMM. Delayed scaling finalizes the rings inside the quantize
the quantize kernels (amax folded into the window, next step's scale kernels (amax folded into the window, next scale published on device);
published on device); dynamic scaling measures the current amax itself. dynamic scaling measures the current amax itself. Training quantizes the
weight every step (the optimizer bumps its version, so there is no cast
cache, matching ``cached_cast``-less behavior).
""" """
if bias is None:
bias = torch.empty(0, device=x.device, dtype=x.dtype)
state = fp8_state() state = fp8_state()
fmt = state.fp8_format.fwd() if cfg is None:
if isinstance(state.recipe, DynamicScaling): cfg = _current_config()
meta = None fmt = cfg.fp8_format.fwd()
sx = _dynamic_scale(x.reshape(-1, w.size(1)), state.recipe, fmt) margin = cfg.recipe.margin
sw = _dynamic_scale(w, state.recipe, fmt) if bias is None:
out, amax_x, amax_w = linear_forward_fp8(x, w, bias, sx, sw, fmt) bias = _empty_bias(x)
if isinstance(cfg.recipe, DynamicScaling): # measure-then-quantize, no state
sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
sw = _dynamic_scale(w, cfg.recipe, fmt)
out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt)
return out
meta = state.get_weight_meta(w)
if not meta.w.initialized:
meta.w.seed(w, fmt)
if not meta.x.initialized:
meta.x.seed(x, fmt)
# The kernel finalizes each ring in-kernel (overwriting the scale slot), so
# the w/x scales are taken from the ring before the quantize.
if w.dtype is not torch.bfloat16: # static pre-quantized weight
w_arg, sw_arg, w_ring = w, meta.w.scale, None
else: else:
meta = state.get_weight_meta(w) w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state
if not meta.w.initialized: out, _x8, _w8, _ax, _aw = linear_forward_fp8(
meta.w.seed(w, fmt) x,
if not meta.x.initialized: w_arg,
meta.x.seed(x, fmt) bias,
# In-kernel ring finalization: the kernels write hist[idx] and the meta.x.scale,
# next scale; idx rotates host-side (the device counter self-rearms). sw_arg,
w_is_fp8 = w.dtype != torch.bfloat16 fmt,
out, amax_x, amax_w = linear_forward_fp8( None,
x, meta.x.state,
w, meta.x.idx,
bias, margin,
meta.x.scale, w_ring,
meta.w.scale, meta.w.idx,
fmt, margin,
None, )
meta.x.state, meta.x.advance()
meta.x.idx, if w_ring is not None:
state.recipe.margin, meta.w.advance()
None if w_is_fp8 else meta.w.state,
meta.w.idx,
state.recipe.margin,
)
meta.x.advance()
if not w_is_fp8:
meta.w.advance()
return out return out
class _LinearFp8(torch.autograd.Function): class _LinearFp8(torch.autograd.Function):
"""The fp8 linear forward/backward pair (standard Function style). """The fp8 linear forward/backward pair (standard Function style).
The forward runs inside the ``fp8_autocast`` region and captures the The forward runs inside ``fp8_autocast`` and captures the active
active fmt/recipe/meta on ``ctx``; the backward reads only the captured fmt/recipe/meta on ``ctx``; the backward reads only that captured state, so
state, so ``loss.backward()`` may run after the context exits. The ``loss.backward()`` may run after the context exits. The gradient is
gradient is quantized once (E5M2 in hybrid mode) and the dX / dW GEMMs quantized once (E5M2 in hybrid) and both dX/dW GEMMs share it; the output
share that quantization; output masks come from ``needs_input_grad``. masks come from ``needs_input_grad``.
""" """
@staticmethod @staticmethod
def forward(ctx, x, w, bias): def forward(ctx, x, w, bias):
out = fp8_linear_forward(x, w, bias) cfg = _current_config()
state = fp8_state() out = fp8_linear_forward(x, w, bias, cfg)
ctx.save_for_backward(x, w) ctx.save_for_backward(x, w)
ctx.fmt_bwd = state.fp8_format.bwd() ctx.fmt_bwd = cfg.fp8_format.bwd()
ctx.recipe = state.recipe ctx.recipe = cfg.recipe
ctx.is_dynamic = isinstance(state.recipe, DynamicScaling) ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
ctx.meta = None if ctx.is_dynamic else state.get_weight_meta(w) ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w)
return out return out
@staticmethod @staticmethod
@@ -310,32 +404,33 @@ class _LinearFp8(torch.autograd.Function):
def backward(ctx, g): def backward(ctx, g):
x, w = ctx.saved_tensors x, w = ctx.saved_tensors
fmt = ctx.fmt_bwd fmt = ctx.fmt_bwd
# Per-recipe scale/ring selection; both branches share one call below.
if ctx.is_dynamic: if ctx.is_dynamic:
sg = _dynamic_scale(g, ctx.recipe, fmt) sg = _dynamic_scale(g, ctx.recipe, fmt)
sw = _dynamic_scale(w, ctx.recipe, fmt) sw = _dynamic_scale(w, ctx.recipe, fmt)
sx = _dynamic_scale(x, ctx.recipe, fmt) sx = _dynamic_scale(x, ctx.recipe, fmt)
grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( ring, idx = None, 0
g, x, w, list(ctx.needs_input_grad), sg, sw, sx, fmt
)
else: else:
meta = ctx.meta meta = ctx.meta
if not meta.g.initialized: if not meta.g.initialized:
meta.g.seed(g, fmt) meta.g.seed(g, fmt)
# The g quantize kernel finalizes the gradient's ring in-kernel. sg, ring, idx = meta.g.scale, meta.g.state, meta.g.idx
grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( sw, sx = meta.w.scale, meta.x.scale
g, grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8(
x, g,
w, x,
list(ctx.needs_input_grad), w,
meta.g.scale, list(ctx.needs_input_grad),
meta.w.scale, sg,
meta.x.scale, sw,
fmt, sx,
meta.g.state, fmt,
meta.g.idx, ring,
ctx.recipe.margin, idx,
) ctx.recipe.margin,
meta.g.advance() )
if not ctx.is_dynamic:
meta.g.advance() # the g quantize kernel finalized the ring in-kernel
return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
@@ -345,31 +440,29 @@ class _LinearFp8(torch.autograd.Function):
def fp8_linear_enable(enabled: bool = True) -> None: def fp8_linear_enable(enabled: bool = True) -> None:
"""Toggle fp8 dispatch for aten::linear (global; backward runs on engine """Toggle fp8 dispatch for aten::linear globally (the out-of-region default;
worker threads, so a thread-local flag would be lost during backward).""" ``fp8_autocast`` regions override it thread-locally)."""
fp8_state().enabled = enabled fp8_state().default_enabled = enabled
def fp8_linear_enabled() -> bool: def fp8_linear_enabled() -> bool:
return fp8_state().enabled """Whether fp8 dispatch is active right now (region config or global)."""
return _active() is not None
def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool: def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool:
"""Shape guard for the fp8 linear path. """Shape guard for the fp8 path. Unlike a strict 16-alignment requirement,
the kernels handle unaligned M/N via boundary checks (slower but correct) —
Unlike a strict 16-alignment requirement, the fp8 kernels handle unaligned so no whole-call bf16 fallback for small decode batches. Only the K-dimension
M/N via boundary checks (slower but correct) — so no whole-call bf16 contraction must match and the weight must be 2D."""
fallback for small decode batches. Only the K-dimension contraction must
match, and the weight must be 2D.
"""
return x.dim() >= 2 and w.dim() == 2 and x.size(-1) == w.size(1) return x.dim() >= 2 and w.dim() == 2 and x.size(-1) == w.size(1)
def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
if ( if (
fp8_linear_enabled() _active() is not None
and x.dtype == torch.bfloat16 and x.dtype is torch.bfloat16
and w.dtype == torch.bfloat16 and w.dtype is torch.bfloat16
and _fp8_supported(x, w) and _fp8_supported(x, w)
): ):
return _LinearFp8.apply(x, w, bias) return _LinearFp8.apply(x, w, bias)
@@ -384,9 +477,8 @@ def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
_lib = Library("aten", "IMPL", "CUDA") _lib = Library("aten", "IMPL", "CUDA")
_lib.impl("linear", _linear_cuda_impl) _lib.impl("linear", _linear_cuda_impl)
# Also replace torch's generated linear autograd formula (which would call # Also replace torch's generated linear autograd formula (which would call
# aten::linear_backward after the fp8_autocast region exits). The fp8 # aten::linear_backward after the fp8_autocast region exits). The fp8 backward
# backward is owned by _LinearFp8 with its state captured at forward time, # is owned by _LinearFp8 with state captured at forward time, so loss.backward()
# so loss.backward() works wherever it is called; the same CUDA registration # works wherever it is called; the CUDA registration still covers inference_mode.
# still covers inference_mode, where autograd keys are skipped entirely.
_lib_autograd = Library("aten", "IMPL", "AutogradCUDA") _lib_autograd = Library("aten", "IMPL", "AutogradCUDA")
_lib_autograd.impl("linear", _linear_cuda_impl) _lib_autograd.impl("linear", _linear_cuda_impl)
astrai/extension/ops/fp8.py
+33 -25
View File
@@ -4,7 +4,7 @@ Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
- ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → FP8 with fused amax - ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → FP8 with fused amax
- ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output) - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
- ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, amax_x, amax_w)`` - ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, x8, w8, amax_x, amax_w)``
- ``linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt) -> (gx, gw, gb, amax_g)`` - ``linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt) -> (gx, gw, gb, amax_g)``
Scale semantics: scales are *quantization steps* — the value divided out when Scale semantics: scales are *quantization steps* — the value divided out when
@@ -16,6 +16,8 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
this module is stateless. this module is stateless.
""" """
from typing import List, Optional, Tuple
import torch import torch
from torch.library import custom_op from torch.library import custom_op
@@ -39,7 +41,7 @@ def _fmt_dtype(fmt: str) -> torch.dtype:
@custom_op("custom::fp8_quantize", mutates_args=()) @custom_op("custom::fp8_quantize", mutates_args=())
def fp8_quantize( def fp8_quantize(
x: torch.Tensor, scale: torch.Tensor, fmt: int x: torch.Tensor, scale: torch.Tensor, fmt: int
) -> tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
"""BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``.""" """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``."""
@@ -73,7 +75,7 @@ def fp8_gemm(
sa: torch.Tensor, sa: torch.Tensor,
sb: torch.Tensor, sb: torch.Tensor,
out_dtype: int = 0, out_dtype: int = 0,
out_scale: torch.Tensor | None = None, out_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor: ) -> torch.Tensor:
"""FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation. """FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation.
@@ -106,7 +108,9 @@ def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None):
return acc.to(torch.bfloat16) return acc.to(torch.bfloat16)
def quantize_bf16(x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"): def quantize_bf16(
x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"
) -> Tuple[torch.Tensor, torch.Tensor]:
"""BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``. """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``.
``scale`` is the quantization step (device scalar); ``fmt`` selects ``scale`` is the quantization step (device scalar); ``fmt`` selects
@@ -122,7 +126,7 @@ def mm_fp8(
sa: torch.Tensor, sa: torch.Tensor,
sb: torch.Tensor, sb: torch.Tensor,
out_dtype: str = "bf16", out_dtype: str = "bf16",
out_scale: torch.Tensor | None = None, out_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor: ) -> torch.Tensor:
"""Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``. """Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``.
@@ -139,24 +143,28 @@ def mm_fp8(
def linear_forward_fp8( def linear_forward_fp8(
x, x: torch.Tensor,
w, w: torch.Tensor,
bias, bias: Optional[torch.Tensor],
sx, sx: torch.Tensor,
sw, sw: torch.Tensor,
fmt: str = "e4m3", fmt: str = "e4m3",
bias_scale=None, bias_scale: Optional[torch.Tensor] = None,
x_ring=None, x_ring: Optional[torch.Tensor] = None,
x_ring_idx: int = 0, x_ring_idx: int = 0,
x_ring_margin: int = 0, x_ring_margin: int = 0,
w_ring=None, w_ring: Optional[torch.Tensor] = None,
w_ring_idx: int = 0, w_ring_idx: int = 0,
w_ring_margin: 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. """Pure FP8 linear forward: quantize x/w to ``fmt``, pre-quantized GEMM.
Returns ``(out, amax_x, amax_w)``. ``bias`` may be ``None``. For static Returns ``(out, x8, w8, amax_x, amax_w)`` — the quantized operands are
fp8 inference, ``w`` and ``bias`` may arrive pre-quantized to ``fmt`` 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`` / (produced by :func:`quantize_bf16` with their scales as ``sw`` /
``bias_scale``); a pre-quantized ``bias`` requires ``bias_scale``, and ``bias_scale``); a pre-quantized ``bias`` requires ``bias_scale``, and
its ``amax_w`` comes back 0. The bias is fused into the GEMM epilogue. its ``amax_w`` comes back 0. The bias is fused into the GEMM epilogue.
@@ -190,18 +198,18 @@ def linear_forward_fp8(
def linear_backward_fp8( def linear_backward_fp8(
g, g: torch.Tensor,
x, x: torch.Tensor,
w, w: torch.Tensor,
masks, masks: List[bool],
sg, sg: torch.Tensor,
sw, sw: torch.Tensor,
sx, sx: torch.Tensor,
fmt: str = "e5m2", fmt: str = "e5m2",
g_ring=None, g_ring: Optional[torch.Tensor] = None,
g_ring_idx: int = 0, g_ring_idx: int = 0,
g_ring_margin: 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)``. """FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``.
The gradient (and the transposed w/x operands) are quantized to ``fmt`` The gradient (and the transposed w/x operands) are quantized to ``fmt``
csrc/kernels/fp8/common.h
+6
View File
@@ -61,6 +61,12 @@ struct Fp8GemmTraits {
static constexpr __nv_fp8_interpretation_t kNvFormat = static constexpr __nv_fp8_interpretation_t kNvFormat =
kIsE5M2 ? __NV_E5M2 : __NV_E4M3; kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f; static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
// Derived launch geometry: 64x32 warp tiles give the CTA thread count.
// The shared-memory budget is layout-aware (crosswise operands add K-
// major staging + a canonical buffer), so it lives in Fp8GemmSmem in
// gemm.cuh together with the resident-CTA hint for __launch_bounds__.
static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
}; };
// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional // Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional
csrc/kernels/fp8/gemm.cuh
+399 -166
View File
@@ -191,20 +191,11 @@ __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) {
((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + (col & 15)); ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + (col & 15));
} }
// Stage-load one GEMM operand into the canonical flat [rows * K] shared tile // Stage-load a CONGRUOUS operand (stored [rows][contract], contract-
// (addressing via tile_at, so stores land in the swizzled layout). The // contiguous — the only cp.async-able shape for the canonical tile) into the
// transpose is folded into the staging step via a CUTLASS-style crosswise // flat [rows * K] shared tile via tile_at's swizzle. Crosswise operands go
// layout: RowMajor (stored [rows][contract]) copies 16-byte K-contiguous runs // through stage_crosswise_tile + transpose_crosswise_tile instead.
// with cp.async, while ColMajor (stored [contract][rows]) reads 16-byte runs template <typename T8, int K, int RowsTile, int kThreads>
// along the operand's contiguous non-contract dim and scatters them across
// the tile's rows. Crosswise runs cannot use cp.async (the 16 destination
// bytes land on 16 different rows), so their global loads are plain LDGs —
// issued as one batch per row group before the first scatter so their
// latencies overlap instead of serializing behind the shared stores.
// RowsTile is the tile's row capacity (kBlockM / kBlockN) and kThreads the
// CTA size; the runtime `rows` bound may be smaller (tail predication).
// `block_row` is this block's origin in the operand's row dim.
template <typename T8, int K, typename Layout, int RowsTile, int kThreads>
__device__ __forceinline__ void __device__ __forceinline__ void
load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows, load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
int64_t contract, int64_t ld, int tid, int64_t k_base, int64_t contract, int64_t ld, int tid, int64_t k_base,
@@ -213,101 +204,29 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
static_assert(RowsTile * kChunks % kThreads == 0, static_assert(RowsTile * kChunks % kThreads == 0,
"tile chunks must divide evenly across threads"); "tile chunks must divide evenly across threads");
constexpr int kCpt = RowsTile * kChunks / kThreads; // chunks per thread constexpr int kCpt = RowsTile * kChunks / kThreads; // chunks per thread
if constexpr (std::is_same_v<Layout, ColMajor>) { // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of
// Operand stored [contract][rows]: contiguous along the non-contract // one row (K=64: a contiguous 32B pair; K=32: a single chunk).
// dim. Each thread scatters one 16-byte run per K/32 pass; when the constexpr int kCpr = kChunks / kCpt; // chunks per row slice
// tile has more 16-row groups than warps (RowsTile > kThreads/2), const int r = tid / kCpr;
// each thread covers several groups. const int c0 = (tid % kCpr) * kCpt * 16;
constexpr int kWarpsTile = kThreads / 32; const int64_t row = block_row + r;
constexpr int kGroups = RowsTile / 16; const bool row_ok = row < rows;
constexpr int kPasses = K / 32; // k_base and every c are multiples of 16, so the per-chunk sources
static_assert(kGroups % kWarpsTile == 0, // share the row base's alignment.
"row groups must divide evenly across warps"); const auto* src = operand + row * ld + k_base;
const int kl = tid & 31; // byte column within a 32B pass const bool chunk_aligned = (reinterpret_cast<uintptr_t>(src) & 15) == 0;
// r0 is always a multiple of 16 (block_row is a multiple of RowsTile
// and each group covers 16 rows), so every run shares the base+ld
// alignment: one uniform check instead of one per pass.
const bool run_aligned =
((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
#pragma unroll #pragma unroll
for (int g = 0; g < kGroups / kWarpsTile; ++g) { for (int j = 0; j < kCpt; ++j) {
const int rg = (tid >> 5) + g * kWarpsTile; const int c = c0 + j * 16;
const int64_t r0 = block_row + rg * 16; T8* dst = tile_at<K>(tile, r, c);
const bool rows_full = r0 + 15 < rows; // pass-invariant if (row_ok && chunk_aligned && k_base + c + 15 < contract) {
// Batch every 16B run load of this row group before the first astrai::cp_async_16(dst, src + c, true);
// scatter: the LDGs are independent, and the byte-granular } else {
// shared stores would otherwise serialize behind each one. // Tail chunk (or misaligned base): predicated scalar fill.
uint4 v[kPasses];
bool fast[kPasses];
#pragma unroll #pragma unroll
for (int pass = 0; pass < kPasses; ++pass) { for (int i = 0; i < 16; ++i)
const int64_t k_idx = k_base + kl + pass * 32; dst[i] =
fast[pass] = rows_full && run_aligned && k_idx < contract; row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f);
if (fast[pass])
v[pass] =
*reinterpret_cast<const uint4*>(operand + k_idx * ld + r0);
}
#pragma unroll
for (int pass = 0; pass < kPasses; ++pass) {
const int col = kl + pass * 32;
if (fast[pass]) {
const auto* bytes = reinterpret_cast<const T8*>(&v[pass]);
// Scatter 16 bytes along the tile rows through tile_at's
// swizzle. Rows sharing a physical chunk form groups of
// (8 / kChunks) consecutive rows (see tile_at), so each
// group is one tile_at address plus a K-byte row stride.
constexpr int kGrp = 8 / kChunks;
#pragma unroll
for (int j = 0; j < 16 / kGrp; ++j) {
T8* p = tile_at<K>(tile, rg * 16 + j * kGrp, col);
#pragma unroll
for (int i = 0; i < kGrp; ++i)
p[i * K] = bytes[j * kGrp + i];
}
} else if (k_base + col < contract) {
// Row-tail or misaligned run: byte-granular gather with
// per-row predication (the k column itself is in range).
#pragma unroll
for (int i = 0; i < 16; ++i) {
const int64_t r_idx = r0 + i;
*tile_at<K>(tile, rg * 16 + i, col) =
r_idx < rows ? operand[(k_base + col) * ld + r_idx]
: T8(0.0f);
}
} else {
// Contract tail: straight zero-fill, no global traffic.
#pragma unroll
for (int i = 0; i < 16; ++i)
*tile_at<K>(tile, rg * 16 + i, col) = T8(0.0f);
}
}
}
} else {
// Operand stored [rows][contract]: contiguous along the contract dim.
// Linear chunk mapping: thread covers kCpt consecutive 16B chunks of
// one row (K=64: a contiguous 32B pair; K=32: a single chunk).
constexpr int kCpr = kChunks / kCpt; // chunks per row slice
const int r = tid / kCpr;
const int c0 = (tid % kCpr) * kCpt * 16;
const int64_t row = block_row + r;
const bool row_ok = row < rows;
// k_base and every c are multiples of 16, so the per-chunk sources
// share the row base's alignment.
const auto* src = operand + row * ld + k_base;
const bool chunk_aligned = (reinterpret_cast<uintptr_t>(src) & 15) == 0;
#pragma unroll
for (int j = 0; j < kCpt; ++j) {
const int c = c0 + j * 16;
T8* dst = tile_at<K>(tile, r, c);
if (row_ok && chunk_aligned && k_base + c + 15 < contract) {
astrai::cp_async_16(dst, src + c, true);
} else {
// Tail chunk (or misaligned base): predicated scalar fill.
#pragma unroll
for (int i = 0; i < 16; ++i)
dst[i] =
row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f);
}
} }
} }
} }
@@ -320,15 +239,195 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// Swizzled 16B-chunk address (tile_at's layout) as a raw shared-memory // Swizzled 16B-chunk address (tile_at's layout) as a raw shared-memory
// pointer for ldmatrix. Valid for kK in {32, 64} (the swizzle itself lives // pointer for ldmatrix. Valid for kK in {32, 64, 128} (the swizzle itself
// only in tile_at; this wrapper just converts the element address). // lives only in tile_at; this wrapper just converts the element address).
template <typename T8, int kK> template <typename T8, int kK>
__device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk) { __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk) {
static_assert(kK == 32 || kK == 64, static_assert(kK == 32 || kK == 64 || kK == 128,
"fragment swizzle offsets assume kK in {32, 64}"); "fragment swizzle offsets assume kK in {32, 64, 128}");
return __cvta_generic_to_shared(tile_at<kK>(tile, row, chunk << 4)); return __cvta_generic_to_shared(tile_at<kK>(tile, row, chunk << 4));
} }
// Crosswise operands (stored [contract][rows], rows-contiguous) cannot be
// cp.async'd into the canonical [rows][contract] tile — a 16B global run
// holds one contract byte for each of 16 rows. They stage K-major instead
// (byte (p, r) at p*RowsTile + r), where the very same runs land contiguously
// and cp.async applies unchanged; a per-tile smem->smem transpose (below)
// then produces the canonical swizzled tile the MMA fragments read. This
// keeps the whole global→shared path asynchronous — the synchronous
// LDG+byte-scatter staging this replaces left the kernel long-scoreboard
// bound (ncu: 4.6 stalled loads per issue vs 0.4 on the congruous path).
template <typename T8, int K, int RowsTile, int kThreads>
__device__ __forceinline__ void
stage_crosswise_tile(T8* staging, const T8* __restrict__ operand, int64_t rows,
int64_t contract, int64_t ld, int tid, int64_t k_base,
int64_t block_row) {
constexpr int kRuns = K * RowsTile / 16; // 16B runs per tile
// r0 is a multiple of 16 and p*ld keeps 16B alignment whenever ld has it,
// so one uniform verdict covers every run.
const bool run_aligned =
((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
for (int run = tid; run < kRuns; run += kThreads) {
const int pl = run % K; // local contract byte (column of the run)
const int rg = run / K; // 16-row group
const int64_t r0 = block_row + (int64_t)rg * 16;
T8* dst = staging + pl * RowsTile + rg * 16;
if (run_aligned && r0 + 15 < rows && k_base + pl < contract)
astrai::cp_async_16(dst, operand + (k_base + pl) * ld + r0, true);
else {
// Row tail, contract tail or misaligned base: predicated fill.
#pragma unroll
for (int i = 0; i < 16; ++i) {
const int64_t r = r0 + i;
dst[i] = r < rows && k_base + pl < contract
? operand[(k_base + pl) * ld + r]
: T8(0.0f);
}
}
}
}
// Direct (synchronous) crosswise load into a canonical rotating stage:
// LDG.128 x4 (4 consecutive contract bytes x 16 rows) + in-register PRMT
// transpose + 16 STS.32. Used for crosswise operands whose global data is
// typically L2-resident (the A side of dW): the staging detour's extra
// shared-memory round trip costs more than the latency it hides there,
// while crosswise B operands (DRAM-streamed weights of dX) take the
// asynchronous stage_crosswise_tile path instead.
template <typename T8, int K, int RowsTile, int kThreads>
__device__ __forceinline__ void
load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows,
int64_t contract, int64_t ld, int tid, int64_t k_base,
int64_t block_row) {
constexpr int kQuads = K / 4; // 4-byte contract quads per tile
constexpr int kGroups = RowsTile / 16;
constexpr int kTChunks = kQuads * kGroups; // 64B chunks per tile
// r0 is always a multiple of 16 (block_row is a multiple of RowsTile and
// each group covers 16 rows), and p*ld keeps the base 16B-aligned
// whenever ld is, so every run of a chunk shares one alignment verdict.
const bool run_aligned =
((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
for (int chunk = tid; chunk < kTChunks; chunk += kThreads) {
const int quad = chunk / kGroups;
const int rg = chunk % kGroups;
const int64_t r0 = block_row + rg * 16;
const bool rows_full = r0 + 15 < rows;
if (rows_full && run_aligned) {
const int64_t p0 = k_base + quad * 4;
uint4 v[4];
#pragma unroll
for (int s = 0; s < 4; ++s) {
// Contract tail: a run past k carries zero bytes; they flow
// through the PRMT transpose like any other value.
if (p0 + s < contract)
v[s] = *reinterpret_cast<const uint4*>(
operand + (p0 + s) * ld + r0);
else
v[s] = make_uint4(0u, 0u, 0u, 0u);
}
const unsigned* bytes = reinterpret_cast<const unsigned*>(v);
#pragma unroll
for (int i = 0; i < 16; ++i) {
// word i = row r0+i's quad: byte i of each of the four runs
// [v0.b(i), v1.b(i), v2.b(i), v3.b(i)]. Byte i of a uint4
// lives in its (i>>2)-th 32-bit register.
const unsigned nib = i & 3;
const unsigned sel = nib | ((nib + 4) << 4);
const unsigned w01 =
__byte_perm(bytes[0 + (i >> 2)], bytes[4 + (i >> 2)], sel);
const unsigned w23 =
__byte_perm(bytes[8 + (i >> 2)], bytes[12 + (i >> 2)], sel);
*reinterpret_cast<unsigned*>(tile_at<K>(tile, rg * 16 + i,
quad * 4)) =
__byte_perm(w01, w23, 0x5410u);
}
} else {
// Row-tail or misaligned chunk: byte-granular gather with
// per-row predication; contract-tail columns zero-fill.
#pragma unroll
for (int s = 0; s < 4; ++s) {
const int col = quad * 4 + s;
if (k_base + col >= contract) {
#pragma unroll
for (int i = 0; i < 16; ++i)
*tile_at<K>(tile, rg * 16 + i, col) = T8(0.0f);
continue;
}
#pragma unroll
for (int i = 0; i < 16; ++i) {
const int64_t r_idx = r0 + i;
*tile_at<K>(tile, rg * 16 + i, col) =
r_idx < rows
? operand[(k_base + col) * ld + r_idx]
: T8(0.0f);
}
}
}
}
}
// K-major staging -> canonical [rows][kK] swizzled tile, one chunk at a time.
// Each chunk (indexed within a k_seg region of `quads_per_seg` quads) covers
// 4 consecutive contract bytes x 16 rows: four LDS.128 grab the staging runs,
// PRMT byte selects transpose them in registers, and sixteen STS.32 land the
// row quads through tile_at's swizzle — 4x fewer store instructions than a
// byte-granular scatter. Chunk-at-a-time lets the caller pool work across
// operands; the region restriction lets the main loop overlap one region's
// transpose with another region's MMAs (a whole-tile serial transpose put
// the crosswise GEMMs at 25% tensor utilization).
template <typename T8, int K, int RowsTile>
__device__ __forceinline__ void
transpose_crosswise_region(T8* tile, const T8* staging, int idx, int quad0) {
constexpr int kGroups = RowsTile / 16;
const int quad = quad0 + idx / kGroups;
const int rg = idx % kGroups;
// The four runs sit RowsTile bytes apart (one per contract byte of the
// quad); each run is 16 contiguous staging bytes = 16 rows.
const char* run0 = reinterpret_cast<const char*>(
staging + quad * 4 * RowsTile + rg * 16);
uint4 v[4];
#pragma unroll
for (int s = 0; s < 4; ++s)
v[s] = *reinterpret_cast<const uint4*>(run0 + s * RowsTile);
const unsigned* bytes = reinterpret_cast<const unsigned*>(v);
#pragma unroll
for (int i = 0; i < 16; ++i) {
// word i = row r0+i's quad: byte i of each of the four runs
// [v0.b(i), v1.b(i), v2.b(i), v3.b(i)]. Byte i of a uint4 lives in
// its (i>>2)-th 32-bit register.
const unsigned nib = i & 3;
const unsigned sel = nib | ((nib + 4) << 4);
const unsigned w01 =
__byte_perm(bytes[0 + (i >> 2)], bytes[4 + (i >> 2)], sel);
const unsigned w23 =
__byte_perm(bytes[8 + (i >> 2)], bytes[12 + (i >> 2)], sel);
*reinterpret_cast<unsigned*>(tile_at<K>(tile, rg * 16 + i, quad * 4)) =
__byte_perm(w01, w23, 0x5410u);
}
}
// Layout-aware shared-memory budget and occupancy hint. A congruous or
// direct-crosswise operand needs its kStages rotating canonical buffers; a
// staged-crosswise operand (crosswise B with kBStaged) needs kStages K-major
// staging buffers plus ONE canonical buffer (rewritten every tile by the
// in-kernel transpose). The 48KB static-smem watermark picks the resident-CTA
// hint for __launch_bounds__ (sm_89: 100KB smem per SM, so two CTAs fit while
// each stays within the static budget).
template <typename Traits, typename LayoutA, typename LayoutB, bool StagedB>
struct Fp8GemmSmem {
// Crosswise = the stage-load's view: A's tag directly, B's transposed.
// A-crosswise always loads direct (L2-typical activations); B-crosswise
// stages only when its contract dim is long enough to stream DRAM.
static constexpr bool kCrossA = std::is_same_v<LayoutA, ColMajor>;
static constexpr bool kCrossB = std::is_same_v<LayoutB, RowMajor>;
static constexpr bool kBStagePath = kCrossB && StagedB;
static constexpr int kBytes =
Traits::kStages * Traits::kBlockM * Traits::kK +
(kBStagePath ? Traits::kStages + 1 : Traits::kStages) *
Traits::kBlockN * Traits::kK;
static constexpr int kMinCtas = kBytes <= 48 * 1024 ? 2 : 1;
};
// LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h): // LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h):
// A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] / // A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] /
// ColMajor = [N][K]. The kernel always computes // ColMajor = [N][K]. The kernel always computes
@@ -342,25 +441,43 @@ __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk
// (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and // (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and
// exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the
// launcher dispatches to it there (see launch_fp8_gemm). // launcher dispatches to it there (see launch_fp8_gemm).
template <typename Traits, bool OutFp8 = false, typename LayoutA = RowMajor, template <typename Traits, bool OutFp8 = false,
typename LayoutB = RowMajor> typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
__global__ void bool kBStaged = true>
__launch_bounds__((Traits::kBlockM / 64) * (Traits::kBlockN / 32) * 32, 2) __global__ void __launch_bounds__(Traits::kCtaThreads,
Fp8GemmSmem<Traits, LayoutA, LayoutB,
kBStaged>::kMinCtas)
fp8_gemm_kernel(FP8Params p) { fp8_gemm_kernel(FP8Params p) {
using T8 = std::conditional_t<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>; using T8 = std::conditional_t<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>;
constexpr int kBlockM = Traits::kBlockM; constexpr int kBlockM = Traits::kBlockM;
constexpr int kBlockN = Traits::kBlockN; constexpr int kBlockN = Traits::kBlockN;
constexpr int kK = Traits::kK; constexpr int kK = Traits::kK;
constexpr int kStages = Traits::kStages; constexpr int kStages = Traits::kStages;
constexpr int kCtaThreads = (kBlockM / 64) * (kBlockN / 32) * 32; constexpr int kCtaThreads = Traits::kCtaThreads;
constexpr bool kCrossA = Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kCrossA;
constexpr bool kCrossB = Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kCrossB;
constexpr bool kBStagePath =
Fp8GemmSmem<Traits, LayoutA, LayoutB, kBStaged>::kBStagePath;
static_assert(kStages >= 1 && kStages <= 8, static_assert(kStages >= 1 && kStages <= 8,
"FP8 GEMM stages must be in the range [1, 8]"); "FP8 GEMM stages must be in [1, 8]");
// Tiles are flat [rows * kK] with a 16B-chunk XOR swizzle (tile_at): // Tiles are flat [rows * kK] with a 16B-chunk XOR swizzle (tile_at):
// ldmatrix reads whole 16B chunks through the same mapping the staging // ldmatrix reads whole 16B chunks through the same mapping the staging
// writes, and the swizzle removes the bank conflict the unswizzled // writes, and the swizzle removes the bank conflict the unswizzled
// 8-word row stride caused (see tile_at). // 8-word row stride caused (see tile_at). The stages live in dynamic
__shared__ __align__(16) T8 a_smem[kStages][kBlockM * kK]; // shared memory so deep pipelines (kStages * (kBlockM + kBlockN) * kK >
__shared__ __align__(16) T8 b_smem[kStages][kBlockN * kK]; // 48KB static limit) opt in via cudaFuncSetAttribute in the launcher.
extern __shared__ __align__(16) char fp8_gemm_smem[];
// Per operand: congruous or direct-crosswise = kStages rotating canonical
// buffers; staged-crosswise (B) = kStages K-major staging buffers (filled
// by cp.async, one per tile in flight) followed by one canonical buffer
// the per-tile transpose rewrites.
constexpr int kAStageBytes = kBlockM * kK;
constexpr int kBStageBytes = kBlockN * kK;
constexpr int kStB = kStages; // B staging ring size (see above)
T8* const a_base = reinterpret_cast<T8*>(fp8_gemm_smem);
T8* const b_base =
reinterpret_cast<T8*>(fp8_gemm_smem + kStages * kAStageBytes);
T8* const b_canon = b_base + kStages * kBStageBytes; // crosswise B only
const auto* a = reinterpret_cast<const T8*>(p.a_ptr); const auto* a = reinterpret_cast<const T8*>(p.a_ptr);
const auto* b = reinterpret_cast<const T8*>(p.b_ptr); const auto* b = reinterpret_cast<const T8*>(p.b_ptr);
@@ -379,12 +496,14 @@ __global__ void
// before advancing along N. All CTAs of one group share the same B column // before advancing along N. All CTAs of one group share the same B column
// stripe, so B tiles stay hot in L2 across the wave (the default // stripe, so B tiles stay hot in L2 across the wave (the default
// N-fastest order makes each wave touch every B tile instead). // N-fastest order makes each wave touch every B tile instead).
// Measured win for the A-crosswise layouts (10-21% at K>=2048) and loss // kGroupRaster is a template knob (the launcher defaults it to the
// for A-congruous (-17..20%, A's cp.async stream prefers the N-fastest // measured best per layout: grouped for A-crosswise (dW) and for the
// order) — so the branch follows LayoutA. // congruous NT forward — whose big B operand gains the most from the
// shared stripe — plain for dX's crosswise-B layouts, where it measured
// neutral).
constexpr int kGroupM = 8; constexpr int kGroupM = 8;
int block_m, block_n; int block_m, block_n;
if constexpr (std::is_same_v<LayoutA, ColMajor>) { if constexpr (kGroupRaster) {
const int blocks_m = gridDim.y; const int blocks_m = gridDim.y;
const int bid = blockIdx.y * gridDim.x + blockIdx.x; const int bid = blockIdx.y * gridDim.x + blockIdx.x;
const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM; const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM;
@@ -409,18 +528,48 @@ __global__ void
const float sb = *p.scale_b; const float sb = *p.scale_b;
float acc[4][4][4] = {}; // [nt][mt][acc] float acc[4][4][4] = {}; // [nt][mt][acc]
// Both operands are staged into the canonical [M][kK] / [N][kK] shared // Both operands end up in the canonical [M][kK] / [N][kK] shared tiles
// tiles regardless of their global layout (see load_operand_tile), so the // the MMA fragments read, regardless of their global layout. A's tag
// MMA fragment reads below stay unchanged across the four layout // already names the operand view ([M][K] = [rows][contract]); B's tag is
// combinations. A's tag already names the operand view ([M][K] = // relative to the canonical [K][N], so the stage-load sees its transpose
// [rows][contract]); B's tag is relative to the canonical [K][N], so the // (transpose_layout_t, see common.h). Congruous operands cp.async
// stage-load sees its transpose (transpose_layout_t, see common.h). // straight into their rotating canonical buffers; crosswise operands
auto load_tile = [&](int stage, int64_t k_base) { // cp.async into K-major staging (zero transformation) and get a per-tile
load_operand_tile<T8, kK, LayoutA, kBlockM, kCtaThreads>( // smem->smem transpose below.
a_smem[stage], a, m, k, a_ld, tid, k_base, (int64_t)block_m * kBlockM); auto load_tile = [&](int tile, int stage, int64_t k_base) {
load_operand_tile<T8, kK, transpose_layout_t<LayoutB>, kBlockN, if constexpr (kCrossA) {
kCtaThreads>(b_smem[stage], b, n, k, b_ld, tid, k_base, load_crosswise_direct<T8, kK, kBlockM, kCtaThreads>(
(int64_t)block_n * kBlockN); a_base + stage * kAStageBytes, a, m, k, a_ld, tid, k_base,
(int64_t)block_m * kBlockM);
} else {
load_operand_tile<T8, kK, kBlockM, kCtaThreads>(
a_base + stage * kAStageBytes, a, m, k, a_ld, tid, k_base,
(int64_t)block_m * kBlockM);
}
if constexpr (kBStagePath) {
stage_crosswise_tile<T8, kK, kBlockN, kCtaThreads>(
b_base + (tile % kStB) * kBStageBytes, b, n, k, b_ld, tid,
k_base, (int64_t)block_n * kBlockN);
} else if constexpr (kCrossB) {
load_crosswise_direct<T8, kK, kBlockN, kCtaThreads>(
b_base + stage * kBStageBytes, b, n, k, b_ld, tid, k_base,
(int64_t)block_n * kBlockN);
} else {
load_operand_tile<T8, kK, kBlockN, kCtaThreads>(
b_base + stage * kBStageBytes, b, n, k, b_ld, tid, k_base,
(int64_t)block_n * kBlockN);
}
};
// smem->smem transpose of one k_seg region (kSegQuads contract quads) of
// this tile's staged-crosswise B into its single canonical buffer.
auto transpose_tile = [&](int tile, int seg) {
if constexpr (!kBStagePath) return;
constexpr int kSegQuads = kK / 4 / (kK / kMmaK); // quads per k_seg
constexpr int kBRegion = kSegQuads * (kBlockN / 16);
const T8* b_stg = b_base + (tile % kStB) * kBStageBytes;
for (int idx = tid; idx < kBRegion; idx += kCtaThreads)
transpose_crosswise_region<T8, kK, kBlockN>(b_canon, b_stg, idx,
seg * kSegQuads);
}; };
const int64_t tile_count = (k + kK - 1) / kK; const int64_t tile_count = (k + kK - 1) / kK;
@@ -445,7 +594,7 @@ __global__ void
#pragma unroll #pragma unroll
for (int stage = 0; stage < kStages; ++stage) { for (int stage = 0; stage < kStages; ++stage) {
if (stage < tile_count) { if (stage < tile_count) {
load_tile(stage, static_cast<int64_t>(stage) * kK); load_tile(stage, stage, static_cast<int64_t>(stage) * kK);
astrai::cp_async_commit_group(); astrai::cp_async_commit_group();
} }
} }
@@ -463,48 +612,77 @@ __global__ void
// before any thread reads tiles written by other threads. // before any thread reads tiles written by other threads.
__syncthreads(); __syncthreads();
// Staged-crosswise B: produce the canonical tile one k_seg region at
// a time so each region's transpose overlaps the previous region's
// MMA sequence (the transposes are pure shared-memory traffic — B's
// global path stayed fully asynchronous above).
constexpr int kSegs = kK / kMmaK;
if constexpr (kBStagePath) {
transpose_tile(tile_index, 0);
// Barrier 2: region 0 visible to every thread before its
// fragment loads. (Compiled out for congruous/direct layouts.)
__syncthreads();
}
const T8* a_tile = a_base + (size_t)stage * kAStageBytes;
const T8* b_tile =
kBStagePath ? b_canon : b_base + (size_t)stage * kBStageBytes;
// 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg — // 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg —
// 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in // 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in
// the 128x64-tile version (the kernel was LSU-issue-bound there). // the 128x64-tile version (the kernel was LSU-issue-bound there).
constexpr int kSegs = kK / kMmaK; // B fragments double-buffer across k_segs while B is congruous (no
// B fragments double-buffered across k_segs: the next k_seg's B load // region writes in flight); a crosswise B reloads per k_seg after
// is issued before the current k_seg's MMA sequence, so its LDS // the region's transpose became visible.
// latency hides behind the A pipeline + tensor-pipe work (same trick
// as the A mt+1 prefetch below; costs kSegs x 8 registers).
unsigned b_frag[2][4][2]; unsigned b_frag[2][4][2];
if constexpr (!kBStagePath) {
#pragma unroll #pragma unroll
for (int nt = 0; nt < 4; ++nt) { for (int nt = 0; nt < 4; ++nt) {
const int row = b_row0 + nt * 8 + r7; const int row = b_row0 + nt * 8 + r7;
astrai::ldmatrix_x2_lane(b_frag[0][nt], astrai::ldmatrix_x2_lane(b_frag[0][nt],
frag_addr<T8, kK>(b_smem[stage], row, rh8)); frag_addr<T8, kK>(b_tile, row, rh8));
}
} }
#pragma unroll #pragma unroll
for (int k_seg = 0; k_seg < kSegs; ++k_seg) { for (int k_seg = 0; k_seg < kSegs; ++k_seg) {
const int bcur = k_seg & 1, bnext = bcur ^ 1; const int bcur = k_seg & 1, bnext = bcur ^ 1;
if (k_seg + 1 < kSegs) { if constexpr (kBStagePath) {
#pragma unroll
for (int nt = 0; nt < 4; ++nt) {
const int row = b_row0 + nt * 8 + r7;
astrai::ldmatrix_x2_lane(
b_frag[bcur][nt],
frag_addr<T8, kK>(b_tile, row, k_seg * 2 + rh8));
}
} else if (k_seg + 1 < kSegs) {
#pragma unroll #pragma unroll
for (int nt = 0; nt < 4; ++nt) { for (int nt = 0; nt < 4; ++nt) {
const int row = b_row0 + nt * 8 + r7; const int row = b_row0 + nt * 8 + r7;
astrai::ldmatrix_x2_lane( astrai::ldmatrix_x2_lane(
b_frag[bnext][nt], b_frag[bnext][nt],
frag_addr<T8, kK>(b_smem[stage], row, frag_addr<T8, kK>(b_tile, row, (k_seg + 1) * 2 + rh8));
(k_seg + 1) * 2 + rh8));
} }
} }
// Region k_seg+1's transpose overlaps this region's MMA work
// (disjoint canonical regions, no race).
if constexpr (kBStagePath) {
if (k_seg + 1 < kSegs)
transpose_tile(tile_index, k_seg + 1);
}
// Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1
// is issued before the MMAs consuming row mt, so the LDS fixed // is issued before the MMAs consuming row mt, so the LDS fixed
// latency hides behind tensor-pipe work (cuts the `wait` stall, // latency hides behind tensor-pipe work (cuts the `wait` stall,
// ~2.3 cycles/issue before this). Costs 4 extra registers. // ~2.3 cycles/issue before this). Costs 4 extra registers.
unsigned a_frag[5][4]; unsigned a_frag[5][4];
astrai::ldmatrix_x4_lane( astrai::ldmatrix_x4_lane(
a_frag[0], frag_addr<T8, kK>(a_smem[stage], a_row0 + rh8 * 8 + r7, a_frag[0], frag_addr<T8, kK>(a_tile, a_row0 + rh8 * 8 + r7,
k_seg * 2 + rh16)); k_seg * 2 + rh16));
#pragma unroll #pragma unroll
for (int mt = 0; mt < 4; ++mt) { for (int mt = 0; mt < 4; ++mt) {
if (mt < 3) if (mt < 3)
astrai::ldmatrix_x4_lane( astrai::ldmatrix_x4_lane(
a_frag[mt + 1], a_frag[mt + 1],
frag_addr<T8, kK>(a_smem[stage], frag_addr<T8, kK>(a_tile,
a_row0 + (mt + 1) * 16 + rh8 * 8 + r7, a_row0 + (mt + 1) * 16 + rh8 * 8 + r7,
k_seg * 2 + rh16)); k_seg * 2 + rh16));
#pragma unroll #pragma unroll
@@ -512,12 +690,19 @@ __global__ void
astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt], b_frag[bcur][nt], astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt], b_frag[bcur][nt],
acc[nt][mt]); acc[nt][mt]);
} }
// Barrier 3: region k_seg+1's transposes complete and become
// visible before the next k_seg reads them.
if constexpr (kBStagePath) {
if (k_seg + 1 < kSegs) __syncthreads();
}
} }
// Barrier 2: every thread finished reading this stage's tiles before // Barrier 4: every thread finished reading this stage's tiles before
// the prefetch for the (i+kStages)-th tile overwrites them. // the prefetch for the (i+kStages)-th tile overwrites them (and the
// next iteration's transposes rewrite the canonical buffer).
__syncthreads(); __syncthreads();
if (tile_index + kStages < tile_count) { if (tile_index + kStages < tile_count) {
load_tile(stage, (tile_index + kStages) * kK); load_tile(tile_index + kStages, stage,
(tile_index + kStages) * kK);
astrai::cp_async_commit_group(); astrai::cp_async_commit_group();
} }
} }
@@ -593,26 +778,74 @@ void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
fp8_quantize_kernel<Fmt><<<blocks, kThreads, 0, stream>>>(p); fp8_quantize_kernel<Fmt><<<blocks, kThreads, 0, stream>>>(p);
} }
// Launch one kernel instantiation with its shared-memory budget: stages live
// in dynamic smem, so budgets beyond the 48KB static limit opt in once per
// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
// memory"). Templated on the kernel *value* (auto NTTP) so every
// instantiation owns its own armed flag — same-signature kernels must not
// share it (the attribute is per-function).
template <auto Kernel, typename... Args>
void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
cudaStream_t stream, Args... args) {
if (smem_bytes > 48 * 1024) {
static bool armed = false; // per instantiation
if (!armed) {
cudaFuncSetAttribute(Kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_bytes);
armed = true;
}
}
Kernel<<<grid, block, smem_bytes, stream>>>(args...);
}
// Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles). // Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles).
// kK selects the K tile (32 or 64; 64 halves the __syncthreads count per K // kK selects the K tile (32 / 64 / 128; larger kK halves the __syncthreads
// and doubles the MMA work per stage, at 2x the smem per stage — measured // count per K and doubles the MMA work per stage at more smem per stage).
// 10-35% across shapes, so 64 is the default). Stages=2 with kK=64 keeps the // Stages is the cp.async pipeline depth (smem = Stages * (BM + BN) * kK
// pipeline at 32KB smem; deeper pipelines only win on K >= 4096 squares and // bytes for congruous layouts; deep pipelines are dynamic-smem backed, 1
// lose elsewhere. LayoutA/LayoutB mirror the kernel template (defaults keep // CTA/SM past 48KB). GroupRaster defaults to the historically-measured best
// the NN layout: out = a @ b). m <= 64 dispatches to the 64x128 CTA — a // per LayoutA (grouped for A-crosswise, plain for A-congruous). m <= 64
// 128-row CTA would waste half its MMA work on predicated-off rows. // dispatches to the 64x128 CTA — a 128-row CTA would waste half its MMA work
// on predicated-off rows.
// Crosswise B takes the asynchronous staging+transpose pipeline only when
// the contract dim is long enough that B streams from DRAM (dX-class GEMMs,
// k = N_ffn); short-K crosswise GEMMs (dW: k = M tokens) read L2-resident
// operands, where the staging round trip costs more shared-memory traffic
// than the latency it hides (measured: dW ~37 TF direct vs ~29 TF staged,
// dX ~39 TF staged vs ~38 direct).
constexpr int64_t kCrossStageMinK = 8192;
template <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor, template <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor,
typename LayoutB = RowMajor, int kK = 64, int Stages = 2> typename LayoutB = RowMajor, int kK = 64, int Stages = 2,
bool GroupRaster = std::is_same_v<LayoutA, ColMajor> || std::is_same_v<LayoutB, ColMajor>>
void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
dim3 grid((p.n + 127) / 128, (p.m + 127) / 128); dim3 grid((p.n + 127) / 128, (p.m + 127) / 128);
const bool b_staged = p.k >= kCrossStageMinK;
if (p.m <= 64) { if (p.m <= 64) {
using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>; using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB> if (b_staged)
<<<grid, (64 / 64) * (128 / 32) * 32, 0, stream>>>(p); launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
GroupRaster, true>>(
Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
dim3(Traits::kCtaThreads), stream, p);
else
launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
GroupRaster, false>>(
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
dim3(Traits::kCtaThreads), stream, p);
} else { } else {
using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>; using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB> if (b_staged)
<<<grid, (128 / 64) * (128 / 32) * 32, 0, stream>>>(p); launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
GroupRaster, true>>(
Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
dim3(Traits::kCtaThreads), stream, p);
else
launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
GroupRaster, false>>(
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
dim3(Traits::kCtaThreads), stream, p);
} }
} }
csrc/kernels/fp8/ops.cu
+41 -16
View File
@@ -157,7 +157,10 @@ std::tuple<torch::Tensor, torch::Tensor> quantize_bf16(torch::Tensor x,
auto x8 = torch::empty_like( auto x8 = torch::empty_like(
x_c, x_c.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}, x_c.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; FP8QuantizeParams p;
pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax, pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax,
nullptr, 0, 0, x_c.numel()); nullptr, 0, 0, x_c.numel());
@@ -230,18 +233,25 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
return out; return out;
} }
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8( std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx, torch::Tensor>
torch::Tensor sw, int64_t fmt, c10::optional<torch::Tensor> bias_scale, linear_forward_fp8(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
c10::optional<torch::Tensor> x_ring, int64_t x_ring_idx, torch::Tensor sx, torch::Tensor sw, int64_t fmt,
int64_t x_ring_margin, c10::optional<torch::Tensor> w_ring, c10::optional<torch::Tensor> bias_scale,
int64_t w_ring_idx, int64_t w_ring_margin) { 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 // Pure FP8 forward: quantize x/w (fmt: 0 = E4M3, 1 = E5M2), then the
// pre-quantized GEMM; the dequantized BF16 output gets the bias added. // pre-quantized GEMM; the dequantized BF16 output gets the bias added.
// amax_x / amax_w come from the quantize kernels (zero-initialized here; // Returns (out, x8, w8, amax_x, amax_w): the quantized operands are
// a pre-quantized w reports amax_w = 0 — nothing to feed a delayed ring). // handed back so the policy layer can cache the weight quantization
// w may itself be pre-quantized fp8 storage matching fmt (static // (torch autocast's cached_cast analog — w8 is reused while the weight
// inference weights): the weight quantize is skipped, amax_w stays 0. // 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 // When x_ring / w_ring are given (delayed scaling), the quantize kernels
// finalize them in-kernel: the returned amax is already folded into the // finalize them in-kernel: the returned amax is already folded into the
// ring window and the next step's scale is published on device. // ring window and the next step's scale is published on device.
@@ -276,8 +286,16 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
if (b_prequant) check_scale(*bias_scale, x, "bias_scale"); if (b_prequant) check_scale(*bias_scale, x, "bias_scale");
} }
auto x8 = torch::empty({m, k}, x_c.options().dtype(f8opt)); auto x8 = torch::empty({m, k}, x_c.options().dtype(f8opt));
auto amax_x = torch::zeros({1}, x.options().dtype(torch::kFloat32)); // Each amax slot feeds a cross-block atomic_max, so it must start at 0.
auto amax_w = torch::zeros({1}, x.options().dtype(torch::kFloat32)); // 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 out = torch::empty({m, n}, x_c.options());
auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst, auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst,
@@ -322,7 +340,7 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1); std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
shape.push_back(n); shape.push_back(n);
return {out.reshape(shape), amax_x, amax_w}; return {out.reshape(shape), x8, w8, amax_x, amax_w};
} }
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
@@ -360,7 +378,13 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
auto grad_input = torch::empty_like(x); auto grad_input = torch::empty_like(x);
auto grad_weight = torch::empty_like(w); auto grad_weight = torch::empty_like(w);
auto grad_bias = torch::empty({0}, g.options()); auto grad_bias = torch::empty({0}, g.options());
auto amax_g = torch::zeros({1}, g.options().dtype(torch::kFloat32)); // amax_g feeds a cross-block atomic_max in the g quantize kernel; zero it
// on the stream (empty + memset, not torch::zeros — see the forward).
// Only needed when a g quantize runs (mask[0]||mask[1]); the bias-only
// fallback below overwrites it via .copy_, so a wasted memset elsewhere
// is harmless.
auto amax_g = torch::empty({1}, g.options().dtype(torch::kFloat32));
cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), stream.stream());
auto f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2) auto f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2)
: g.options().dtype(torch::kFloat8_e4m3fn); : g.options().dtype(torch::kFloat8_e4m3fn);
@@ -453,7 +477,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
"matching fmt (static inference path; fp8 bias requires bias_scale);" "matching fmt (static inference path; fp8 bias requires bias_scale);"
" x_ring/w_ring optionally finalize a delayed-scaling ring " " x_ring/w_ring optionally finalize a delayed-scaling ring "
"([hist | scale | counter] float32 buffer) in-kernel; returns " "([hist | scale | counter] float32 buffer) in-kernel; returns "
"(out, amax_x, amax_w)"); "(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"), 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("x"), py::arg("w"), py::arg("masks"), py::arg("sg"),
py::arg("sw"), py::arg("sx"), py::arg("fmt"), py::arg("sw"), py::arg("sx"), py::arg("fmt"),
tests/extension/test_fp8_mma.py
+71 -9
View File
@@ -6,15 +6,22 @@ policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks
of the custom ops) run without a GPU. of the custom ops) run without a GPU.
""" """
import threading
import pytest import pytest
import torch import torch
import torch.nn.functional as F
import astrai.extension.fp8 as f8mod
from astrai.extension.fp8 import ( from astrai.extension.fp8 import (
DelayedScaling, DelayedScaling,
DynamicScaling, DynamicScaling,
FP8Format, FP8Format,
FP8TensorMeta, FP8TensorMeta,
_ScaleRing,
fp8_autocast, fp8_autocast,
fp8_linear_enable,
fp8_linear_enabled,
fp8_state, fp8_state,
) )
from astrai.extension.ops.fp8 import ( from astrai.extension.ops.fp8 import (
@@ -91,8 +98,6 @@ def test_quantize_ring_in_kernel_finalize():
"""The quantize kernel finalizes the delayed-scaling ring in-kernel: the """The quantize kernel finalizes the delayed-scaling ring in-kernel: the
measured amax lands in hist[idx], the window reduces to the next step's measured amax lands in hist[idx], the window reduces to the next step's
scale on device, and the counter re-arms for the next launch.""" scale on device, and the counter re-arms for the next launch."""
from astrai.extension.fp8 import _ScaleRing
torch.manual_seed(21) torch.manual_seed(21)
dev = torch.device("cuda") dev = torch.device("cuda")
ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0)) ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0))
@@ -141,7 +146,7 @@ def test_fp8_linear_forward_and_backward():
bias = torch.randn(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) scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad)
out, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w) out, x8, w8, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w)
grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3" grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3"
) )
@@ -205,7 +210,7 @@ def test_fp8_linear_static_fp8_weight_and_bias():
w8, _ = quantize_bf16(weight, sw, "e4m3") w8, _ = quantize_bf16(weight, sw, "e4m3")
b8, _ = quantize_bf16(bias, sb, "e4m3") b8, _ = quantize_bf16(bias, sb, "e4m3")
out, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb) out, x8, w8_back, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb)
qx = _quantize(x, sx) qx = _quantize(x, sx)
qw = _quantize(weight, sw) qw = _quantize(weight, sw)
@@ -214,9 +219,10 @@ def test_fp8_linear_static_fp8_weight_and_bias():
torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1)) 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 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 # bf16 bias stays bf16 on the same fused-epilogue path
out_bf16bias, _, _ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3") out_bf16bias, *_ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3")
expected_b = (qx @ qw.t() * sx * sw + bias.float()).to(torch.bfloat16) 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) torch.testing.assert_close(out_bf16bias, expected_b, atol=0.125, rtol=0.01)
@@ -226,10 +232,6 @@ def test_fp8_linear_backward_outside_autocast():
"""aten::linear records an fp8 autograd node inside fp8_autocast; the """aten::linear records an fp8 autograd node inside fp8_autocast; the
backward runs fp8 kernels even after the context exits (loss.backward() backward runs fp8 kernels even after the context exits (loss.backward()
placement is free), instead of falling back to bf16 mm.""" placement is free), instead of falling back to bf16 mm."""
import torch.nn.functional as F
import astrai.extension.fp8 as f8mod
torch.manual_seed(5) torch.manual_seed(5)
x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
weight = torch.randn( weight = torch.randn(
@@ -424,3 +426,63 @@ def test_mm_fp8_fp8_output_cpu():
assert out8.dtype == torch.float8_e4m3fn assert out8.dtype == torch.float8_e4m3fn
ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn) ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn)
assert torch.equal(out8, ref) assert torch.equal(out8, ref)
# --------------------------------------------------------------------------
# torch-autocast parity: context semantics (nesting, thread locality, switch)
# --------------------------------------------------------------------------
def _linear():
"""Shared helper: a small bf16 linear operand set on CUDA (grad-tracking
so aten::linear records an autograd node)."""
torch.manual_seed(31)
x = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
w = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
return x, w
@skip_no_fp8
def test_nested_disabled_region_redispatches_bf16():
"""A nested fp8_autocast(enabled=False) region temporarily restores the
bf16 aten::linear path (torch's nested-disable semantics), and fp8
resumes when it exits."""
x, w = _linear()
with fp8_autocast(enabled=True):
out_fp8 = F.linear(x, w)
with fp8_autocast(enabled=False):
out_bf16 = F.linear(x, w)
assert type(out_bf16.grad_fn).__name__ != "_LinearFp8Backward"
assert out_bf16.dtype == torch.bfloat16
out_again = F.linear(x, w)
assert type(out_again.grad_fn).__name__ == "_LinearFp8Backward"
@skip_no_fp8
def test_global_switch_routes_without_region():
"""fp8_linear_enable(True) routes aten::linear to fp8 outside any region
(the persistent default); disabling restores bf16."""
x, w = _linear()
state = fp8_state()
try:
fp8_linear_enable(True)
out = F.linear(x, w)
assert type(out.grad_fn).__name__ == "_LinearFp8Backward"
fp8_linear_enable(False)
out = F.linear(x, w)
assert type(out.grad_fn).__name__ != "_LinearFp8Backward"
finally:
state.reset()
def test_autocast_state_is_thread_local():
"""torch parity: the active config is thread-local — another thread does
not see an open region (CPU-only check of the flag, no kernels)."""
seen = {}
with fp8_autocast(enabled=True):
assert fp8_linear_enabled()
t = threading.Thread(target=lambda: seen.update(enabled=fp8_linear_enabled()))
t.start()
t.join()
assert seen["enabled"] is False
assert not fp8_linear_enabled()