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

perf: fold the delayed-scaling ring update into the quantize kernel

Browse Source 
- the kernel's last block folds amax into the history window and publishes the next scale in-kernel (atomicAdd ticket + fences), replacing the host update chain
- quantize bindings split into quantize(transposed) / quantize_dual with fixed arities and a QuantLayout enum; the python adapter becomes a thin attention-style wrapper over pybind (Optional ring_state at the boundary, no torch.library custom_ops)
- tests: in-kernel fold vs host reference (exact), dual/transposed orientation byte-equality

Benchmark: L20 (sm_89), 1.2B model, full train step. Per-linear fixed overhead 28.8us -> 8.8us; fp8 vs bf16: M=512 77.5ms, M=2048 144.5ms (1.15x), M=8192 527.4ms (1.28x); losses bit-identical.
This commit is contained in:
ViperEkura
2026-08-31 14:24:51 +08:00
parent 1cf7d6c76b
commit 962c10c52b
6 changed files with 324 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/fp8.py
+37 -26
View File
@@ -36,7 +36,7 @@ from typing import Dict, List, Optional
import torch
from torch.library import Library
from astrai.extension.ops.fp8 import mm_fp8, quantize
from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
# Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -92,10 +92,12 @@ class DynamicScaling(FP8Recipe):
class _ScaleRing:
"""One operand's delayed-scaling state: a float32 buffer
``[hist[n] | scale | counter]`` (views). ``update`` folds the amax
returned by the quantize primitive into ``hist[idx]`` and publishes the
next scale from the window; ``idx`` advances host-side each step. The
trailing slot is a legacy counter kept for state-buffer compatibility.
``[hist[n] | scale | legacy | amax | done]`` (views). The quantize
kernel folds its fused amax into ``hist[idx]`` and publishes the next
scale from the window in its own last block (``fold_args`` passes the
buffer + recipe constants); ``idx`` advances host-side each use. The
``amax``/``done`` tail slots are kernel scratch (self-cleaning across
launches); the legacy slot keeps state-buffer compatibility.
"""
__slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -103,7 +105,7 @@ class _ScaleRing:
def __init__(self, device: torch.device, recipe: FP8Recipe):
self.recipe = recipe
n = recipe.history_len
self.state = torch.zeros(n + 2, device=device, dtype=torch.float32)
self.state = torch.zeros(n + 4, device=device, dtype=torch.float32)
self.hist = self.state[:n]
self.scale = self.state[n : n + 1]
self.idx = 0
@@ -119,9 +121,14 @@ class _ScaleRing:
self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
self.initialized = True
def update(self, amax: torch.Tensor, fmt: str) -> None:
self.hist[self.idx].copy_(amax.reshape(()))
self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
def fold_args(self, fmt: str) -> dict:
"""Keyword arguments for quantize()'s in-kernel history fold."""
return {
"ring_state": self.state,
"hist_idx": self.idx,
"fp8_max": FP8_MAX[fmt],
"pow2_margin": float(2**self.recipe.margin),
}
class FP8TensorMeta:
@@ -322,11 +329,11 @@ def fp8_linear_forward(
Composed from the two stateless primitives: quantize x/w with the active
scales, run the pre-quantized GEMM with the bias fused into its epilogue.
Delayed scaling folds
the returned amax into the history ring and publishes the next scale;
dynamic scaling measures the current amax itself. Training quantizes the
weight every step (the optimizer bumps its version, so there is no cast
cache, matching ``cached_cast``-less behavior).
Delayed scaling lets the quantize kernel fold the fused amax into the
history ring and publish the next scale in its own last block; dynamic
scaling measures the current amax itself. Training quantizes the weight
every step (the optimizer bumps its version, so there is no cast cache,
matching ``cached_cast``-less behavior).
"""
state = fp8_state()
if cfg is None:
@@ -351,19 +358,19 @@ def fp8_linear_forward(
if not meta.x.initialized:
meta.x.seed(x, fmt)
sx, sw = meta.x.scale.clone(), meta.w.scale.clone()
x8, amax_x = quantize(x, sx.reciprocal(), fmt)
# The clones feed this call's kernels (stream-ordered before the in-kernel
# fold overwrites the ring scale slots); the fp8 quantize kernel folds the
# amax into the history window and publishes the next scale itself.
x8, _ = quantize(x, sx.reciprocal(), fmt, **meta.x.fold_args(fmt))
if _is_fp8(w.dtype):
w8, amax_w = w, None
w8 = w
else:
w8, amax_w = quantize(w, sw.reciprocal(), fmt)
w8, _ = quantize(w, sw.reciprocal(), fmt, **meta.w.fold_args(fmt))
out = mm_fp8(
x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
).reshape(*x.shape[:-1], w.size(0))
meta.x.update(amax_x, fmt)
if amax_w is not None:
meta.w.update(amax_w, fmt)
meta.x.advance()
if amax_w is not None:
if not _is_fp8(w.dtype):
meta.w.advance()
return out, sx, sw
@@ -411,22 +418,26 @@ class _LinearFp8(torch.autograd.Function):
# quantize outputs: g8 [m,n] with w8T [k,n] (trans_b=True) gives
# grad_x, g8T [n,m] with x8T [k,m] gives grad_w — no NN-swap or TT
# crosswise kernel in the training path. g is consumed in both
# orientations, so one dual-layout pass feeds both.
g8, g8T, amax_g = quantize(g2, sg.reciprocal(), fmt, layout=2)
x8T, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, layout=1)
# orientations, so quantize_dual's single pass feeds both.
# The g quantize folds the gradient amax into its ring in-kernel;
# the x8T/w8T orientation copies discard amax (those rings were
# folded at forward time).
g8, g8T, _ = quantize_dual(g2, sg.reciprocal(), fmt, **meta.g.fold_args(fmt))
x8T, _ = quantize(
x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, transposed=True
)
if _is_fp8(w.dtype):
# Pre-quantized weight has no transposed copy: keep the swap
# path for grad_x (grad_w is unaffected).
grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape)
else:
w8T, _ = quantize(w, sw.reciprocal(), fmt, layout=1)
w8T, _ = quantize(w, sw.reciprocal(), fmt, transposed=True)
grad_x = mm_fp8(g8, w8T, sg * sw, trans_b=True).reshape(x.shape)
grad_w = mm_fp8(g8T, x8T, sg * sx, trans_b=True) # g8.T @ x8
# bias-free linears must not pay the column-sum
# reduce: g2.sum(0) is another full read of the gradient.
grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None
if not ctx.is_dynamic:
meta.g.update(amax_g, fmt)
meta.g.advance()
return grad_x, grad_w, grad_b
astrai/extension/ops/fp8.py
+65 -200
View File
@@ -1,14 +1,20 @@
"""FP8 CUDA kernel interface adapter (the only module touching the pybind).
Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
Attention-style thin wrappers: one Python entry per binding, called directly
— no torch.library dispatch layer. Optional arguments (``ring_state``,
``bias``) keep native Optional semantics at the pybind boundary, and
in-place buffer updates (the delayed-scaling ring fold, like attention's
KV-cache appends) happen on-stream without mutation declarations. CUDA-only:
non-CUDA or unsupported inputs raise from the binding's TORCH_CHECKs.
- ``quantize(x, scale, fmt) -> (x8, amax)`` — BF16/FP16/FP32 → FP8 with fused amax
- ``quantize(x, scale, fmt, transposed=False) -> (x8|x8T, amax)`` — BF16/FP16/FP32
→ FP8 with fused amax (``transposed`` picks the orientation; arity is fixed)
- ``quantize_dual(x, scale, fmt) -> (x8, x8T, amax)`` — both orientations, one read
- ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
Scale semantics: scales are *quantization steps* — the value divided out when
quantizing (``x8 = x / scale``). Every primitive computes its own inverse
internally; callers never pass ``scale_inv``. ``amax`` values are *returned*,
never passed as output arguments. ``fmt`` is ``"e4m3"`` or ``"e5m2"``.
``scale`` is the quantization multiplier (device scalar); ``fmt`` is
``"e4m3"`` or ``"e5m2"``. ``amax`` values are *returned*, never passed as
output arguments.
Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
this module is stateless.
@@ -17,7 +23,6 @@ this module is stateless.
from typing import Optional, Tuple
import torch
from torch.library import custom_op
from astrai.extension.loader import get_module
@@ -32,192 +37,63 @@ def _fmt_int(fmt: str) -> int:
raise ValueError(f"unsupported fp8 format {fmt!r} (expected 'e4m3' or 'e5m2')")
def _fmt_name(fmt: int) -> str:
if fmt == 0:
return "e4m3"
if fmt == 1:
return "e5m2"
raise ValueError(f"unsupported quantization type {fmt!r}")
def _fmt_dtype(fmt: str) -> torch.dtype:
return torch.float8_e5m2 if _fmt_int(fmt) else torch.float8_e4m3fn
@custom_op("custom::fp8_quantize", mutates_args=())
def fp8_quantize(
x: torch.Tensor, scale: torch.Tensor, fmt: int
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; ``scale`` is a multiplier."""
@fp8_quantize.register_fake
def _fp8_quantize_fake(x, scale, fmt):
dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
return (
torch.empty(x.shape, device=x.device, dtype=dtype),
torch.empty(1, device=x.device, dtype=torch.float32),
)
_QUANT_INPUT_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
@custom_op("custom::fp8_quantize_t", mutates_args=())
def fp8_quantize_t(
x: torch.Tensor, scale: torch.Tensor, fmt: int
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Transposed-output variant of fp8_quantize: returns ``(x8T, amax)``
where ``x8T`` is the [cols][rows] row-major transpose of the quantized
input (the K-contiguous operand orientation for NT GEMMs)."""
@fp8_quantize_t.register_fake
def _fp8_quantize_t_fake(x, scale, fmt):
dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
rows, cols = x.shape[-2], x.shape[-1]
return (
torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
torch.empty(1, device=x.device, dtype=torch.float32),
)
@fp8_quantize_t.register_kernel("cuda")
def _fp8_quantize_t_cuda(x, scale, fmt):
if x.dtype not in _QUANT_INPUT_DTYPES:
raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
return get_module("fp8_ops").quantize(x, scale, int(fmt), 1)
@fp8_quantize_t.register_kernel("cpu")
def _fp8_quantize_t_cpu(x, scale, fmt):
x8, amax = _fp8_quantize_cpu(x, scale, fmt)
return x8.transpose(-2, -1).contiguous(), amax
@custom_op("custom::fp8_quantize_dual", mutates_args=())
def fp8_quantize_dual(
x: torch.Tensor, scale: torch.Tensor, fmt: int
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Dual-orientation quantize: one read of ``x`` produces both the
row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
consumed by GEMMs on both orientations (backward ``g``)."""
@fp8_quantize_dual.register_fake
def _fp8_quantize_dual_fake(x, scale, fmt):
dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
rows, cols = x.shape[-2], x.shape[-1]
return (
torch.empty(x.shape, device=x.device, dtype=dtype),
torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
torch.empty(1, device=x.device, dtype=torch.float32),
)
@fp8_quantize_dual.register_kernel("cuda")
def _fp8_quantize_dual_cuda(x, scale, fmt):
if x.dtype not in _QUANT_INPUT_DTYPES:
raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
return get_module("fp8_ops").quantize(x, scale, int(fmt), 2)
@fp8_quantize_dual.register_kernel("cpu")
def _fp8_quantize_dual_cpu(x, scale, fmt):
x8, amax = _fp8_quantize_cpu(x, scale, fmt)
return x8, x8.transpose(-2, -1).contiguous(), amax
@fp8_quantize.register_kernel("cuda")
def _fp8_quantize_cuda(x, scale, fmt):
if x.dtype not in _QUANT_INPUT_DTYPES:
raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
return get_module("fp8_ops").quantize(x, scale, int(fmt))
@fp8_quantize.register_kernel("cpu")
def _fp8_quantize_cpu(x, scale, fmt):
x8 = (x.float() * scale).to(_fmt_dtype(_fmt_name(fmt)))
amax = x.abs().amax().float().reshape(1).clamp_min(1e-12)
return x8, amax
@custom_op("custom::fp8_gemm", mutates_args=())
def fp8_gemm(
a: torch.Tensor,
b: torch.Tensor,
scale: torch.Tensor,
trans_a: int = 0,
trans_b: int = 0,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""FP8 GEMM: ``a @ b * scale (+ bias)`` with FP32 accumulation.
2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules).
``bias`` (bf16, length n) fuses into the epilogue in fp32 before the
single bf16 rounding. The result is always BF16; FP8 output is a
separate quantize operation.
"""
@fp8_gemm.register_fake
def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0, bias=None):
dtype = torch.bfloat16
rows = a.size(2) if trans_a else a.size(1)
cols = b.size(1) if trans_b else b.size(2)
batches = [t.size(0) for t in (a, b) if t.dim() == 3]
shape = (max(batches), rows, cols) if batches else (rows, cols)
return torch.empty(shape, device=a.device, dtype=dtype)
@fp8_gemm.register_kernel("cuda")
def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0, bias=None):
if a.dtype != b.dtype or a.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
raise TypeError(
f"fp8 GEMM requires matching fp8 inputs, got {a.dtype}/{b.dtype}"
)
return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
@fp8_gemm.register_kernel("cpu")
def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
aa = a.float().transpose(-2, -1) if trans_a else a.float()
bb = b.float().transpose(-2, -1) if trans_b else b.float()
acc = aa @ bb * scale
if bias is not None and bias.numel() > 0:
acc = acc + bias.float()
return acc.to(torch.bfloat16)
def quantize(
x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3", layout: int = 0
) -> tuple:
x: torch.Tensor,
scale: torch.Tensor,
fmt: str = "e4m3",
transposed: bool = False,
ring_state: Optional[torch.Tensor] = None,
hist_idx: int = 0,
fp8_max: float = 448.0,
pow2_margin: float = 1.0,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Float (bf16/fp16/fp32) -> FP8 quantize with fused amax.
``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. ``layout``
picks the output orientation: 0 = row-major ``(x8, amax)``; 1 =
transposed ``[cols][rows]`` ``(x8T, amax)`` — the K-contiguous operand
orientation NT GEMMs want; 2 = both from one read ``(x8, x8T, amax)``
(for tensors consumed in both orientations, e.g. backward ``g``).
E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
``transposed=True`` swaps ``x8`` for ``x8T``, the ``[cols][rows]``
row-major transpose of the quantized input — the K-contiguous operand
orientation NT GEMMs want — at the same 2-tuple arity.
``ring_state`` (a 1D float32 CUDA buffer laid out
``[hist n | scale | legacy | amax | done]``) switches on the in-kernel
delayed-scaling fold: the kernel's last block folds the amax into
``hist[hist_idx]`` and publishes the next scale as
``max(hist) / fp8_max / pow2_margin`` — the returned ``amax`` is then the
self-cleaned persistent slot (reads zero). None keeps the classic
fresh-amax return.
"""
# Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a
# 512-wide GEMM): real CUDA tensors of a supported dtype go straight to
# the extension. Fake/subclass tensors and non-CUDA inputs keep the
# custom_op route so torch.compile / meta / fake-tensor tracing and the
# CPU fallback behave exactly as before.
if (
type(x) is torch.Tensor
and x.is_cuda
and x.dtype in _QUANT_INPUT_DTYPES
and fmt in _FMT_TO_INT
):
return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt], layout)
if layout == 0:
return fp8_quantize(x, scale, _fmt_int(fmt))
if layout == 1:
return fp8_quantize_t(x, scale, _fmt_int(fmt))
return fp8_quantize_dual(x, scale, _fmt_int(fmt))
return get_module("fp8_ops").quantize(
x,
scale,
_fmt_int(fmt),
transposed,
ring_state,
hist_idx,
fp8_max,
pow2_margin,
)
def quantize_dual(
x: torch.Tensor,
scale: torch.Tensor,
fmt: str = "e4m3",
ring_state: Optional[torch.Tensor] = None,
hist_idx: int = 0,
fp8_max: float = 448.0,
pow2_margin: float = 1.0,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Dual-orientation quantize: one read of ``x`` produces both the
row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
consumed by GEMMs in both orientations (backward ``g``).
``ring_state`` switches on the in-kernel delayed-scaling fold exactly as
in :func:`quantize`.
"""
return get_module("fp8_ops").quantize_dual(
x, scale, _fmt_int(fmt), ring_state, hist_idx, fp8_max, pow2_margin
)
def mm_fp8(
@@ -237,15 +113,4 @@ def mm_fp8(
kernel epilogue in fp32 — no separate elementwise pass. The result is
BF16; FP8 output is a separate quantize operation.
"""
# Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep
# validation identical on the direct route (bias may be None — the
# binding resolves it to the no-bias path).
if (
type(a) is torch.Tensor
and a.is_cuda
and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
):
return get_module("fp8_ops").mm_fp8(
a, b, scale, int(trans_a), int(trans_b), bias
)
return fp8_gemm(a, b, scale, trans_a, trans_b, bias)
return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)