From 962c10c52b78f604879ac2c5eb72b7c241c2fab4 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 29 Aug 2026 09:16:22 +0800 Subject: [PATCH] perf: fold the delayed-scaling ring update into the quantize kernel - 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. --- astrai/extension/fp8.py | 63 ++++---- astrai/extension/ops/fp8.py | 265 ++++++++------------------------ csrc/kernels/fp8/common.h | 29 +++- csrc/kernels/fp8/ops.cu | 165 +++++++++++++------- csrc/kernels/fp8/quantize.cuh | 38 +++-- tests/extension/test_fp8_mma.py | 84 +++++++--- 6 files changed, 324 insertions(+), 320 deletions(-) diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index 65194a6..f2ca308 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -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 diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index abe2427..8896744 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -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) diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 342fb2d..64c89b8 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -54,19 +54,40 @@ struct Fp8GemmTraits { "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. struct FP8QuantizeParams { const void* __restrict__ input_ptr = nullptr; void* __restrict__ output_ptr = nullptr; void* __restrict__ output_transposed_ptr = nullptr; // [cols][rows] - // Output layout: 0 = row-major only, 1 = transposed only, 2 = both from - // 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; + QuantLayout out_layout = QuantLayout::RowMajor; const float* __restrict__ scale = nullptr; // device multiplier 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 // buffer as [rows][cols] row-major. int total = 0; diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index 21da867..14bfc2e 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -1,4 +1,4 @@ -// CUDA bindings for the two stateless FP8 primitives. +// CUDA bindings for the stateless FP8 quantize/GEMM primitives. #include #include @@ -94,13 +94,17 @@ void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p, launch_for_dtype(x, p, stream); } -} // namespace - -// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return), -// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a -// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path. -py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt, - int64_t layout) { +// Shared binding body for the two quantize entry points: RowMajor / +// Transposed (single output) serve quantize(), Dual (both orientations from +// one read) serves quantize_dual(). A ring tensor switches +// on the in-kernel delayed-scaling fold: state layout +// [hist n | scale | legacy | amax | done-as-int], and the returned amax is +// the (self-cleaned) persistent slot. Without it, amax is reduced into a +// 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.scalar_type() == torch::kBFloat16 || 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(FP8Format::E4M3) || fmt == static_cast(FP8Format::E5M2), "unsupported quantization type: expected E4M3 (0) or E5M2 (1)"); - TORCH_CHECK(layout >= 0 && layout <= 2, - "layout must be 0 (row-major), 1 (transposed) or 2 (both)"); - TORCH_CHECK(layout == 0 || x.dim() >= 2, + TORCH_CHECK(layout == QuantLayout::RowMajor || x.dim() >= 2, "transposed quantize layouts need a 2D+ tensor"); check_scale(scale, 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 out_opts = input.options().dtype( fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn); - // amax is reduced via atomicMax of non-negative values; a driver memset - // arms it cheaper than the zeros() fill kernel (one fewer tensor-op - // dispatch + kernel launch on every quantize call). - auto amax = torch::empty({1}, input.options().dtype(torch::kFloat32)); - cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream()); + torch::Tensor amax; + float *ring_hist = nullptr, *ring_scale_out = nullptr; + unsigned int* ring_done = nullptr; + int ring_len = 0; + if (!ring.is_none()) { + auto st = ring.cast(); + 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(); + amax = st.narrow(0, n + 2, 1); + ring_hist = base; + ring_scale_out = base + n; + ring_done = reinterpret_cast(base + n + 3); + ring_len = static_cast(n); + } else { + amax = torch::empty({1}, input.options().dtype(torch::kFloat32)); + cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream()); + } FP8QuantizeParams p; p.input_ptr = input.data_ptr(); p.scale = scale.data_ptr(); p.amax = amax.data_ptr(); + 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(hist_idx); + p.fp8_max = static_cast(fp8_max); + p.pow2_margin = static_cast(pow2_margin); + } p.total = static_cast(input.numel()); - p.out_layout = static_cast(layout); + p.out_layout = layout; p.rows = static_cast(input.size(-2)); p.cols = static_cast(input.size(-1)); torch::Tensor output, output_t; - if (layout == 0 || layout == 2) { + if (layout != QuantLayout::Transposed) { output = torch::empty_like(input, out_opts); 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); p.output_transposed_ptr = output_t.data_ptr(); } const bool e5m2 = fmt == static_cast(FP8Format::E5M2); - if (layout != 0) - launch_quantize_for(input, p, e5m2, stream.stream()); - else + if (layout == QuantLayout::RowMajor) launch_quantize_for(input, p, e5m2, stream.stream()); + else + launch_quantize_for(input, p, e5m2, stream.stream()); C10_CUDA_CHECK(cudaGetLastError()); - if (layout == 2) return py::make_tuple(output, output_t, amax); - return py::make_tuple(layout == 1 ? output_t : output, amax); + if (layout == QuantLayout::Dual) + 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, - 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.scalar_type() == torch::kFloat8_e4m3fn || 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), "a and b must be 2D or 3D (batched)"); 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 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(); + } catch (const py::cast_error&) { + TORCH_CHECK(false, "bias must be a torch.Tensor or None"); + } + } check_scale(scale, a); check_fp8_device(a); 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; int64_t a_ld, b_ld, a_bstride, b_bstride; - const bool tag_a = - resolve_operand(a, trans_a != 0, a_ld, a_bstride, a_st); - const bool tag_b = - resolve_operand(b, trans_b != 0, b_ld, b_bstride, b_st); + const bool tag_a = resolve_operand(a, trans_a, a_ld, a_bstride, a_st); + const bool tag_b = resolve_operand(b, trans_b, b_ld, b_bstride, b_st); // 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 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(b_ld); // Fused epilogue bias (bf16, broadcast over rows and batches). An // undefined or 0-element tensor keeps the plain scaled output. - if (bias.defined() && bias.numel() > 0) { - TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kBFloat16, + if (bias_t.defined() && bias_t.numel() > 0) { + TORCH_CHECK(bias_t.is_cuda() && bias_t.scalar_type() == torch::kBFloat16, "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); - TORCH_CHECK(bias.is_contiguous(), "fp8 gemm bias must be contiguous"); - p.bias_ptr = bias.data_ptr(); + TORCH_CHECK(bias_t.is_contiguous(), "fp8 gemm bias must be contiguous"); + p.bias_ptr = bias_t.data_ptr(); } p.batch = static_cast(batch); 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; } -// 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) { m.def("quantize", &quantize, py::arg("x"), py::arg("scale"), - py::arg("fmt"), py::arg("layout") = 0); - m.def( - "mm_fp8", - [](torch::Tensor a, torch::Tensor b, torch::Tensor scale, - int64_t trans_a, int64_t trans_b, py::object bias) { - torch::Tensor t; - if (!bias.is_none()) { - // (py::isinstance is false for real tensors - // here — torch's caster registers no pybind type info — so - // validate by attempting the cast itself.) - try { - t = bias.cast(); - } 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()); + py::arg("fmt"), py::arg("transposed") = false, + py::arg("ring") = py::none(), py::arg("hist_idx") = 0, + py::arg("fp8_max") = 448.0, py::arg("pow2_margin") = 1.0); + m.def("quantize_dual", &quantize_dual, py::arg("x"), py::arg("scale"), + py::arg("fmt"), py::arg("ring") = py::none(), + py::arg("hist_idx") = 0, py::arg("fp8_max") = 448.0, + py::arg("pow2_margin") = 1.0); + m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"), + py::arg("trans_a") = false, py::arg("trans_b") = false, + py::arg("bias") = py::none()); } diff --git a/csrc/kernels/fp8/quantize.cuh b/csrc/kernels/fp8/quantize.cuh index 1d51189..97f2bbd 100644 --- a/csrc/kernels/fp8/quantize.cuh +++ b/csrc/kernels/fp8/quantize.cuh @@ -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 // 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 -__device__ __forceinline__ void publish_amax(float* amax, float v) { +__device__ __forceinline__ void publish_amax(const FP8QuantizeParams& p, + float v) { v = warp_reduce_max(v); __shared__ float slots[kWarps]; 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) { #pragma unroll 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 -// stores, fused amax over raw values. +// Elementwise quantize kernel (QuantLayout::RowMajor): vectorized 16B loads +// -> fp8 stores, fused amax over raw values. template __global__ void fp8_quantize_kernel(FP8QuantizeParams p) { const float mult = *p.scale; @@ -174,10 +190,11 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) { local_amax = fmaxf(local_amax, fabsf(v)); x8[i] = cvt_fp8(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 // 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 @@ -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(p.output_ptr); #pragma unroll 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] = 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 -// 1/2) over the vectorized elementwise one. The transpose kernel vectorizes +// Unified quantize launcher: Tiled selects the transpose kernel +// (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 // rows, so the host side picks only the grid. template diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index da6458b..4c3dbbb 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -2,8 +2,9 @@ The kernel-level tests exercise the two stateless primitives (``quantize`` for bf16/fp16/fp32 -> FP8, ``mm_fp8`` for the pre-quantized GEMM with transposed -operands); the policy-level tests (recipes, autocast context, per-tensor meta, -CPU fallbacks of the custom ops) run without a GPU. +operands); the policy-level tests (recipes, autocast context, per-tensor +meta) run without a GPU. The primitives themselves are CUDA-only +(attention-style direct wrappers — no torch.library dispatch layer). """ import threading @@ -24,7 +25,7 @@ from astrai.extension.fp8 import ( fp8_linear_enabled, 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 @@ -448,37 +449,72 @@ def test_fp8_tensor_meta_delayed_update(): meta.w.seed(w, "e4m3") assert meta.w.initialized torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1)) - # [hist | scale] packing: views alias the single state buffer. - assert meta.w.state.numel() == 4 + 2 + # [hist | scale | legacy | amax | done] packing: views alias one buffer. + assert meta.w.state.numel() == 4 + 4 assert meta.w.hist.data_ptr() == meta.w.state.data_ptr() assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr() meta.w.advance() assert meta.w.idx == 1 - # update folds a fresh amax into the window and publishes the next scale - amax = torch.tensor([8.0]) - meta.w.update(amax, "e4m3") - torch.testing.assert_close(meta.w.scale, torch.tensor([8.0 / 448.0])) + # fold_args hands the kernel the buffer, the slot and the recipe constants + args = meta.w.fold_args("e4m3") + assert args["ring_state"] is meta.w.state and args["hist_idx"] == 1 + assert args["fp8_max"] == 448.0 and args["pow2_margin"] == 1.0 -def test_quantize_cpu_fallback(): - """CPU fallback of the quantize primitive (scale semantics + amax).""" - x = torch.randn(16, 32, dtype=torch.bfloat16) - scale = torch.tensor([0.5]) # quantize multiplier - x8, amax = quantize(x, scale, "e4m3") - assert x8.dtype == torch.float8_e4m3fn - ref = (x.float() * 0.5).to(torch.float8_e4m3fn) - assert torch.equal(x8, ref) +@skip_no_fp8 +@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"]) +def test_quantize_dual_and_transposed_orientations(fmt): + """quantize_dual yields both orientations from one read; quantize's + transposed switch keeps the 2-tuple arity with the [cols][rows] layout.""" + torch.manual_seed(11) + x = torch.randn(37, 67, device="cuda", dtype=torch.bfloat16) * 3 + 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)) -def test_mm_fp8_cpu_fallback(): - a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn) - b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn) - scale = torch.tensor([1.0]) - out = mm_fp8(a8, b8, scale) - ref = (a8.float() @ b8.float() * 1.0).to(torch.bfloat16) - torch.testing.assert_close(out, ref) +@skip_no_fp8 +@pytest.mark.parametrize("fmt,fmax", [("e4m3", 448.0), ("e5m2", 57344.0)]) +@pytest.mark.parametrize("margin", [0, 1]) +def test_quantize_ring_fold_matches_host_update(fmt, fmax, margin): + """The in-kernel delayed-scaling fold matches a host-side reference.""" + dev = torch.device("cuda") + 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 # --------------------------------------------------------------------------