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