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 import torch
from torch.library import Library 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). # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0} FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@@ -92,10 +92,12 @@ class DynamicScaling(FP8Recipe):
class _ScaleRing: class _ScaleRing:
"""One operand's delayed-scaling state: a float32 buffer """One operand's delayed-scaling state: a float32 buffer
``[hist[n] | scale | counter]`` (views). ``update`` folds the amax ``[hist[n] | scale | legacy | amax | done]`` (views). The quantize
returned by the quantize primitive into ``hist[idx]`` and publishes the kernel folds its fused amax into ``hist[idx]`` and publishes the next
next scale from the window; ``idx`` advances host-side each step. The scale from the window in its own last block (``fold_args`` passes the
trailing slot is a legacy counter kept for state-buffer compatibility. 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") __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
@@ -103,7 +105,7 @@ 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
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.hist = self.state[:n]
self.scale = self.state[n : n + 1] self.scale = self.state[n : n + 1]
self.idx = 0 self.idx = 0
@@ -119,9 +121,14 @@ class _ScaleRing:
self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt)) self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
self.initialized = True self.initialized = True
def update(self, amax: torch.Tensor, fmt: str) -> None: def fold_args(self, fmt: str) -> dict:
self.hist[self.idx].copy_(amax.reshape(())) """Keyword arguments for quantize()'s in-kernel history fold."""
self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt)) return {
"ring_state": self.state,
"hist_idx": self.idx,
"fp8_max": FP8_MAX[fmt],
"pow2_margin": float(2**self.recipe.margin),
}
class FP8TensorMeta: class FP8TensorMeta:
@@ -322,11 +329,11 @@ def fp8_linear_forward(
Composed from the two stateless primitives: quantize x/w with the active 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. scales, run the pre-quantized GEMM with the bias fused into its epilogue.
Delayed scaling folds Delayed scaling lets the quantize kernel fold the fused amax into the
the returned amax into the history ring and publishes the next scale; history ring and publish the next scale in its own last block; dynamic
dynamic scaling measures the current amax itself. Training quantizes the scaling measures the current amax itself. Training quantizes the weight
weight every step (the optimizer bumps its version, so there is no cast every step (the optimizer bumps its version, so there is no cast cache,
cache, matching ``cached_cast``-less behavior). matching ``cached_cast``-less behavior).
""" """
state = fp8_state() state = fp8_state()
if cfg is None: if cfg is None:
@@ -351,19 +358,19 @@ def fp8_linear_forward(
if not meta.x.initialized: if not meta.x.initialized:
meta.x.seed(x, fmt) meta.x.seed(x, fmt)
sx, sw = meta.x.scale.clone(), meta.w.scale.clone() 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): if _is_fp8(w.dtype):
w8, amax_w = w, None w8 = w
else: else:
w8, amax_w = quantize(w, sw.reciprocal(), fmt) w8, _ = quantize(w, sw.reciprocal(), fmt, **meta.w.fold_args(fmt))
out = mm_fp8( out = mm_fp8(
x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
).reshape(*x.shape[:-1], w.size(0)) ).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() meta.x.advance()
if amax_w is not None: if not _is_fp8(w.dtype):
meta.w.advance() meta.w.advance()
return out, sx, sw 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 # 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 # 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 # crosswise kernel in the training path. g is consumed in both
# orientations, so one dual-layout pass feeds both. # orientations, so quantize_dual's single pass feeds both.
g8, g8T, amax_g = quantize(g2, sg.reciprocal(), fmt, layout=2) # The g quantize folds the gradient amax into its ring in-kernel;
x8T, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, layout=1) # 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): if _is_fp8(w.dtype):
# Pre-quantized weight has no transposed copy: keep the swap # Pre-quantized weight has no transposed copy: keep the swap
# path for grad_x (grad_w is unaffected). # path for grad_x (grad_w is unaffected).
grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape) grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape)
else: 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_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 grad_w = mm_fp8(g8T, x8T, sg * sx, trans_b=True) # g8.T @ x8
# bias-free linears must not pay the column-sum # bias-free linears must not pay the column-sum
# reduce: g2.sum(0) is another full read of the gradient. # 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 grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None
if not ctx.is_dynamic: if not ctx.is_dynamic:
meta.g.update(amax_g, fmt)
meta.g.advance() meta.g.advance()
return grad_x, grad_w, grad_b 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). """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) - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
Scale semantics: scales are *quantization steps* — the value divided out when ``scale`` is the quantization multiplier (device scalar); ``fmt`` is
quantizing (``x8 = x / scale``). Every primitive computes its own inverse ``"e4m3"`` or ``"e5m2"``. ``amax`` values are *returned*, never passed as
internally; callers never pass ``scale_inv``. ``amax`` values are *returned*, output arguments.
never passed as output arguments. ``fmt`` is ``"e4m3"`` or ``"e5m2"``.
Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``; Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
this module is stateless. this module is stateless.
@@ -17,7 +23,6 @@ this module is stateless.
from typing import Optional, Tuple from typing import Optional, Tuple
import torch import torch
from torch.library import custom_op
from astrai.extension.loader import get_module 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')") 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( def quantize(
x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3", layout: int = 0 x: torch.Tensor,
) -> tuple: 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. """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax.
``scale`` is the quantization multiplier (device scalar); ``fmt`` selects ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. ``layout`` E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
picks the output orientation: 0 = row-major ``(x8, amax)``; 1 = ``transposed=True`` swaps ``x8`` for ``x8T``, the ``[cols][rows]``
transposed ``[cols][rows]`` ``(x8T, amax)`` — the K-contiguous operand row-major transpose of the quantized input — the K-contiguous operand
orientation NT GEMMs want; 2 = both from one read ``(x8, x8T, amax)`` orientation NT GEMMs want — at the same 2-tuple arity.
(for tensors consumed in both orientations, e.g. backward ``g``).
``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 return get_module("fp8_ops").quantize(
# 512-wide GEMM): real CUDA tensors of a supported dtype go straight to x,
# the extension. Fake/subclass tensors and non-CUDA inputs keep the scale,
# custom_op route so torch.compile / meta / fake-tensor tracing and the _fmt_int(fmt),
# CPU fallback behave exactly as before. transposed,
if ( ring_state,
type(x) is torch.Tensor hist_idx,
and x.is_cuda fp8_max,
and x.dtype in _QUANT_INPUT_DTYPES pow2_margin,
and fmt in _FMT_TO_INT )
):
return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt], layout)
if layout == 0: def quantize_dual(
return fp8_quantize(x, scale, _fmt_int(fmt)) x: torch.Tensor,
if layout == 1: scale: torch.Tensor,
return fp8_quantize_t(x, scale, _fmt_int(fmt)) fmt: str = "e4m3",
return fp8_quantize_dual(x, scale, _fmt_int(fmt)) 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( def mm_fp8(
@@ -237,15 +113,4 @@ def mm_fp8(
kernel epilogue in fp32 — no separate elementwise pass. The result is kernel epilogue in fp32 — no separate elementwise pass. The result is
BF16; FP8 output is a separate quantize operation. BF16; FP8 output is a separate quantize operation.
""" """
# Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
# 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)
csrc/kernels/fp8/common.h
+25 -4
View File
@@ -54,19 +54,40 @@ struct Fp8GemmTraits {
"warp tile must be a multiple of the m16n8 MMA shape"); "warp tile must be a multiple of the m16n8 MMA shape");
}; };
// Quantize output orientation: RowMajor = x8 only; Transposed = the
// [cols][rows] x8T only; Dual = both from a single read. Transposed/Dual
// produce K-contiguous operands so crosswise consumers (backward
// grad_x / grad_w) route through the NT fast path.
enum class QuantLayout : int {
RowMajor = 0,
Transposed = 1,
Dual = 2,
};
// Quantize-kernel parameter POD: float input -> FP8 with fused amax. // Quantize-kernel parameter POD: float input -> FP8 with fused amax.
struct FP8QuantizeParams { struct FP8QuantizeParams {
const void* __restrict__ input_ptr = nullptr; const void* __restrict__ input_ptr = nullptr;
void* __restrict__ output_ptr = nullptr; void* __restrict__ output_ptr = nullptr;
void* __restrict__ output_transposed_ptr = nullptr; // [cols][rows] void* __restrict__ output_transposed_ptr = nullptr; // [cols][rows]
// Output layout: 0 = row-major only, 1 = transposed only, 2 = both from QuantLayout out_layout = QuantLayout::RowMajor;
// a single read. Modes 1/2 produce K-contiguous operands so crosswise
// consumers (backward grad_x / grad_w) route through the NT fast path.
int out_layout = 0;
const float* __restrict__ scale = nullptr; // device multiplier const float* __restrict__ scale = nullptr; // device multiplier
float* __restrict__ amax = nullptr; // raw-domain max out float* __restrict__ amax = nullptr; // raw-domain max out
// Optional delayed-scaling ring fold: when fold_ring is set, the kernel's
// last-finishing block folds the final amax into hist[hist_idx], reduces
// the window and publishes the next scale — replacing the host-side
// update chain. amax then points at a persistent self-cleaning slot
// (zeroed by the same last block) inside the caller's ring state.
bool fold_ring = false;
float* __restrict__ hist = nullptr; // [hist_len] amax history window
float* __restrict__ scale_out = nullptr;
unsigned int* __restrict__ done = nullptr; // block-completion counter
int hist_len = 0;
int hist_idx = 0;
float fp8_max = 448.0f; // scale = max(hist) / fp8_max / pow2_margin
float pow2_margin = 1.0f;
// Element count (elementwise kernel); the tiled kernel views the same // Element count (elementwise kernel); the tiled kernel views the same
// buffer as [rows][cols] row-major. // buffer as [rows][cols] row-major.
int total = 0; int total = 0;
csrc/kernels/fp8/ops.cu
+109 -56
View File
@@ -1,4 +1,4 @@
// CUDA bindings for the two stateless FP8 primitives. // CUDA bindings for the stateless FP8 quantize/GEMM primitives.
#include <ATen/cuda/CUDAContext.h> #include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h> #include <c10/cuda/CUDAGuard.h>
@@ -94,13 +94,17 @@ void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p,
launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream); launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
} }
} // namespace // Shared binding body for the two quantize entry points: RowMajor /
// Transposed (single output) serve quantize(), Dual (both orientations from
// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return), // one read) serves quantize_dual(). A ring tensor switches
// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a // on the in-kernel delayed-scaling fold: state layout
// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path. // [hist n | scale | legacy | amax | done-as-int], and the returned amax is
py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt, // the (self-cleaned) persistent slot. Without it, amax is reduced into a
int64_t layout) { // fresh buffer armed by a driver memset — cheaper than the zeros() fill
// kernel.
py::object quantize_impl(torch::Tensor x, torch::Tensor scale, int64_t fmt,
QuantLayout layout, py::object ring, int64_t hist_idx,
double fp8_max, double pow2_margin) {
TORCH_CHECK(x.is_cuda(), "CUDA tensors required"); TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
TORCH_CHECK(x.scalar_type() == torch::kBFloat16 || TORCH_CHECK(x.scalar_type() == torch::kBFloat16 ||
x.scalar_type() == torch::kHalf || x.scalar_type() == torch::kHalf ||
@@ -109,9 +113,7 @@ py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
TORCH_CHECK(fmt == static_cast<int64_t>(FP8Format::E4M3) || TORCH_CHECK(fmt == static_cast<int64_t>(FP8Format::E4M3) ||
fmt == static_cast<int64_t>(FP8Format::E5M2), fmt == static_cast<int64_t>(FP8Format::E5M2),
"unsupported quantization type: expected E4M3 (0) or E5M2 (1)"); "unsupported quantization type: expected E4M3 (0) or E5M2 (1)");
TORCH_CHECK(layout >= 0 && layout <= 2, TORCH_CHECK(layout == QuantLayout::RowMajor || x.dim() >= 2,
"layout must be 0 (row-major), 1 (transposed) or 2 (both)");
TORCH_CHECK(layout == 0 || x.dim() >= 2,
"transposed quantize layouts need a 2D+ tensor"); "transposed quantize layouts need a 2D+ tensor");
check_scale(scale, x); check_scale(scale, x);
check_fp8_device(x); check_fp8_device(x);
@@ -120,41 +122,94 @@ py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
auto input = x.contiguous(); auto input = x.contiguous();
auto out_opts = input.options().dtype( auto out_opts = input.options().dtype(
fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn); fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn);
// amax is reduced via atomicMax of non-negative values; a driver memset torch::Tensor amax;
// arms it cheaper than the zeros() fill kernel (one fewer tensor-op float *ring_hist = nullptr, *ring_scale_out = nullptr;
// dispatch + kernel launch on every quantize call). unsigned int* ring_done = nullptr;
auto amax = torch::empty({1}, input.options().dtype(torch::kFloat32)); int ring_len = 0;
cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream()); if (!ring.is_none()) {
auto st = ring.cast<torch::Tensor>();
TORCH_CHECK(st.is_cuda() && st.dim() == 1 &&
st.scalar_type() == torch::kFloat32,
"ring state must be a 1D float32 CUDA tensor");
const int64_t n = st.numel() - 4;
TORCH_CHECK(n > 0 && hist_idx >= 0 && hist_idx < n,
"ring state too small or hist_idx out of range");
float* base = st.data_ptr<float>();
amax = st.narrow(0, n + 2, 1);
ring_hist = base;
ring_scale_out = base + n;
ring_done = reinterpret_cast<unsigned int*>(base + n + 3);
ring_len = static_cast<int>(n);
} else {
amax = torch::empty({1}, input.options().dtype(torch::kFloat32));
cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream());
}
FP8QuantizeParams p; FP8QuantizeParams p;
p.input_ptr = input.data_ptr(); p.input_ptr = input.data_ptr();
p.scale = scale.data_ptr<float>(); p.scale = scale.data_ptr<float>();
p.amax = amax.data_ptr<float>(); p.amax = amax.data_ptr<float>();
if (ring_hist) {
p.fold_ring = true;
p.hist = ring_hist;
p.scale_out = ring_scale_out;
p.done = ring_done;
p.hist_len = ring_len;
p.hist_idx = static_cast<int>(hist_idx);
p.fp8_max = static_cast<float>(fp8_max);
p.pow2_margin = static_cast<float>(pow2_margin);
}
p.total = static_cast<int>(input.numel()); p.total = static_cast<int>(input.numel());
p.out_layout = static_cast<int>(layout); p.out_layout = layout;
p.rows = static_cast<int>(input.size(-2)); p.rows = static_cast<int>(input.size(-2));
p.cols = static_cast<int>(input.size(-1)); p.cols = static_cast<int>(input.size(-1));
torch::Tensor output, output_t; torch::Tensor output, output_t;
if (layout == 0 || layout == 2) { if (layout != QuantLayout::Transposed) {
output = torch::empty_like(input, out_opts); output = torch::empty_like(input, out_opts);
p.output_ptr = output.data_ptr(); p.output_ptr = output.data_ptr();
} }
if (layout >= 1) { if (layout != QuantLayout::RowMajor) {
output_t = torch::empty({input.size(-1), input.size(-2)}, out_opts); output_t = torch::empty({input.size(-1), input.size(-2)}, out_opts);
p.output_transposed_ptr = output_t.data_ptr(); p.output_transposed_ptr = output_t.data_ptr();
} }
const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2); const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
if (layout != 0) if (layout == QuantLayout::RowMajor)
launch_quantize_for<true>(input, p, e5m2, stream.stream());
else
launch_quantize_for<false>(input, p, e5m2, stream.stream()); launch_quantize_for<false>(input, p, e5m2, stream.stream());
else
launch_quantize_for<true>(input, p, e5m2, stream.stream());
C10_CUDA_CHECK(cudaGetLastError()); C10_CUDA_CHECK(cudaGetLastError());
if (layout == 2) return py::make_tuple(output, output_t, amax); if (layout == QuantLayout::Dual)
return py::make_tuple(layout == 1 ? output_t : output, amax); return py::make_tuple(output, output_t, amax);
return py::make_tuple(
layout == QuantLayout::Transposed ? output_t : output, amax);
}
} // namespace
// Single-orientation quantize binding: row-major x8, or its [cols][rows]
// transpose when transposed is set — the K-contiguous operand orientation
// NT GEMMs want. Returns (x8|x8T, amax).
py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
bool transposed, py::object ring, int64_t hist_idx,
double fp8_max, double pow2_margin) {
const QuantLayout layout =
transposed ? QuantLayout::Transposed : QuantLayout::RowMajor;
return quantize_impl(x, scale, fmt, layout, ring, hist_idx, fp8_max,
pow2_margin);
}
// Dual-orientation quantize binding: one read of x produces both the
// row-major x8 and its transpose (plus amax), for tensors consumed by GEMMs
// in both orientations (backward g). Returns (x8, x8T, amax).
py::object quantize_dual(torch::Tensor x, torch::Tensor scale, int64_t fmt,
py::object ring, int64_t hist_idx, double fp8_max,
double pow2_margin) {
return quantize_impl(x, scale, fmt, QuantLayout::Dual, ring, hist_idx,
fp8_max, pow2_margin);
} }
torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale, torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
int64_t trans_a, int64_t trans_b, torch::Tensor bias) { bool trans_a, bool trans_b, py::object bias) {
TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required"); TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn || TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn ||
a.scalar_type() == torch::kFloat8_e5m2, a.scalar_type() == torch::kFloat8_e5m2,
@@ -164,6 +219,18 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
(b.dim() == 2 || b.dim() == 3), (b.dim() == 2 || b.dim() == 3),
"a and b must be 2D or 3D (batched)"); "a and b must be 2D or 3D (batched)");
TORCH_CHECK(a.device() == b.device(), "a and b must share device"); TORCH_CHECK(a.device() == b.device(), "a and b must share device");
// Python None and an omitted argument both mean "no bias" — an undefined
// tensor below. (py::isinstance<torch::Tensor> is false for real tensors
// here — torch's caster registers no pybind type info — so validate by
// attempting the cast itself.)
torch::Tensor bias_t;
if (!bias.is_none()) {
try {
bias_t = bias.cast<torch::Tensor>();
} catch (const py::cast_error&) {
TORCH_CHECK(false, "bias must be a torch.Tensor or None");
}
}
check_scale(scale, a); check_scale(scale, a);
check_fp8_device(a); check_fp8_device(a);
const at::cuda::OptionalCUDAGuard guard(a.device()); const at::cuda::OptionalCUDAGuard guard(a.device());
@@ -180,10 +247,8 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
torch::Tensor a_st, b_st; torch::Tensor a_st, b_st;
int64_t a_ld, b_ld, a_bstride, b_bstride; int64_t a_ld, b_ld, a_bstride, b_bstride;
const bool tag_a = const bool tag_a = resolve_operand(a, trans_a, a_ld, a_bstride, a_st);
resolve_operand(a, trans_a != 0, a_ld, a_bstride, a_st); const bool tag_b = resolve_operand(b, trans_b, b_ld, b_bstride, b_st);
const bool tag_b =
resolve_operand(b, trans_b != 0, b_ld, b_bstride, b_st);
// GEMM dims from the user flags; storage layout never swaps them. // GEMM dims from the user flags; storage layout never swaps them.
const int64_t m = trans_a ? a.size(-1) : a.size(-2); const int64_t m = trans_a ? a.size(-1) : a.size(-2);
const int64_t k = trans_a ? a.size(-2) : a.size(-1); const int64_t k = trans_a ? a.size(-2) : a.size(-1);
@@ -207,13 +272,13 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
p.b_ld = static_cast<int>(b_ld); p.b_ld = static_cast<int>(b_ld);
// Fused epilogue bias (bf16, broadcast over rows and batches). An // Fused epilogue bias (bf16, broadcast over rows and batches). An
// undefined or 0-element tensor keeps the plain scaled output. // undefined or 0-element tensor keeps the plain scaled output.
if (bias.defined() && bias.numel() > 0) { if (bias_t.defined() && bias_t.numel() > 0) {
TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kBFloat16, TORCH_CHECK(bias_t.is_cuda() && bias_t.scalar_type() == torch::kBFloat16,
"fp8 gemm bias must be a CUDA bf16 tensor"); "fp8 gemm bias must be a CUDA bf16 tensor");
TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n, TORCH_CHECK(bias_t.dim() == 1 && bias_t.size(0) == n,
"fp8 gemm bias must be 1D of length n=", n); "fp8 gemm bias must be 1D of length n=", n);
TORCH_CHECK(bias.is_contiguous(), "fp8 gemm bias must be contiguous"); TORCH_CHECK(bias_t.is_contiguous(), "fp8 gemm bias must be contiguous");
p.bias_ptr = bias.data_ptr(); p.bias_ptr = bias_t.data_ptr();
} }
p.batch = static_cast<int>(batch); p.batch = static_cast<int>(batch);
p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride; p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride;
@@ -227,28 +292,16 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
return output; return output;
} }
// mm_fp8 binding: Python None and an omitted argument both mean "no bias",
// so every Python layer can pass its bias argument through untouched.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("quantize", &quantize, py::arg("x"), py::arg("scale"), m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
py::arg("fmt"), py::arg("layout") = 0); py::arg("fmt"), py::arg("transposed") = false,
m.def( py::arg("ring") = py::none(), py::arg("hist_idx") = 0,
"mm_fp8", py::arg("fp8_max") = 448.0, py::arg("pow2_margin") = 1.0);
[](torch::Tensor a, torch::Tensor b, torch::Tensor scale, m.def("quantize_dual", &quantize_dual, py::arg("x"), py::arg("scale"),
int64_t trans_a, int64_t trans_b, py::object bias) { py::arg("fmt"), py::arg("ring") = py::none(),
torch::Tensor t; py::arg("hist_idx") = 0, py::arg("fp8_max") = 448.0,
if (!bias.is_none()) { py::arg("pow2_margin") = 1.0);
// (py::isinstance<torch::Tensor> is false for real tensors m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
// here — torch's caster registers no pybind type info — so py::arg("trans_a") = false, py::arg("trans_b") = false,
// validate by attempting the cast itself.) py::arg("bias") = py::none());
try {
t = bias.cast<torch::Tensor>();
} catch (const py::cast_error&) {
TORCH_CHECK(false, "bias must be a torch.Tensor or None");
}
}
return mm_fp8(a, b, scale, trans_a, trans_b, t);
},
py::arg("a"), py::arg("b"), py::arg("scale"), py::arg("trans_a") = 0,
py::arg("trans_b") = 0, py::arg("bias") = py::none());
} }
csrc/kernels/fp8/quantize.cuh
+28 -10
View File
@@ -108,8 +108,13 @@ __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
// Block-wide amax reduce -> one atomic per block: warp-reduce, park one // Block-wide amax reduce -> one atomic per block: warp-reduce, park one
// value per warp, thread 0 folds. kWarps must cover the block's warp count. // value per warp, thread 0 folds. kWarps must cover the block's warp count.
// With p.fold_ring, the last-finishing block additionally folds the final
// amax into the history window and publishes the next scale (atomicAdd
// ticket + fences), re-zeroing the amax slot and the counter for the next
// launch — the host-side delayed-scaling update chain disappears.
template <int kWarps> template <int kWarps>
__device__ __forceinline__ void publish_amax(float* amax, float v) { __device__ __forceinline__ void publish_amax(const FP8QuantizeParams& p,
float v) {
v = warp_reduce_max(v); v = warp_reduce_max(v);
__shared__ float slots[kWarps]; __shared__ float slots[kWarps];
const int tid = threadIdx.y * blockDim.x + threadIdx.x; const int tid = threadIdx.y * blockDim.x + threadIdx.x;
@@ -118,12 +123,23 @@ __device__ __forceinline__ void publish_amax(float* amax, float v) {
if (tid == 0) { if (tid == 0) {
#pragma unroll #pragma unroll
for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]); for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
atomic_max_float(amax, v); atomic_max_float(p.amax, v);
if (!p.fold_ring) return;
__threadfence();
const unsigned int ticket = atomicAdd(p.done, 1u);
__threadfence();
if (ticket != gridDim.x - 1u) return;
p.hist[p.hist_idx] = *p.amax;
float peak = p.hist[0];
for (int i = 1; i < p.hist_len; ++i) peak = fmaxf(peak, p.hist[i]);
*p.scale_out = fmaxf(peak / p.fp8_max / p.pow2_margin, 1e-12f);
*p.amax = 0.0f;
*p.done = 0u;
} }
} }
// Elementwise quantize kernel (out_layout 0): vectorized 16B loads -> fp8 // Elementwise quantize kernel (QuantLayout::RowMajor): vectorized 16B loads
// stores, fused amax over raw values. // -> fp8 stores, fused amax over raw values.
template <FP8Format Fmt, typename InT> template <FP8Format Fmt, typename InT>
__global__ void fp8_quantize_kernel(FP8QuantizeParams p) { __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
const float mult = *p.scale; const float mult = *p.scale;
@@ -174,10 +190,11 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
local_amax = fmaxf(local_amax, fabsf(v)); local_amax = fmaxf(local_amax, fabsf(v));
x8[i] = cvt_fp8<Fmt>(v * mult); x8[i] = cvt_fp8<Fmt>(v * mult);
} }
if (p.amax) publish_amax<8>(p.amax, local_amax); if (p.amax) publish_amax<8>(p, local_amax);
} }
// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] input // Tiled transpose quantize (QuantLayout::Transposed/Dual): reads the
// [rows][cols] input
// once and writes the fp8 bytes transposed ([cols][rows], so the contract // once and writes the fp8 bytes transposed ([cols][rows], so the contract
// dim lands K-contiguous for NT GEMM operands) and, in mode 2, the row-major // dim lands K-contiguous for NT GEMM operands) and, in mode 2, the row-major
// copy too. 64x32 tiles, one native pair load per row (a full 128B warp // copy too. 64x32 tiles, one native pair load per row (a full 128B warp
@@ -233,7 +250,7 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
} }
} }
} }
if (p.out_layout == 2) { if (p.out_layout == QuantLayout::Dual) {
uint8_t* out = static_cast<uint8_t*>(p.output_ptr); uint8_t* out = static_cast<uint8_t*>(p.output_ptr);
#pragma unroll #pragma unroll
for (int j = 0; j < 4; ++j) for (int j = 0; j < 4; ++j)
@@ -266,11 +283,12 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] = out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
tile[threadIdx.y * 8 + i][threadIdx.x]; tile[threadIdx.y * 8 + i][threadIdx.x];
} }
if (p.amax) publish_amax<8>(p.amax, local_amax); if (p.amax) publish_amax<8>(p, local_amax);
} }
// Unified quantize launcher: Tiled selects the transpose kernel (out_layout // Unified quantize launcher: Tiled selects the transpose kernel
// 1/2) over the vectorized elementwise one. The transpose kernel vectorizes // (QuantLayout::Transposed/Dual) over the vectorized elementwise one. The
// transpose kernel vectorizes
// pair loads in-kernel and falls back to scalar loads at unaligned/ragged // pair loads in-kernel and falls back to scalar loads at unaligned/ragged
// rows, so the host side picks only the grid. // rows, so the host side picks only the grid.
template <FP8Format Fmt, typename InT, bool Tiled = false> template <FP8Format Fmt, typename InT, bool Tiled = false>
tests/extension/test_fp8_mma.py
+60 -24
View File
@@ -2,8 +2,9 @@
The kernel-level tests exercise the two stateless primitives (``quantize`` for The kernel-level tests exercise the two stateless primitives (``quantize`` for
bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed
operands); the policy-level tests (recipes, autocast context, per-tensor meta, operands); the policy-level tests (recipes, autocast context, per-tensor
CPU fallbacks of the custom ops) run without a GPU. meta) run without a GPU. The primitives themselves are CUDA-only
(attention-style direct wrappers — no torch.library dispatch layer).
""" """
import threading import threading
@@ -24,7 +25,7 @@ from astrai.extension.fp8 import (
fp8_linear_enabled, fp8_linear_enabled,
fp8_state, fp8_state,
) )
from astrai.extension.ops.fp8 import mm_fp8, quantize from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
from tests.conftest import skip_no_fp8 from tests.conftest import skip_no_fp8
@@ -448,37 +449,72 @@ def test_fp8_tensor_meta_delayed_update():
meta.w.seed(w, "e4m3") meta.w.seed(w, "e4m3")
assert meta.w.initialized assert meta.w.initialized
torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1)) torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1))
# [hist | scale] packing: views alias the single state buffer. # [hist | scale | legacy | amax | done] packing: views alias one buffer.
assert meta.w.state.numel() == 4 + 2 assert meta.w.state.numel() == 4 + 4
assert meta.w.hist.data_ptr() == meta.w.state.data_ptr() assert meta.w.hist.data_ptr() == meta.w.state.data_ptr()
assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr() assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr()
meta.w.advance() meta.w.advance()
assert meta.w.idx == 1 assert meta.w.idx == 1
# update folds a fresh amax into the window and publishes the next scale # fold_args hands the kernel the buffer, the slot and the recipe constants
amax = torch.tensor([8.0]) args = meta.w.fold_args("e4m3")
meta.w.update(amax, "e4m3") assert args["ring_state"] is meta.w.state and args["hist_idx"] == 1
torch.testing.assert_close(meta.w.scale, torch.tensor([8.0 / 448.0])) assert args["fp8_max"] == 448.0 and args["pow2_margin"] == 1.0
def test_quantize_cpu_fallback(): @skip_no_fp8
"""CPU fallback of the quantize primitive (scale semantics + amax).""" @pytest.mark.parametrize("fmt", ["e4m3", "e5m2"])
x = torch.randn(16, 32, dtype=torch.bfloat16) def test_quantize_dual_and_transposed_orientations(fmt):
scale = torch.tensor([0.5]) # quantize multiplier """quantize_dual yields both orientations from one read; quantize's
x8, amax = quantize(x, scale, "e4m3") transposed switch keeps the 2-tuple arity with the [cols][rows] layout."""
assert x8.dtype == torch.float8_e4m3fn torch.manual_seed(11)
ref = (x.float() * 0.5).to(torch.float8_e4m3fn) x = torch.randn(37, 67, device="cuda", dtype=torch.bfloat16) * 3
assert torch.equal(x8, ref) mult = _scale(x).reciprocal()
x8, amax = quantize(x, mult, fmt)
x8T, _ = quantize(x, mult, fmt, transposed=True)
d8, d8T, _ = quantize_dual(x, mult, fmt)
assert x8T.shape == (67, 37)
assert torch.equal(x8.view(torch.uint8), d8.view(torch.uint8))
assert torch.equal(x8T.view(torch.uint8), d8T.view(torch.uint8))
assert torch.equal(x8T.t().contiguous().view(torch.uint8), x8.view(torch.uint8))
torch.testing.assert_close(amax, x.abs().amax().float().reshape(1)) torch.testing.assert_close(amax, x.abs().amax().float().reshape(1))
def test_mm_fp8_cpu_fallback(): @skip_no_fp8
a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn) @pytest.mark.parametrize("fmt,fmax", [("e4m3", 448.0), ("e5m2", 57344.0)])
b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn) @pytest.mark.parametrize("margin", [0, 1])
scale = torch.tensor([1.0]) def test_quantize_ring_fold_matches_host_update(fmt, fmax, margin):
out = mm_fp8(a8, b8, scale) """The in-kernel delayed-scaling fold matches a host-side reference."""
ref = (a8.float() @ b8.float() * 1.0).to(torch.bfloat16) dev = torch.device("cuda")
torch.testing.assert_close(out, ref) n, idx = 4, 2
torch.manual_seed(3)
x = torch.randn(128, 96, dtype=torch.bfloat16, device=dev) * 3
mult = torch.tensor([0.01], device=dev)
pow2m = float(2**margin)
# Reference: legacy quantize + the host fold it used to return amax for.
x8_ref, amax = quantize(x, mult, fmt)
hist = torch.full((n,), 1.0, device=dev)
hist[idx] = amax.to(torch.float32)
scale = (hist.max() / fmax / pow2m).clamp_min(1e-12).reshape(1)
# Fused: same window, fold inside the quantize kernel's last block.
ring = torch.zeros(n + 4, device=dev)
ring[:n].fill_(1.0)
x8, _ = quantize(
x,
mult,
fmt,
ring_state=ring,
hist_idx=idx,
fp8_max=fmax,
pow2_margin=pow2m,
)
assert torch.equal(x8.view(torch.uint8), x8_ref.view(torch.uint8))
torch.testing.assert_close(ring[:n], hist, rtol=0, atol=0)
torch.testing.assert_close(ring[n : n + 1], scale, rtol=0, atol=0)
assert float(ring[n + 2]) == 0.0 # amax slot self-cleaned
assert int(ring[n + 3].view(torch.int32)) == 0 # done counter reset
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------