diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index fccc326..57c401c 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -36,10 +36,7 @@ from typing import Dict, List, Optional import torch from torch.library import Library -from astrai.extension.ops.fp8 import ( - linear_backward_fp8, - linear_forward_fp8, -) +from astrai.extension.ops.fp8 import mm_fp8, quantize # Max representable value per FP8 format (E4M3: 448, E5M2: 57344). FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0} @@ -95,11 +92,10 @@ class DynamicScaling(FP8Recipe): class _ScaleRing: """One operand's delayed-scaling state: a float32 buffer - ``[hist[n] | scale | counter]`` (views). The quantize kernel's last-finishing - block records the measured amax into ``hist[idx]``, reduces the window and - publishes the next scale entirely on device — the Python-side write/max/write - chain is gone. The counter slot stays int32-zero (float bits) between - launches; ``idx`` advances host-side each step. + ``[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. """ __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized") @@ -114,7 +110,7 @@ class _ScaleRing: self.initialized = False def advance(self) -> None: - """Rotate to the next history slot after an in-kernel finalize.""" + """Rotate to the next history slot after metadata update.""" self.idx = (self.idx + 1) % self.hist.numel() def seed(self, t: torch.Tensor, fmt: str) -> None: @@ -123,12 +119,15 @@ 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)) + class FP8TensorMeta: - """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/ - ``g``). Fused kernels record amax while quantizing, so the scale used at step - N reflects amax from steps < N. DynamicScaling never allocates a meta — it - measures the current amax inline. + """Per-weight delayed-scaling state for ``w``, ``x`` and ``g``. + + DynamicScaling never allocates a meta; it measures the current amax inline. """ __slots__ = ("w", "x", "g") @@ -213,7 +212,11 @@ class FP8State: return meta def reset(self) -> None: + """Restore construction defaults (switch, recipe, format) and drop all + per-weight metas — a full state reset for tests / reconfiguration.""" self.default_enabled = False + self.default_recipe = DelayedScaling() + self.default_format = FP8Format.HYBRID self._metas.clear() @@ -307,6 +310,11 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor return recipe.scale_from_history(amax, fmt) +def _is_fp8(dtype: torch.dtype) -> bool: + """A pre-quantized weight takes the GEMM directly (no re-quantize).""" + return dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + + _zero_bias: Dict[Optional[int], torch.Tensor] = {} @@ -326,9 +334,9 @@ def fp8_linear_forward( ): """Scaled fp8 linear forward (called from the aten::linear impl). - Pure FP8 path for both recipes: quantize x/w with the active scales, run the - pre-quantized GEMM. Delayed scaling finalizes the rings inside the quantize - kernels (amax folded into the window, next scale published on device); + Composed from the two stateless primitives: quantize x/w with the active + scales, run the pre-quantized GEMM, add the bias. 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). @@ -337,49 +345,39 @@ def fp8_linear_forward( if cfg is None: cfg = _current_config() fmt = cfg.fp8_format.fwd() - margin = cfg.recipe.margin if bias is None: bias = _empty_bias(x) - if isinstance(cfg.recipe, DynamicScaling): # measure-then-quantize, no state + if isinstance(cfg.recipe, DynamicScaling): sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt) sw = _dynamic_scale(w, cfg.recipe, fmt) - out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt) - return out, sx, sw + x8, _ = quantize(x, sx.reciprocal(), fmt) + w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0] + out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape( + *x.shape[:-1], w.size(0) + ) + return (out + bias if bias.numel() else out), sx, sw meta = state.get_weight_meta(w) if not meta.w.initialized: meta.w.seed(w, fmt) if not meta.x.initialized: meta.x.seed(x, fmt) - # The quantize kernels finalize each ring in-kernel and overwrite the ring's - # scale slot, which ALIASES meta.*.scale (a view into the state buffer). - # Snapshot the scales first so the GEMM dequantizes with the SAME scale the - # operands were quantized with, and so the backward can reuse this step's - # scale (gradient consistency with the forward). The ring finalize still - # publishes the next step's scale into the original slot. - if w.dtype is not torch.bfloat16: # static pre-quantized weight - w_arg, sw_arg, w_ring = w, meta.w.scale, None + sx, sw = meta.x.scale.clone(), meta.w.scale.clone() + x8, amax_x = quantize(x, sx.reciprocal(), fmt) + if _is_fp8(w.dtype): + w8, amax_w = w, None else: - w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state - sx = meta.x.scale.clone() - sw = sw_arg.clone() - out, _x8, _w8, _ax, _aw = linear_forward_fp8( - x, - w_arg, - bias, - sx, - sw, - fmt, - None, - meta.x.state, - meta.x.idx, - margin, - w_ring, - meta.w.idx, - margin, + w8, amax_w = quantize(w, sw.reciprocal(), fmt) + out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape( + *x.shape[:-1], w.size(0) ) + if bias.numel(): + out = out + bias + meta.x.update(amax_x, fmt) + if amax_w is not None: + meta.w.update(amax_w, fmt) meta.x.advance() - if w_ring is not None: + if amax_w is not None: meta.w.advance() return out, sx, sw @@ -410,37 +408,25 @@ class _LinearFp8(torch.autograd.Function): def backward(ctx, g): x, w, _sx_fwd, _sw_fwd = ctx.saved_tensors fmt = ctx.fmt_bwd - # Per-recipe scale/ring selection; both branches share one call below. if ctx.is_dynamic: sg = _dynamic_scale(g, ctx.recipe, fmt) sw = _dynamic_scale(w, ctx.recipe, fmt) sx = _dynamic_scale(x, ctx.recipe, fmt) - ring, idx = None, 0 else: meta = ctx.meta if not meta.g.initialized: meta.g.seed(g, fmt) - # Snapshot the g scale before its ring finalize overwrites the slot - # (same aliasing as the forward); reuse the forward's w/x scales so - # the backward quantizes with the scale the forward actually used. sg = meta.g.scale.clone() - ring, idx = meta.g.state, meta.g.idx sw, sx = _sw_fwd, _sx_fwd - grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8( - g, - x, - w, - list(ctx.needs_input_grad), - sg, - sw, - sx, - fmt, - ring, - idx, - ctx.recipe.margin, - ) + g8, amax_g = quantize(g, sg.reciprocal(), fmt) + x8, _ = quantize(x, sx.reciprocal(), fmt) + w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0] + grad_x = mm_fp8(g8, w8, sg * sw) # g8[m,n] @ w8[n,k] natural + grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True) # g8.T @ x8 + grad_b = g.sum(0).to(torch.bfloat16) if not ctx.is_dynamic: - meta.g.advance() # the g quantize kernel finalized the ring in-kernel + meta.g.update(amax_g, fmt) + meta.g.advance() return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index 549f2a1..ad34e48 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -2,10 +2,8 @@ Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives: -- ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → FP8 with fused amax +- ``quantize(x, scale, fmt) -> (x8, amax)`` — BF16/FP16/FP32 → FP8 with fused amax - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output) -- ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, x8, w8, amax_x, amax_w)`` -- ``linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt) -> (gx, gw, gb, amax_g)`` Scale semantics: scales are *quantization steps* — the value divided out when quantizing (``x8 = x / scale``). Every primitive computes its own inverse @@ -16,7 +14,7 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``; this module is stateless. """ -from typing import List, Optional, Tuple +from typing import Tuple import torch from torch.library import custom_op @@ -34,6 +32,14 @@ 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 @@ -42,28 +48,31 @@ def _fmt_dtype(fmt: str) -> torch.dtype: def fp8_quantize( x: torch.Tensor, scale: torch.Tensor, fmt: int ) -> Tuple[torch.Tensor, torch.Tensor]: - """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``.""" + """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 else torch.float8_e4m3fn + 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) + + @fp8_quantize.register_kernel("cuda") def _fp8_quantize_cuda(x, scale, fmt): - if x.dtype != torch.bfloat16: - raise TypeError(f"fp8 quantize requires bf16 input, got {x.dtype}") - return get_module("fp8_ops").quantize_bf16(x, scale, int(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("e5m2" if fmt else "e4m3")) + x8 = (x.float() * scale).to(_fmt_dtype(_fmt_name(fmt))) amax = x.abs().amax().float().reshape(1).clamp_min(1e-12) return x8, amax @@ -72,50 +81,51 @@ def _fp8_quantize_cpu(x, scale, fmt): def fp8_gemm( a: torch.Tensor, b: torch.Tensor, - sa: torch.Tensor, - sb: torch.Tensor, - out_dtype: int = 0, - out_scale: Optional[torch.Tensor] = None, + scale: torch.Tensor, + trans_a: int = 0, + trans_b: int = 0, ) -> torch.Tensor: - """FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation. + """FP8 GEMM: ``a @ b * scale`` with FP32 accumulation. - ``out_dtype``: 0 = BF16 (default), 1 = FP8 E4M3 (requires ``out_scale``, - the quantization step for the output — mirrors ``torch._scaled_mm``). + The result is always BF16; FP8 output is a separate quantize operation. """ @fp8_gemm.register_fake -def _fp8_gemm_fake(a, b, sa, sb, out_dtype=0, out_scale=None): - dtype = torch.float8_e4m3fn if out_dtype else torch.bfloat16 - return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=dtype) +def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0): + dtype = torch.bfloat16 + return torch.empty( + (a.size(1) if trans_a else a.size(0), b.size(0) if trans_b else b.size(1)), + device=a.device, + dtype=dtype, + ) @fp8_gemm.register_kernel("cuda") -def _fp8_gemm_cuda(a, b, sa, sb, out_dtype=0, out_scale=None): +def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0): 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, sa, sb, int(out_dtype), out_scale) + return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b) @fp8_gemm.register_kernel("cpu") -def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None): - acc = a.float() @ b.float() * sa * sb - if out_dtype: - os_ = 1.0 if out_scale is None else out_scale - return (acc * os_).to(torch.float8_e4m3fn) +def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0): + aa = a.float().t() if trans_a else a.float() + bb = b.float().t() if trans_b else b.float() + acc = aa @ bb * scale return acc.to(torch.bfloat16) -def quantize_bf16( +def quantize( x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3" ) -> Tuple[torch.Tensor, torch.Tensor]: - """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``. + """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; returns + ``(x8, amax)``. - ``scale`` is the quantization step (device scalar); ``fmt`` selects - E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor — the caller - never clears it. + ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects + E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. """ return fp8_quantize(x, scale, _fmt_int(fmt)) @@ -123,123 +133,14 @@ def quantize_bf16( def mm_fp8( a: torch.Tensor, b: torch.Tensor, - sa: torch.Tensor, - sb: torch.Tensor, - out_dtype: str = "bf16", - out_scale: Optional[torch.Tensor] = None, + scale: torch.Tensor, + trans_a: bool = False, + trans_b: bool = False, ) -> torch.Tensor: - """Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``. + """Pre-quantized FP8 GEMM: ``a @ b * scale``. - ``a``/``b`` must be FP8 tensors of the same format (E4M3 or E5M2); - ``sa``/``sb`` are their quantization steps. ``out_dtype`` is ``"bf16"`` - (default) or ``"e4m3"`` — FP8 output for layer-to-layer pipelines, which - requires ``out_scale`` (the output quantization step). + ``a``/``b`` must be FP8 tensors of the same format. ``scale`` is their + combined dequantization scale. The result is BF16; FP8 output is a separate + quantize operation. """ - if out_dtype not in ("bf16", "e4m3"): - raise ValueError( - f"unsupported out_dtype {out_dtype!r} (expected 'bf16' or 'e4m3')" - ) - return fp8_gemm(a, b, sa, sb, int(out_dtype == "e4m3"), out_scale) - - -def linear_forward_fp8( - x: torch.Tensor, - w: torch.Tensor, - bias: Optional[torch.Tensor], - sx: torch.Tensor, - sw: torch.Tensor, - fmt: str = "e4m3", - bias_scale: Optional[torch.Tensor] = None, - x_ring: Optional[torch.Tensor] = None, - x_ring_idx: int = 0, - x_ring_margin: int = 0, - w_ring: Optional[torch.Tensor] = None, - w_ring_idx: int = 0, - w_ring_margin: int = 0, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Pure FP8 linear forward: quantize x/w to ``fmt``, pre-quantized GEMM. - - Returns ``(out, x8, w8, amax_x, amax_w)`` — the quantized operands are - handed back so the policy layer can cache the weight quantization while - the weight tensor is unchanged (torch autocast's cached_cast analog). - ``x8`` is ``[M, K]`` and ``w8`` is ``[N, K]`` (the passed-in ``w`` itself - on the pre-quantized path). ``bias`` may be ``None``. For static fp8 - inference, ``w`` and ``bias`` may arrive pre-quantized to ``fmt`` - (produced by :func:`quantize_bf16` with their scales as ``sw`` / - ``bias_scale``); a pre-quantized ``bias`` requires ``bias_scale``, and - its ``amax_w`` comes back 0. The bias is fused into the GEMM epilogue. - ``x_ring`` / ``w_ring`` (delayed scaling) are ``[hist | scale | counter]`` - float32 buffers the quantize kernels finalize in-kernel: the measured - amax lands in ``hist[idx]`` and the next step's scale is published on - device, replacing the eager hist/max/scale update chain. - """ - fmt8 = _fmt_dtype(fmt) - if x.dtype != torch.bfloat16 or w.dtype not in (torch.bfloat16, fmt8): - raise TypeError( - f"fp8 forward requires bf16 x and bf16-or-{fmt} w, got {x.dtype}/{w.dtype}" - ) - if bias is None: - bias = torch.empty(0, device=x.device, dtype=x.dtype) - return get_module("fp8_ops").linear_forward_fp8( - x, - w, - bias, - sx, - sw, - _fmt_int(fmt), - bias_scale, - x_ring, - x_ring_idx, - x_ring_margin, - w_ring, - w_ring_idx, - w_ring_margin, - ) - - -def linear_backward_fp8( - g: torch.Tensor, - x: torch.Tensor, - w: torch.Tensor, - masks: List[bool], - sg: torch.Tensor, - sw: torch.Tensor, - sx: torch.Tensor, - fmt: str = "e5m2", - g_ring: Optional[torch.Tensor] = None, - g_ring_idx: int = 0, - g_ring_margin: int = 0, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``. - - ``g``/``x``/``w`` may each be bf16 (quantized to ``fmt`` here) or already - pre-quantized fp8 matching ``fmt`` — a pre-quantized operand skips its - quantize kernel and is read directly by the GEMM (the ``cached_cast`` - analog for the backward, symmetric with :func:`linear_forward_fp8`'s - pre-quantized weight path). ``fmt`` defaults to E5M2 (larger dynamic range - for gradients); the two GEMMs run as FP8 tensor-core products sharing a - single gradient quantization. ``g_ring`` (delayed scaling) is a - ``[hist | scale | counter]`` buffer the g quantize kernel finalizes - in-kernel (see :func:`linear_forward_fp8`); a pre-quantized ``g`` does not - finalize it and reports ``amax_g = 0``. - """ - f8 = _fmt_dtype(fmt) - for name, t in (("g", g), ("x", x), ("w", w)): - if t.dtype not in (torch.bfloat16, f8): - raise TypeError( - f"fp8 backward requires bf16 or pre-quantized {fmt} inputs, " - f"got {name}={t.dtype}" - ) - return get_module("fp8_ops").linear_backward_fp8( - g, - x, - w, - list(masks), - sg, - sw, - sx, - _fmt_int(fmt), - g_ring, - g_ring_idx, - g_ring_margin, - ) + return fp8_gemm(a, b, scale, trans_a, trans_b) diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 1cb5d83..68e09b3 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -69,32 +69,17 @@ struct Fp8GemmTraits { static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32; }; -// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional -// delayed-scaling ring finalization. Separate from FP8Params so each -// operator owns exactly the fields it touches (the GEMM never reads amax / -// ring state). Same NSDMI rationale: amax / ring_state gate optional paths -// via null checks. Still an aggregate, still trivially copyable. +// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8 +// with fused amax. struct FP8QuantizeParams { - // BF16 input and FP8 output buffers; scale_a is the quantization step - // (device scalar). amax_a (may be null) is zero-initialized by the - // binding and receives the raw-domain absolute maximum. - const void* __restrict__ a_ptr = nullptr; - void* __restrict__ out_ptr = nullptr; - const float* __restrict__ scale_a = nullptr; - float* __restrict__ amax_a = nullptr; + // Float input and FP8 output buffers; scale is the quantization + // multiplier (device scalar). amax (may be null) is zero-initialized by + // the binding and receives the raw-domain absolute maximum. + const void* __restrict__ input_ptr = nullptr; + void* __restrict__ output_ptr = nullptr; - // Optional delayed-scaling ring finalization. ring_state packs - // [hist[ring_len] | scale | counter] with ring_len = numel - 2. When - // non-null and amax_a is set, the last-finishing block records the - // measured amax into hist[ring_idx], reduces the window and publishes - // the next step's scale (max(hist) / fp8_max / 2^ring_margin) — the - // fused replacement for the eager hist-write / max / scale-write chain, - // at zero extra launches. The counter slot is a persistent zero-armed - // int32 (float bits) electing the last block each launch. - float* ring_state = nullptr; - int ring_len = 0; - int ring_idx = 0; - int ring_margin = 0; + const float* __restrict__ scale = nullptr; + float* __restrict__ amax = nullptr; // Element count (only the elementwise quantize kernel uses it). int total = 0; @@ -103,24 +88,15 @@ struct FP8QuantizeParams { // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows // through the pre-quantized GEMM kernels. Each kernel touches only the // fields it needs; buffers are raw pointers packed by the torch binding. -// Pointer members default to null (same NSDMI rationale as AttentionParams: -// bias / out_scale gate optional paths via null checks, so a partially -// packed struct must never hold garbage non-null pointers). Still an -// aggregate, still trivially copyable. +// Pointer members default to null so optional paths cannot hold garbage. struct FP8Params { // Inputs: a/b are FP8 for the pre-quantized path. Scales are // quantization steps (device scalars). const void* __restrict__ a_ptr = nullptr; const void* __restrict__ b_ptr = nullptr; - const void* __restrict__ bias = nullptr; - const float* __restrict__ scale_a = nullptr; - const float* __restrict__ scale_b = nullptr; - const float* __restrict__ bias_scale = nullptr; - // Output: BF16 or FP8 (E4M3). out_scale is the output quantization step - // (FP8 output only). void* __restrict__ out_ptr = nullptr; - const float* __restrict__ out_scale = nullptr; + const float* __restrict__ scale = nullptr; // Shapes. `int` covers every realistic LLM shape; the kernels promote // to int64 for all pointer arithmetic. int m, n, k; diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index 55bc428..9eab35c 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -2,7 +2,8 @@ // FP8 GEMM device code — pure CUDA, no torch. Mirrors the attention kernel // layout (attn_*_mma.cuh): kernels take the FP8Params POD, tile shape and // FP8 format ride on compile-time template parameters, and launchers are -// plain functions usable from both the torch binding and pure C tests. +// plain functions usable from both the torch binding and pure C tests. The +// quantize kernel lives in quantize.cuh. #include #include @@ -46,128 +47,11 @@ struct fp8_input { // FP8 MMA lives in the shared astrai::mma_sync template (common/mma.cuh); // instantiate it with fp8_input::type. Accumulates in-place: callers // pass the same accumulator array as both `d` and `c`. -// warp_reduce_max / atomic_max_float (quantize amax) live in -// common/reduce.cuh; the cp.async pipeline primitives (predicated 16-byte -// copy, commit_group, wait_group + runtime dispatch) in common/cp_async.cuh. +// warp_reduce_sum / group_reduce_sum (GEMM) live in common/reduce.cuh; the +// cp.async pipeline primitives (predicated 16-byte copy, commit_group, +// wait_group + runtime dispatch) in common/cp_async.cuh. // --------------------------------------------------------------------------- -// Quantize kernel: BF16 -> FP8 (E4M3 or E5M2), fused amax over raw values. -// --------------------------------------------------------------------------- - -// Convert one packed bf16 pair to one packed fp8 pair. amax sees the *raw* -// (unscaled) values; the stored bytes see value * inv. Bit-identical to the -// scalar __nv_fp8_*(q) constructor path (round-nearest-even + satfinite). -template -__device__ __forceinline__ unsigned quantize2(unsigned pair, float inv, - float& amax) { - const float lo = __bfloat162float(__ushort_as_bfloat16(pair & 0xffffu)); - const float hi = __bfloat162float(__ushort_as_bfloat16(pair >> 16)); - amax = fmaxf(amax, fmaxf(fabsf(lo), fabsf(hi))); - constexpr __nv_fp8_interpretation_t kFmt = - Fmt == FP8Format::E5M2 ? __NV_E5M2 : __NV_E4M3; - return static_cast(__nv_cvt_float2_to_fp8x2( - make_float2(lo * inv, hi * inv), __NV_SATFINITE, kFmt)); -} - -template -__global__ void fp8_quantize_kernel(FP8QuantizeParams p) { - const float inv = 1.0f / *p.scale_a; - const auto* x = reinterpret_cast(p.a_ptr); - void* x8 = p.out_ptr; - float* amax = p.amax_a; - float local_amax = 0.0f; - const int64_t stride = (int64_t)blockDim.x * gridDim.x; - - // Vectorized body: 8 bf16 (16B load) -> 8 fp8 (8B store) per step. Torch - // allocations are >=16B aligned and the binding passes freshly allocated - // contiguous buffers, so element 0 keeps the uint4/uint2 accesses - // natural; a misaligned base (contiguous view with an odd storage - // offset) falls back to the scalar loop below via total_vec = 0. - const bool aligned = - ((reinterpret_cast(x) | reinterpret_cast(x8)) & 15) == - 0; - const int64_t total_vec = aligned ? p.total / 8 : 0; - const uint4* xv = reinterpret_cast(x); - uint2* o8 = reinterpret_cast(x8); - for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total_vec; - i += stride) { - const uint4 v = xv[i]; - const unsigned pair[4] = {v.x, v.y, v.z, v.w}; - unsigned packed[2] = {0u, 0u}; -#pragma unroll - for (int j = 0; j < 4; ++j) - packed[j >> 1] |= quantize2(pair[j], inv, local_amax) - << (16 * (j & 1)); - o8[i] = make_uint2(packed[0], packed[1]); - } - // Scalar tail (and full fallback for misaligned bases). - for (int64_t i = total_vec * 8 + blockIdx.x * blockDim.x + threadIdx.x; - i < p.total; i += stride) { - const float f = __bfloat162float(x[i]); - local_amax = fmaxf(local_amax, fabsf(f)); - if constexpr (Fmt == FP8Format::E5M2) { - reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] = __nv_fp8_e5m2(f * inv); - } else { - reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] = __nv_fp8_e4m3(f * inv); - } - } - if (amax) { - local_amax = warp_reduce_max(local_amax); - __shared__ float slots[32]; - if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax; - __syncthreads(); - if (threadIdx.x == 0) { - float v = 0.0f; - for (int w = 0; w < (blockDim.x >> 5); ++w) v = fmaxf(v, slots[w]); - atomic_max_float(amax, v); - } - } - if (p.ring_state && amax) { - // Delayed-scaling ring finalization as a last-block epilogue (the - // CUDA threadFenceReduction pattern): the fence + counter elect the - // final block once every block's atomic_max above is visible; warp 0 - // folds the fresh amax into the window, reduces it and publishes the - // next step's scale, then re-arms the counter for the next launch. - // __fdiv_rn / ldexpf keep the scale bit-identical to the eager - // (peak / fp8_max) / 2^margin fp32 chain despite --use_fast_math. - __threadfence(); - __shared__ bool ring_last; - if (threadIdx.x == 0) - ring_last = atomicAdd(reinterpret_cast(p.ring_state + - p.ring_len + 1), - 1) == gridDim.x - 1; - __syncthreads(); - if (ring_last && threadIdx.x < 32) { - float* hist = p.ring_state; - const int lane = threadIdx.x; - float v = 0.0f; - if (lane < p.ring_len) v = hist[lane]; - if (lane == p.ring_idx) { - v = *amax; // the global amax is final now - hist[lane] = v; - } - // Windows longer than one warp (atypical) fold the tail. - for (int i = lane + 32; i < p.ring_len; i += 32) { - float h = hist[i]; - if (i == p.ring_idx) { - h = *amax; - hist[i] = h; - } - v = fmaxf(v, h); - } - const float peak = warp_reduce_max(v); - if (lane == 0) { - constexpr float kFmtMax = - Fmt == FP8Format::E5M2 ? 57344.0f : 448.0f; - p.ring_state[p.ring_len] = fmaxf( - ldexpf(__fdiv_rn(peak, kFmtMax), -p.ring_margin), 1e-12f); - __threadfence(); - // Re-arm the counter (0.0f bits == int32 0). - p.ring_state[p.ring_len + 1] = 0.0f; - } - } - } -} // Swizzled address inside a flat [rows * K] staging tile: the 16-byte chunk // index is XORed with a row-dependent slice so a warp's fragment load (8 @@ -446,8 +330,7 @@ struct Fp8GemmSmem { // (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the // launcher dispatches to it there (see launch_fp8_gemm). -template __global__ void __launch_bounds__(Traits::kCtaThreads, Fp8GemmSmem(p.a_ptr); const auto* b = reinterpret_cast(p.b_ptr); auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr); - auto* out_fp8 = reinterpret_cast<__nv_fp8_e4m3*>(p.out_ptr); const int64_t m = p.m, n = p.n, k = p.k; const int64_t a_ld = p.a_ld, b_ld = p.b_ld; @@ -537,8 +419,7 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, (int64_t)block_n * kBlockN + warp_n * 32 + thread_in_group * 2; const int a_row0 = warp_m * 64; // + mt * 16 in the loop const int b_row0 = warp_n * 32; // + nt * 8 - const float sa = *p.scale_a; - const float sb = *p.scale_b; + const float scale = *p.scale; float acc[4][4][4] = {}; // [nt][mt][acc] // Both operands end up in the canonical [M][kK] / [N][kK] shared tiles @@ -742,49 +623,25 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, } } - const float output_scale = sa * sb; - // Fused bias: BF16 raw values, or FP8 storage dequantized by its own - // scale (bias_scale != null selects the FP8 path; the format follows the - // kernel's Traits). Added in real units after the operand dequantization - // and before any output quantization. - const auto* bias16 = static_cast(p.bias); - const auto* bias8 = static_cast(p.bias); - auto bias_val = [&](int64_t col) -> float { - if (p.bias == nullptr || col >= n) return 0.0f; - if (p.bias_scale == nullptr) return __bfloat162float(bias16[col]); - return __half2float(__half(bias8[col])) * *p.bias_scale; - }; + const float output_scale = scale; #pragma unroll for (int nt = 0; nt < 4; ++nt) { const int64_t col = output_col + nt * 8; - const float b0 = bias_val(col); - const float b1 = bias_val(col + 1); // Per-row store: FP8 packs two adjacent columns into one 16-bit // write, BF16 into one 32-bit __nv_bfloat162 (single cvt+pack // instruction); boundary or unaligned columns fall back to scalar // converts so a pack never crosses the row edge or misaligns. auto store_out = [&](int64_t row, float v0, float v1) { if (row >= m) return; - const float r0 = v0 * output_scale + b0; - const float r1 = v1 * output_scale + b1; - if constexpr (OutFp8) { - if (col + 1 < n) { - *reinterpret_cast(out_fp8 + row * n + col) = - static_cast(__nv_cvt_float2_to_fp8x2( - make_float2(r0 * *p.out_scale, r1 * *p.out_scale), - __NV_SATFINITE, __NV_E4M3)); - } else { - out_fp8[row * n + col] = __nv_fp8_e4m3(r0 * *p.out_scale); - } + const float r0 = v0 * output_scale; + const float r1 = v1 * output_scale; + auto* dst = out_bf16 + row * n + col; + if (col + 1 < n && (reinterpret_cast(dst) & 3) == 0) { + *reinterpret_cast<__nv_bfloat162*>(dst) = + __floats2bfloat162_rn(r0, r1); } else { - auto* dst = out_bf16 + row * n + col; - if (col + 1 < n && (reinterpret_cast(dst) & 3) == 0) { - *reinterpret_cast<__nv_bfloat162*>(dst) = - __floats2bfloat162_rn(r0, r1); - } else { - dst[0] = __float2bfloat16(r0); - if (col + 1 < n) dst[1] = __float2bfloat16(r1); - } + dst[0] = __float2bfloat16(r0); + if (col + 1 < n) dst[1] = __float2bfloat16(r1); } }; #pragma unroll @@ -803,16 +660,6 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, // Launchers — pure CUDA (no torch), usable from the binding and pure C tests. // --------------------------------------------------------------------------- -template -void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) { - constexpr int kThreads = 256; - // One block per 256 vectors (8 elements each); at least one block so the - // scalar tail of a tiny / misaligned tensor is still covered. - int64_t blocks = (p.total / 8 + kThreads - 1) / kThreads; - if (blocks < 1) blocks = 1; - fp8_quantize_kernel<<>>(p); -} - // Launch one kernel instantiation with its shared-memory budget: stages live // in dynamic smem, so budgets beyond the 48KB static limit opt in once per // instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared @@ -851,7 +698,7 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block, // dX ~39 TF staged vs ~38 direct). constexpr int64_t kCrossStageMinK = 8192; -template || std::is_same_v> void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { @@ -860,24 +707,24 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { if (p.m <= 64) { using Traits = Fp8GemmTraits; if (b_staged) - launch_with_smem>( Fp8GemmSmem::kBytes, grid, dim3(Traits::kCtaThreads), stream, p); else - launch_with_smem>( Fp8GemmSmem::kBytes, grid, dim3(Traits::kCtaThreads), stream, p); } else { using Traits = Fp8GemmTraits; if (b_staged) - launch_with_smem>( Fp8GemmSmem::kBytes, grid, dim3(Traits::kCtaThreads), stream, p); else - launch_with_smem>( Fp8GemmSmem::kBytes, grid, dim3(Traits::kCtaThreads), stream, p); diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index dad2072..0f1c59f 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -1,76 +1,68 @@ -// FP8 GEMM torch binding: tensor validation, FP8Params packing, template -// dispatch and pybind. Device code lives in gemm.cuh (pure CUDA) — -// mirroring the attn_*.cu / attn_*_mma.cuh split of the attention kernels. +// CUDA bindings for the two stateless FP8 primitives. -#include #include #include -#include +#include + #include #include #include #include -#include "gemm.cuh" #include "../common/device.cuh" +#include "gemm.cuh" +#include "quantize.cuh" using namespace astrai::fp8; namespace { -// FP8Format / FP8Params and the launchers live in astrai::fp8 (common.h / -// gemm.cuh); this TU opens the using-directive above so the binding reads -// them unqualified. - void check_fp8_device(const torch::Tensor& tensor) { static std::mutex mutex; static std::unordered_map supported; const int device = tensor.device().index(); { std::lock_guard lock(mutex); - auto cached = supported.find(device); - if (cached != supported.end()) { - TORCH_CHECK(cached->second, - "fused FP8 MMA requires compute capability 8.9 or newer"); + auto it = supported.find(device); + if (it != supported.end()) { + TORCH_CHECK(it->second, "FP8 MMA requires compute capability 8.9+"); return; } } - const auto* properties = at::cuda::getDeviceProperties(device); - const bool is_supported = - astrai::sm_at_least(properties->major, properties->minor, - astrai::kMinSmForFp8Major, - astrai::kMinSmForFp8Minor); + const bool ok = astrai::sm_at_least( + properties->major, properties->minor, astrai::kMinSmForFp8Major, + astrai::kMinSmForFp8Minor); { std::lock_guard lock(mutex); - supported.emplace(device, is_supported); + supported.emplace(device, ok); } - TORCH_CHECK(is_supported, - "fused FP8 MMA requires compute capability 8.9 or newer"); + TORCH_CHECK(ok, "FP8 MMA requires compute capability 8.9+"); } -void check_scale(const torch::Tensor& scale, const torch::Tensor& input, - const char* name) { +void check_scale(const torch::Tensor& scale, const torch::Tensor& input) { TORCH_CHECK(scale.is_cuda() && scale.device() == input.device() && scale.scalar_type() == torch::kFloat32 && scale.numel() == 1, - name, " must be a CUDA float32 scalar on the input device"); + "scale must be a CUDA float32 scalar on the input device"); } -// ---- FP8Params packing (mirrors attention/entry_utils.cuh pack_* helpers) ---- +void pack_quantize(FP8QuantizeParams& p, const void* input, void* output, + const torch::Tensor& scale, torch::Tensor& amax, + int64_t total) { + p.input_ptr = input; + p.output_ptr = output; + p.scale = scale.data_ptr(); + p.amax = amax.data_ptr(); + p.total = static_cast(total); +} -void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out, - const torch::Tensor& sa, const torch::Tensor& sb, - const torch::Tensor* out_scale, const void* bias, - const torch::Tensor* bias_scale, int64_t m, int64_t n, - int64_t k, int64_t a_ld, int64_t b_ld) { +void pack_gemm(FP8Params& p, const void* a, const void* b, void* output, + const torch::Tensor& scale, int64_t m, int64_t n, int64_t k, + int64_t a_ld, int64_t b_ld) { p.a_ptr = a; p.b_ptr = b; - p.out_ptr = out; - p.scale_a = sa.data_ptr(); - p.scale_b = sb.data_ptr(); - p.out_scale = out_scale ? out_scale->data_ptr() : nullptr; - p.bias = bias; - p.bias_scale = bias_scale ? bias_scale->data_ptr() : nullptr; + p.out_ptr = output; + p.scale = scale.data_ptr(); p.m = static_cast(m); p.n = static_cast(n); p.k = static_cast(k); @@ -78,443 +70,112 @@ void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out, p.b_ld = static_cast(b_ld); } -// Pack the quantize params, optionally wiring the delayed-scaling ring. -// ring (may be null) packs [hist[len] | scale | counter]; len/margin come -// from the active recipe and idx is the caller's slot for this step. -void pack_quantize_params(FP8QuantizeParams& p, const void* x, void* x8, - const torch::Tensor& scale, torch::Tensor* amax, - const torch::Tensor* ring, int64_t ring_idx, - int64_t ring_margin, int64_t total) { - p.a_ptr = x; - p.out_ptr = x8; - p.scale_a = scale.data_ptr(); - p.amax_a = amax ? amax->data_ptr() : nullptr; - if (ring && ring->defined()) { - TORCH_CHECK(ring->is_cuda() && ring->scalar_type() == torch::kFloat32 && - ring->numel() >= 3 && ring->is_contiguous(), - "ring must be a contiguous CUDA float32 tensor packing " - "[hist | scale | counter]"); - p.ring_state = ring->data_ptr(); - p.ring_len = static_cast(ring->numel() - 2); - p.ring_idx = static_cast(ring_idx); - p.ring_margin = static_cast(ring_margin); - } - p.total = static_cast(total); -} - -// ---- GEMM launch dispatch (runtime flags -> compile-time kernel variants) ---- - template -void launch_gemm_variant(const FP8Params& p, cudaStream_t stream) { - static_assert(Variant >= 0 && Variant < 8, - "invalid FP8 GEMM dispatch variant"); - constexpr bool out_fp8 = (Variant & 4) != 0; - // Variant bits 1/0 = trans_a/trans_b -> CUTLASS-style layout tags - // (trans_a ? A ColMajor : RowMajor, same for B; see common.h). +void launch_variant(const FP8Params& p, cudaStream_t stream) { using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>; using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>; - launch_fp8_gemm(p, stream); + launch_fp8_gemm(p, stream); } template -void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool out_fp8, - bool trans_a, bool trans_b) { - // Encode the runtime flags as [output FP8, transpose A, transpose B]. - const int variant = (static_cast(out_fp8) << 2) | - (static_cast(trans_a) << 1) | +void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a, + bool trans_b) { + const int variant = (static_cast(trans_a) << 1) | static_cast(trans_b); switch (variant) { - case 0: launch_gemm_variant(p, stream); break; - case 1: launch_gemm_variant(p, stream); break; - case 2: launch_gemm_variant(p, stream); break; - case 3: launch_gemm_variant(p, stream); break; - case 4: launch_gemm_variant(p, stream); break; - case 5: launch_gemm_variant(p, stream); break; - case 6: launch_gemm_variant(p, stream); break; - case 7: launch_gemm_variant(p, stream); break; + case 0: launch_variant(p, stream); break; + case 1: launch_variant(p, stream); break; + case 2: launch_variant(p, stream); break; + case 3: launch_variant(p, stream); break; } } } // namespace -// --------------------------------------------------------------------------- -// Entry points -// --------------------------------------------------------------------------- - -std::tuple quantize_bf16(torch::Tensor x, - torch::Tensor scale, - int64_t fmt) { - // BF16 -> FP8 quantize with fused amax. fmt: 0 = E4M3, 1 = E5M2. - // Returns (x8, amax); the caller never clears amax (zero-initialized here). - TORCH_CHECK(x.is_cuda() && scale.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16"); - check_scale(scale, x, "scale"); +std::tuple quantize(torch::Tensor x, + torch::Tensor scale, + int64_t fmt) { + TORCH_CHECK(x.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(x.scalar_type() == torch::kBFloat16 || + x.scalar_type() == torch::kHalf || + x.scalar_type() == torch::kFloat32, + "x must be bf16, fp16 or fp32"); + TORCH_CHECK(fmt == static_cast(FP8Format::E4M3) || + fmt == static_cast(FP8Format::E5M2), + "unsupported quantization type: expected E4M3 (0) or E5M2 (1)"); + check_scale(scale, x); check_fp8_device(x); const at::cuda::OptionalCUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); - - auto x_c = x.contiguous(); - auto x8 = torch::empty_like( - x_c, x_c.options().dtype(fmt ? torch::kFloat8_e5m2 - : torch::kFloat8_e4m3fn)); - // amax feeds the cross-block atomic_max; zero it on the stream (empty + - // memset, not torch::zeros — the latter routes through a fill_ dispatcher). - auto amax = torch::empty({1}, x_c.options().dtype(torch::kFloat32)); - cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream()); + auto input = x.contiguous(); + auto output = torch::empty_like( + input, input.options().dtype(fmt ? torch::kFloat8_e5m2 + : torch::kFloat8_e4m3fn)); + auto amax = torch::zeros({1}, input.options().dtype(torch::kFloat32)); FP8QuantizeParams p; - pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax, - nullptr, 0, 0, x_c.numel()); - if (fmt) { - launch_fp8_quantize(p, stream.stream()); + pack_quantize(p, input.data_ptr(), output.data_ptr(), scale, amax, + input.numel()); + const bool e5m2 = fmt == static_cast(FP8Format::E5M2); + if (x.scalar_type() == torch::kHalf) { + if (e5m2) + launch_fp8_quantize(p, stream.stream()); + else + launch_fp8_quantize(p, stream.stream()); + } else if (x.scalar_type() == torch::kFloat32) { + if (e5m2) + launch_fp8_quantize(p, stream.stream()); + else + launch_fp8_quantize(p, stream.stream()); } else { - launch_fp8_quantize(p, stream.stream()); + if (e5m2) + launch_fp8_quantize( + p, stream.stream()); + else + launch_fp8_quantize( + p, stream.stream()); } C10_CUDA_CHECK(cudaGetLastError()); - return {x8, amax}; + return {output, amax}; } -torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa, - torch::Tensor sb, int64_t out_dtype, - c10::optional out_scale, int64_t trans_a, - int64_t trans_b) { - // Pre-quantized FP8 GEMM: out = op(a) @ op(b)^T * (sa * sb), FP32 accum. - // trans_a / trans_b select the operand layout (0 = stored [M,K]/[K,N], - // 1 = transposed [K,M]/[N,K]); the default (0/0) is the plain a @ b. - // out_dtype: 0 = BF16 (default), 1 = FP8 E4M3 (requires out_scale, the - // output quantization step). Both operands share one format. +torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale, + int64_t trans_a, int64_t trans_b) { 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, - "a and b must be fp8 (e4m3fn or e5m2)"); - TORCH_CHECK(a.scalar_type() == b.scalar_type(), - "a and b must share the same fp8 format"); + "a and b must be fp8"); + TORCH_CHECK(a.scalar_type() == b.scalar_type(), "a and b must share format"); TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D"); - TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device"); - check_scale(sa, a, "sa"); - check_scale(sb, a, "sb"); + TORCH_CHECK(a.device() == b.device(), "a and b must share device"); + check_scale(scale, a); check_fp8_device(a); const at::cuda::OptionalCUDAGuard guard(a.device()); auto stream = at::cuda::getCurrentCUDAStream(); - auto a_c = a.contiguous(); auto b_c = b.contiguous(); - const bool ta = (trans_a == 1), tb = (trans_b == 1); - // Physical leading dimension = column count of each contiguous buffer. + const bool ta = trans_a != 0; + const bool tb = trans_b != 0; const int64_t a_ld = a_c.size(1); const int64_t b_ld = b_c.size(1); - // Logical GEMM shape derived from the layout flags. const int64_t m = ta ? a_c.size(1) : a_c.size(0); const int64_t k = ta ? a_c.size(0) : a_c.size(1); const int64_t n = tb ? b_c.size(0) : b_c.size(1); - const int64_t k2 = tb ? b_c.size(1) : b_c.size(0); - TORCH_CHECK(k == k2, "inner dim mismatch"); - const bool out_fp8 = (out_dtype == 1); - TORCH_CHECK(out_dtype == 0 || out_fp8, - "out_dtype must be 0 (bf16) or 1 (fp8 e4m3)"); - torch::Tensor os; - if (out_fp8) { - TORCH_CHECK(out_scale.has_value(), "fp8 output requires out_scale"); - os = out_scale.value(); - check_scale(os, a, "out_scale"); - } - auto out = torch::empty( - {m, n}, - out_fp8 ? a_c.options().dtype(torch::kFloat8_e4m3fn) - : a_c.options().dtype(torch::kBFloat16)); + TORCH_CHECK(k == (tb ? b_c.size(1) : b_c.size(0)), "inner dim mismatch"); + auto output = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16)); FP8Params p; - pack_gemm_params(p, a_c.data_ptr(), b_c.data_ptr(), out.data_ptr(), sa, sb, - out_fp8 ? &os : nullptr, nullptr, nullptr, m, n, k, a_ld, - b_ld); + pack_gemm(p, a_c.data_ptr(), b_c.data_ptr(), output.data_ptr(), scale, m, n, + k, a_ld, b_ld); if (a.scalar_type() == torch::kFloat8_e4m3fn) - dispatch_gemm(p, stream.stream(), out_fp8, ta, tb); + dispatch_gemm(p, stream.stream(), ta, tb); else - dispatch_gemm(p, stream.stream(), out_fp8, ta, tb); + dispatch_gemm(p, stream.stream(), ta, tb); C10_CUDA_CHECK(cudaGetLastError()); - return out; -} - -std::tuple -linear_forward_fp8(torch::Tensor x, torch::Tensor w, torch::Tensor bias, - torch::Tensor sx, torch::Tensor sw, int64_t fmt, - c10::optional bias_scale, - c10::optional x_ring, int64_t x_ring_idx, - int64_t x_ring_margin, c10::optional w_ring, - int64_t w_ring_idx, int64_t w_ring_margin) { - // Pure FP8 forward: quantize x/w (fmt: 0 = E4M3, 1 = E5M2), then the - // pre-quantized GEMM; the dequantized BF16 output gets the bias added. - // Returns (out, x8, w8, amax_x, amax_w): the quantized operands are - // handed back so the policy layer can cache the weight quantization - // (torch autocast's cached_cast analog — w8 is reused while the weight - // tensor is unchanged, and the backward can share x8/w8 when the fwd/bwd - // formats match). amax_x / amax_w come from the quantize kernels - // (zero-initialized here; a pre-quantized w reports amax_w = 0 — nothing - // to feed a delayed ring). w may itself be pre-quantized fp8 storage - // matching fmt (static inference weights): the weight quantize is - // skipped, amax_w stays 0, and w8 returns the passed-in w. - // When x_ring / w_ring are given (delayed scaling), the quantize kernels - // finalize them in-kernel: the returned amax is already folded into the - // ring window and the next step's scale is published on device. - TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required"); - const auto f8opt = fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn; - const bool w_prequant = w.scalar_type() == f8opt; - TORCH_CHECK( - x.scalar_type() == torch::kBFloat16 && - (w.scalar_type() == torch::kBFloat16 || w_prequant), - "x must be bf16; w must be bf16 or pre-quantized fp8 matching fmt"); - TORCH_CHECK(x.device() == w.device(), "x and w must be on the same device"); - check_scale(sx, x, "sx"); - check_scale(sw, x, "sw"); - check_fp8_device(x); - const at::cuda::OptionalCUDAGuard guard(x.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - - auto x_c = x.reshape({-1, w.size(1)}).contiguous(); // [M, K] - auto w_c = w.contiguous(); // [N, K] - int64_t m = x_c.size(0), k = x_c.size(1), n = w_c.size(0); - TORCH_CHECK(w_c.dim() == 2 && w_c.size(1) == k, "inner dim mismatch"); - const bool has_bias = bias.defined() && bias.numel() > 0; - const bool b_prequant = has_bias && bias.scalar_type() == f8opt; - if (has_bias) { - TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() && - bias.numel() == n && - (bias.scalar_type() == torch::kBFloat16 || b_prequant), - "bias must be CUDA bf16 or pre-quantized fp8 matching fmt, " - "with shape [N]"); - TORCH_CHECK(b_prequant == bias_scale.has_value(), - "fp8 bias requires bias_scale (and bf16 bias takes none)"); - if (b_prequant) check_scale(*bias_scale, x, "bias_scale"); - } - auto x8 = torch::empty({m, k}, x_c.options().dtype(f8opt)); - // Each amax slot feeds a cross-block atomic_max, so it must start at 0. - // torch::zeros would route through a fill_ dispatcher (~50us CPU per call - // in the profile); a caching-allocator empty + cudaMemsetAsync is ~2us. - // Zero both up front: the pre-quantized-w path never quantizes w, so its - // amax_w is never atomically written and must not carry stale bytes. The - // returned values are the freshly measured (or 0) amax either way. - auto amax_x = torch::empty({1}, x.options().dtype(torch::kFloat32)); - auto amax_w = torch::empty({1}, x.options().dtype(torch::kFloat32)); - cudaMemsetAsync(amax_x.data_ptr(), 0, sizeof(float), stream.stream()); - cudaMemsetAsync(amax_w.data_ptr(), 0, sizeof(float), stream.stream()); - auto out = torch::empty({m, n}, x_c.options()); - - auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst, - const torch::Tensor& scale, torch::Tensor* amax, - const c10::optional& ring, - int64_t ring_idx, int64_t ring_margin) { - FP8QuantizeParams qp; - pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax, - ring ? &*ring : nullptr, ring_idx, ring_margin, - src.numel()); - if (fmt) { - launch_fp8_quantize(qp, stream.stream()); - } else { - launch_fp8_quantize(qp, stream.stream()); - } - }; - quantize(x_c, x8, sx, &amax_x, x_ring, x_ring_idx, x_ring_margin); - // Static inference weights arrive pre-quantized (w8 storage + its scale); - // only freshly-loaded bf16 weights quantize here. - torch::Tensor w8 = w_prequant - ? w_c - : torch::empty({n, k}, x_c.options().dtype(f8opt)); - if (!w_prequant) - quantize(w_c, w8, sw, &amax_w, w_ring, w_ring_idx, w_ring_margin); - - FP8Params p; - // Forward is the NT layout: A = x8 [M,K] (a_ld = k), B = w8 [N,K] - // (b_ld = k), out = x @ w^T. The bias is fused into the epilogue (bf16 - // raw, or fp8 + bias_scale on the static path). - auto bias_c = has_bias ? bias.contiguous() : bias; - pack_gemm_params(p, x8.data_ptr(), w8.data_ptr(), out.data_ptr(), sx, sw, - nullptr, has_bias ? bias_c.data_ptr() : nullptr, - b_prequant ? &*bias_scale : nullptr, m, n, k, k, k); - if (fmt) { - launch_fp8_gemm( - p, stream.stream()); - } else { - launch_fp8_gemm( - p, stream.stream()); - } - C10_CUDA_CHECK(cudaGetLastError()); - - std::vector shape(x.sizes().begin(), x.sizes().end() - 1); - shape.push_back(n); - return {out.reshape(shape), x8, w8, amax_x, amax_w}; -} - -std::tuple -linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w, - std::vector masks, torch::Tensor sg, - torch::Tensor sw, torch::Tensor sx, int64_t fmt, - c10::optional g_ring, int64_t g_ring_idx, - int64_t g_ring_margin) { - // Pre-quantized FP8 backward: grad is quantized once (E4M3 or E5M2 per - // `fmt`), then dX / dW run as FP8 tensor-core GEMMs sharing g8. - // Returns (grad_input, grad_weight, grad_bias, amax_g). With g_ring - // (delayed scaling), the g quantize kernel finalizes the ring in-kernel - // (amax folded into the window, next step's scale published on device); - // the w/x quantizes for dX / dW never touch rings — each operand's ring - // is finalized exactly once per step (by the forward or this kernel). - TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required"); - // Each operand may be bf16 (quantized here) or already fp8 matching fmt - // (reused from the forward — the fp8 cached_cast analog, symmetric with - // the forward's pre-quantized w path). Pre-quantized operands skip their - // quantize kernel; their scale is still passed for the GEMM dequant. - const auto f8opt = fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn; - const bool g_prequant = g.scalar_type() == f8opt; - const bool x_prequant = x.scalar_type() == f8opt; - const bool w_prequant = w.scalar_type() == f8opt; - TORCH_CHECK( - (g.scalar_type() == torch::kBFloat16 || g_prequant) && - (x.scalar_type() == torch::kBFloat16 || x_prequant) && - (w.scalar_type() == torch::kBFloat16 || w_prequant), - "g, x, and w must be bf16 or pre-quantized fp8 matching fmt"); - TORCH_CHECK(g.device() == x.device() && g.device() == w.device(), - "g, x, and w must be on the same device"); - TORCH_CHECK(masks.size() == 3, "masks must contain three values"); - check_fp8_device(g); - const at::cuda::OptionalCUDAGuard guard(g.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - - auto g_c = g.reshape({-1, w.size(0)}).contiguous(); // [M, N] - auto x_c = x.reshape({-1, x.size(-1)}).contiguous(); // [M, K] - auto w_c = w.contiguous(); // [N, K] - int64_t m = g_c.size(0), n = w_c.size(0), k = w_c.size(1); - TORCH_CHECK(x_c.size(0) == m && x_c.size(1) == k && g_c.size(1) == n, - "backward shape mismatch"); - - auto grad_input = - torch::empty_like(x, x.options().dtype(torch::kBFloat16)); - auto grad_weight = - torch::empty_like(w, w.options().dtype(torch::kBFloat16)); - auto grad_bias = torch::empty({0}, g.options()); - // amax_g feeds a cross-block atomic_max in the g quantize kernel; zero it - // on the stream (empty + memset, not torch::zeros — see the forward). - // Only needed when a g quantize runs (mask[0]||mask[1]); the bias-only - // fallback below overwrites it via .copy_, so a wasted memset elsewhere - // is harmless. - auto amax_g = torch::empty({1}, g.options().dtype(torch::kFloat32)); - cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), stream.stream()); - - auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst, - const torch::Tensor& scale, torch::Tensor* amax, - const c10::optional& ring, - int64_t ring_idx, int64_t ring_margin) { - FP8QuantizeParams qp; - pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax, - ring ? &*ring : nullptr, ring_idx, ring_margin, - src.numel()); - if (fmt) { - launch_fp8_quantize(qp, stream.stream()); - } else { - launch_fp8_quantize(qp, stream.stream()); - } - }; - // Four-layout backward: the gradient and activation tensors keep their - // natural row-major layout, and the kernel reads them transposed where the - // GEMM needs it (the ColMajor layout tags pick the crosswise stage-load). - // No torch-level `.transpose().contiguous()` - // copies are required — dX uses g8 [M,N] as A with w8 [N,K] read transposed - // as B; dW uses g8 transposed as A with x8 transposed as B. - // g is quantized once (amax_g measured here); both GEMMs share g8. - auto run_bwd_gemm = [&](const FP8Params& gp, bool trans_a, bool trans_b) { - if (fmt) - dispatch_gemm(gp, stream.stream(), false, trans_a, - trans_b); - else - dispatch_gemm(gp, stream.stream(), false, trans_a, - trans_b); - }; - - torch::Tensor g8; - if (masks[0] || masks[1]) { - if (g_prequant) { - g8 = g_c; - } else { - g8 = torch::empty({m, n}, g.options().dtype(f8opt)); - quantize(g_c, g8, sg, &amax_g, g_ring, g_ring_idx, g_ring_margin); - } - } - // dX = g @ w: A = g8 [M,N] (contract over N), B = w8 [N,K] read transposed - // (b[p*b_ld + n] = w[p,n]); out = [M,K], a_ld = N, b_ld = K, contract = N. - if (masks[0]) { - torch::Tensor w8; - if (w_prequant) { - w8 = w_c; - } else { - w8 = torch::empty({n, k}, g.options().dtype(f8opt)); - quantize(w_c, w8, sw, nullptr, c10::nullopt, 0, 0); - } - auto grad_input_2d = grad_input.reshape({m, k}); - FP8Params gp; - pack_gemm_params(gp, g8.data_ptr(), w8.data_ptr(), - grad_input_2d.data_ptr(), sg, sw, nullptr, nullptr, - nullptr, m, k, n, n, k); - run_bwd_gemm(gp, false, false); - } - // dW = g^T @ x: A = g8 [M,N] read transposed (a[p*a_ld + m] = g[p,m]), B = - // x8 [M,K] read transposed (b[p*b_ld + n] = x[p,n]); out = [N,K], a_ld = N, - // b_ld = K, contract = M. - if (masks[1]) { - torch::Tensor x8; - if (x_prequant) { - x8 = x_c; - } else { - x8 = torch::empty({m, k}, g.options().dtype(f8opt)); - quantize(x_c, x8, sx, nullptr, c10::nullopt, 0, 0); - } - FP8Params gp; - pack_gemm_params(gp, g8.data_ptr(), x8.data_ptr(), - grad_weight.data_ptr(), sg, sx, nullptr, nullptr, - nullptr, n, k, m, n, k); - run_bwd_gemm(gp, true, false); - } - if (!masks[0] && !masks[1] && !g_prequant) { - amax_g.copy_(g_c.abs().amax().to(torch::kFloat32)); - } - C10_CUDA_CHECK(cudaGetLastError()); - if (masks[2]) { - // A pre-quantized g has no bf16 source to reduce; dequantize with its - // scale before the batch-sum so grad_bias stays in the true gradient - // domain (sg * sum(g8)). - grad_bias = g_prequant - ? (g_c.to(torch::kFloat32) * sg).sum(0).to(torch::kBFloat16) - : g_c.sum(0).to(g.scalar_type()); - } - return {grad_input, grad_weight, grad_bias, amax_g}; + return output; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("quantize_bf16", &quantize_bf16, py::arg("x"), py::arg("scale"), - py::arg("fmt"), - "BF16 to FP8 (E4M3/E5M2) quantize with fused amax; returns (x8, amax)"); - m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("sa"), - py::arg("sb"), py::arg("out_dtype") = 0, - py::arg("out_scale") = py::none(), py::arg("trans_a") = 0, - py::arg("trans_b") = 0, - "Pre-quantized FP8 GEMM: op(a) @ op(b)^T * (sa * sb); out_dtype " - "0=bf16, 1=fp8 e4m3 (requires out_scale); trans_a/trans_b select " - "the operand layout (default 0/0 = a@b)"); - m.def("linear_forward_fp8", &linear_forward_fp8, py::arg("x"), - py::arg("w"), py::arg("bias"), py::arg("sx"), py::arg("sw"), - py::arg("fmt") = 0, py::arg("bias_scale") = py::none(), - py::arg("x_ring") = py::none(), py::arg("x_ring_idx") = 0, - py::arg("x_ring_margin") = 0, py::arg("w_ring") = py::none(), - py::arg("w_ring_idx") = 0, py::arg("w_ring_margin") = 0, - "Pure FP8 linear forward: quantize x/w, pre-quantized GEMM with the " - "bias fused into the epilogue; w and bias may be pre-quantized fp8 " - "matching fmt (static inference path; fp8 bias requires bias_scale);" - " x_ring/w_ring optionally finalize a delayed-scaling ring " - "([hist | scale | counter] float32 buffer) in-kernel; returns " - "(out, x8, w8, amax_x, amax_w) — x8 is [M,K], w8 is [N,K] (the " - "passed-in w on the pre-quantized path)"); - m.def("linear_backward_fp8", &linear_backward_fp8, py::arg("g"), - py::arg("x"), py::arg("w"), py::arg("masks"), py::arg("sg"), - py::arg("sw"), py::arg("sx"), py::arg("fmt"), - py::arg("g_ring") = py::none(), py::arg("g_ring_idx") = 0, - py::arg("g_ring_margin") = 0, - "FP8 linear backward; g_ring optionally finalizes the gradient's " - "delayed-scaling ring in-kernel; returns (grad_input, grad_weight, " - "grad_bias, amax_g)"); + m.def("quantize", &quantize, py::arg("x"), py::arg("scale"), + py::arg("fmt")); + m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"), + py::arg("trans_a") = 0, py::arg("trans_b") = 0); } diff --git a/csrc/kernels/fp8/quantize.cuh b/csrc/kernels/fp8/quantize.cuh new file mode 100644 index 0000000..0a56e5c --- /dev/null +++ b/csrc/kernels/fp8/quantize.cuh @@ -0,0 +1,173 @@ +#pragma once +// FP8 quantize device code — pure CUDA, no torch. Any float input element +// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over +// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the +// FP8QuantizeParams POD, formats and input types ride on template parameters, +// and the launcher is a plain function usable from both the torch binding and +// pure C tests. + +#include +#include +#include +#include +#include + +#include "common.h" +#include "../common/reduce.cuh" + +namespace astrai { +namespace fp8 { + +// Input element type traits: one element -> float, and the vectorized +// unpack of one 16-byte load into kVecElems floats. +template +struct quant_in_traits; + +template <> +struct quant_in_traits<__nv_bfloat16> { + static constexpr int kVecElems = 8; + static __device__ __forceinline__ float to_float(__nv_bfloat16 v) { + return __bfloat162float(v); + } + static __device__ __forceinline__ void load_vec(const uint4& raw, + float* f) { + const unsigned w[4] = {raw.x, raw.y, raw.z, raw.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + f[2 * j] = + __bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu)); + f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16)); + } + } +}; + +template <> +struct quant_in_traits<__half> { + static constexpr int kVecElems = 8; + static __device__ __forceinline__ float to_float(__half v) { + return __half2float(v); + } + static __device__ __forceinline__ void load_vec(const uint4& raw, + float* f) { + const __half2* h2 = reinterpret_cast(&raw); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 p = __half22float2(h2[j]); + f[2 * j] = p.x; + f[2 * j + 1] = p.y; + } + } +}; + +template <> +struct quant_in_traits { + static constexpr int kVecElems = 4; + static __device__ __forceinline__ float to_float(float v) { return v; } + static __device__ __forceinline__ void load_vec(const uint4& raw, + float* f) { + f[0] = __uint_as_float(raw.x); + f[1] = __uint_as_float(raw.y); + f[2] = __uint_as_float(raw.z); + f[3] = __uint_as_float(raw.w); + } +}; + +// Convert one float pair to one packed fp8 pair. The stored bytes see +// value * mult (round-nearest-even + satfinite). +template +__device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) { + constexpr __nv_fp8_interpretation_t kFmt = + Fmt == FP8Format::E5M2 ? __NV_E5M2 : __NV_E4M3; + return static_cast(__nv_cvt_float2_to_fp8x2( + make_float2(a, b), __NV_SATFINITE, kFmt)); +} + +// Quantize kernel: float input -> FP8 (E4M3 or E5M2), fused amax over raw +// values. +template +__global__ void fp8_quantize_kernel(FP8QuantizeParams p) { + const float mult = *p.scale; + const auto* x = static_cast(p.input_ptr); + void* x8 = p.output_ptr; + float* amax = p.amax; + float local_amax = 0.0f; + const int64_t stride = (int64_t)blockDim.x * gridDim.x; + + // Vectorized body: one 16B load -> kVecElems fp8 bytes per step (8 + // elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B + // aligned and the binding passes freshly allocated contiguous buffers, + // so element 0 keeps the uint4 access natural; a misaligned base + // (contiguous view with an odd storage offset) falls back to the scalar + // loop below via total_vec = 0. + constexpr int kVecElems = quant_in_traits::kVecElems; + const bool aligned = + ((reinterpret_cast(x) | + reinterpret_cast(x8)) & + 15) == 0; + const int64_t total_vec = aligned ? p.total / kVecElems : 0; + const uint4* xv = reinterpret_cast(x); + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total_vec; + i += stride) { + float f[kVecElems]; + quant_in_traits::load_vec(xv[i], f); + // One 32-bit word packs two fp8x2 pairs (4 elements). + unsigned packed[kVecElems / 4]; +#pragma unroll + for (int j = 0; j < kVecElems / 4; ++j) { + local_amax = fmaxf( + local_amax, + fmaxf(fmaxf(fabsf(f[4 * j]), fabsf(f[4 * j + 1])), + fmaxf(fabsf(f[4 * j + 2]), fabsf(f[4 * j + 3])))); + const unsigned lo = + cvt_fp8x2(f[4 * j] * mult, f[4 * j + 1] * mult); + const unsigned hi = + cvt_fp8x2(f[4 * j + 2] * mult, f[4 * j + 3] * mult); + packed[j] = (lo & 0xffffu) | (hi << 16); + } + if constexpr (kVecElems == 8) + reinterpret_cast(x8)[i] = + make_uint2(packed[0], packed[1]); + else + reinterpret_cast(x8)[i] = packed[0]; + } + // Scalar tail (and full fallback for misaligned bases). + for (int64_t i = total_vec * kVecElems + blockIdx.x * blockDim.x + + threadIdx.x; + i < p.total; i += stride) { + const float v = quant_in_traits::to_float(x[i]); + local_amax = fmaxf(local_amax, fabsf(v)); + if constexpr (Fmt == FP8Format::E5M2) { + reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] = + __nv_fp8_e5m2(v * mult); + } else { + reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] = + __nv_fp8_e4m3(v * mult); + } + } + if (amax) { + local_amax = warp_reduce_max(local_amax); + __shared__ float slots[32]; + if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax; + __syncthreads(); + if (threadIdx.x == 0) { + float v = 0.0f; + for (int w = 0; w < (blockDim.x >> 5); ++w) + v = fmaxf(v, slots[w]); + atomic_max_float(amax, v); + } + } +} + +template +void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) { + constexpr int kThreads = 256; + // One block per 256 vectors; at least one block so the scalar tail of a + // tiny / misaligned tensor is still covered. + constexpr int kVecElems = quant_in_traits::kVecElems; + int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads; + if (blocks < 1) blocks = 1; + fp8_quantize_kernel<<>>(p); +} + +} // namespace fp8 +} // namespace astrai diff --git a/csrc/tests/fp8_test.cu b/csrc/tests/fp8_test.cu index bc5c608..3ede122 100644 --- a/csrc/tests/fp8_test.cu +++ b/csrc/tests/fp8_test.cu @@ -175,15 +175,13 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n, int k, int a_ld, int b_ld) { __nv_fp8_e4m3 *da, *db; __nv_bfloat16* dout; - float *dsa, *dsb; + float* dscale; cudaMalloc(&da, (size_t)m * k); cudaMalloc(&db, (size_t)n * k); cudaMalloc(&dout, (size_t)m * n * 2); - cudaMalloc(&dsa, 4); - cudaMalloc(&dsb, 4); + cudaMalloc(&dscale, 4); float one = 1.0f; - cudaMemcpy(dsa, &one, 4, cudaMemcpyHostToDevice); - cudaMemcpy(dsb, &one, 4, cudaMemcpyHostToDevice); + cudaMemcpy(dscale, &one, 4, cudaMemcpyHostToDevice); // quantize inputs to e4m3 on host and upload byte-by-byte std::vector qa(m * k), qb(n * k); for (int i = 0; i < m * k; ++i) { @@ -201,14 +199,13 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n, p.a_ptr = da; p.b_ptr = db; p.out_ptr = dout; - p.scale_a = dsa; - p.scale_b = dsb; + p.scale = dscale; p.m = m; p.n = n; p.k = k; p.a_ld = a_ld; p.b_ld = b_ld; - launch_fp8_gemm(p, 0); + launch_fp8_gemm(p, 0); cudaError_t e = cudaDeviceSynchronize(); if (e != cudaSuccess) { printf(" CUDA err: %s\n", cudaGetErrorString(e)); @@ -247,8 +244,7 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n, cudaFree(da); cudaFree(db); cudaFree(dout); - cudaFree(dsa); - cudaFree(dsb); + cudaFree(dscale); return ok; } diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index f959332..8c4cf3f 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -47,16 +47,17 @@ style as attention, but split into **three** files: | File | Role | |------|------| | `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits`, `FP8Params` POD — no torch | -| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_quantize_kernel` (BF16→FP8 + amax), `fp8_gemm_kernel` (pre-quantized GEMM, 128×64 CTA / 64×16 warp / 3-stage cp.async) — no torch | +| `fp8/quantize.cuh` | pure-CUDA device code: `fp8_quantize_kernel` (bf16/fp16/fp32 → FP8 + amax, `quant_in_traits` vectorized unpack) — no torch | +| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_gemm_kernel` (pre-quantized GEMM, 128×128 CTA / 64×32 warp / multi-stage cp.async, transposed-operand layouts) — no torch | | `fp8/ops.cu` | binding only: `check_fp8_device` (sm_89+), param packing, launch dispatch, pybind → module `fp8_ops` | -Scale semantics follow `torch._scaled_mm` (quantization step size: divide by -`scale`; the kernel computes the reciprocal internally — the interface never -takes `*_inv`). `amax` is always returned in the original bf16 domain. +Scale semantics: `quantize` takes the quantization *multiplier*, `mm_fp8` +takes the combined dequant scale (`sa * sb`); the strategy layer passes +`scale.reciprocal()` / `sa * sb` respectively. `amax` is always returned in +the original input domain. Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless -primitives (`quantize_bf16` / `mm_fp8` / `linear_forward_fp8` / -`linear_backward_fp8`) via `torch.library.custom_op`, and +primitives (`quantize` / `mm_fp8`) via `torch.library.custom_op`, and `astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed / dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear` on CUDA). See the FP8 section in `AGENTS.md` for full detail. diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index 61175a0..4d1ff52 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -1,9 +1,9 @@ """FP8 primitives: kernel-level (CUDA) and policy-level (CPU-verifiable) tests. -The kernel-level tests exercise the pure FP8 path (quantize_bf16 + mm_fp8 for -the forward GEMM, quantize + pre-quantized GEMMs for the backward); the -policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks -of the custom ops) run without a GPU. +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. """ import threading @@ -24,12 +24,7 @@ from astrai.extension.fp8 import ( fp8_linear_enabled, fp8_state, ) -from astrai.extension.ops.fp8 import ( - linear_backward_fp8, - linear_forward_fp8, - mm_fp8, - quantize_bf16, -) +from astrai.extension.ops.fp8 import mm_fp8, quantize from tests.conftest import skip_no_fp8 @@ -37,8 +32,11 @@ def _scale(tensor): return (tensor.abs().amax().float() / 448.0).clamp_min(1e-12) -def _quantize(tensor, scale): - return (tensor.float() / scale).to(torch.float8_e4m3fn).float() +def _quantize(tensor, scale, fmt="e4m3"): + """Reference quantize: multiply by the reciprocal (the kernel's exact + arithmetic — a plain divide flips fp8 boundary cases by one ulp).""" + dtype = torch.float8_e5m2 if fmt == "e5m2" else torch.float8_e4m3fn + return (tensor.float() * scale.reciprocal()).to(dtype).float() # -------------------------------------------------------------------------- @@ -57,9 +55,9 @@ def test_fp8_mm_matches_explicit_quantization(m, n, k): b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16) scale_a = _scale(a) scale_b = _scale(b) - a8, _ = quantize_bf16(a, scale_a, "e4m3") - b8, _ = quantize_bf16(b, scale_b, "e4m3") - out = mm_fp8(a8, b8, scale_a, scale_b) + a8, _ = quantize(a, scale_a.reciprocal(), "e4m3") + b8, _ = quantize(b, scale_b.reciprocal(), "e4m3") + out = mm_fp8(a8, b8, scale_a * scale_b) expected = (_quantize(a, scale_a) @ _quantize(b, scale_b) * scale_a * scale_b).to( torch.bfloat16 ) @@ -70,79 +68,59 @@ def test_fp8_mm_matches_explicit_quantization(m, n, k): @skip_no_fp8 -def test_quantize_bf16_returns_amax(): - """quantize_bf16 returns (x8, amax); amax tracks the *raw* values and the - caller never clears it (zero-initialized inside the kernel entry).""" +@pytest.mark.parametrize("in_dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("fmt", ["e4m3", "e5m2"]) +def test_quantize_input_dtypes(in_dtype, fmt): + """quantize accepts bf16/fp16/fp32 inputs; bytes and amax match the + explicit (value * multiplier) reference.""" torch.manual_seed(3) - x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16) + x = torch.randn(64, 128, device="cuda", dtype=torch.float32) * 0.5 + x = x.to(in_dtype) scale = torch.tensor([0.5], device="cuda") - x8, amax = quantize_bf16(x, scale, "e4m3") - assert x8.dtype == torch.float8_e4m3fn + x8, amax = quantize(x, scale, fmt) + out_dtype = torch.float8_e5m2 if fmt == "e5m2" else torch.float8_e4m3fn + assert x8.dtype == out_dtype assert x8.shape == x.shape assert amax.shape == (1,) torch.testing.assert_close(amax, x.abs().amax().float().reshape(1)) - ref = (x.float() / 0.5).to(torch.float8_e4m3fn) + ref = (x.float() * 0.5).to(out_dtype) assert torch.equal(x8, ref) @skip_no_fp8 -def test_quantize_bf16_e5m2_format(): +def test_quantize_e5m2_format(): x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) - x8, amax = quantize_bf16(x, torch.tensor([0.1], device="cuda"), "e5m2") + x8, amax = quantize(x, torch.tensor([10.0], device="cuda"), "e5m2") assert x8.dtype == torch.float8_e5m2 torch.testing.assert_close(amax, x.abs().amax().float().reshape(1)) @skip_no_fp8 -def test_quantize_ring_in_kernel_finalize(): - """The quantize kernel finalizes the delayed-scaling ring in-kernel: the - measured amax lands in hist[idx], the window reduces to the next step's - scale on device, and the counter re-arms for the next launch.""" - torch.manual_seed(21) - dev = torch.device("cuda") - ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0)) - x0 = torch.randn(256, 256, device=dev, dtype=torch.bfloat16) - w = torch.randn(256, 256, device=dev, dtype=torch.bfloat16) - sw = torch.tensor([1.0], device=dev) - ring.seed(x0, "e4m3") - hist0 = ring.hist.clone() +@pytest.mark.parametrize("trans_a", [False, True]) +@pytest.mark.parametrize("trans_b", [False, True]) +def test_mm_fp8_transposed_operands(trans_a, trans_b): + """mm_fp8 handles all four operand layouts via trans_a/trans_b.""" + torch.manual_seed(17) + m, n, k = 19, 13, 37 + a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) # A [M][K] + b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) # B^T [N][K] + sa, sb = _scale(a), _scale(b) + a8, _ = quantize(a, sa.reciprocal(), "e4m3") + b8, _ = quantize(b, sb.reciprocal(), "e4m3") + a_op = a8.t().contiguous() if trans_a else a8 + b_op = b8 if trans_b else b8.t().contiguous() - # Step over three fresh tensors: each launch folds its amax into - # hist[idx] and publishes max(hist)/448 as the next scale. - idx = 0 - for _ in range(3): - x = torch.randn(256, 256, device=dev, dtype=torch.bfloat16) * (2.0 + 4.0 * _) - _ = linear_forward_fp8( - x, - w, - None, - ring.scale, - sw, - "e4m3", - None, - ring.state, - idx, - 0, - ) - torch.cuda.synchronize() - expected_hist = hist0.clone() - expected_hist[idx] = x.abs().amax().float() - torch.testing.assert_close(ring.hist, expected_hist) - expected_scale = (expected_hist.max() / 448.0).reshape(1) - torch.testing.assert_close(ring.scale, expected_scale, rtol=1e-6, atol=1e-12) - # counter re-armed to int32 zero - assert ring.state[-1].view(torch.int32).item() == 0 - hist0 = expected_hist.clone() - idx = (idx + 1) % 4 + out = mm_fp8(a_op, b_op, sa * sb, trans_a=trans_a, trans_b=trans_b) + assert out.shape == (m, n) + expected = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16) + torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) @skip_no_fp8 def test_delayed_scaling_forward_uses_snapshot_scale(): - """Regression: the in-kernel ring finalize overwrites the scale slot, which - aliases the scale the GEMM must dequantize with. The forward must snapshot - the delayed scale first, so a changing amax across steps does not leak the - next-step scale into the output (otherwise out is off by - scale_next / scale_current).""" + """The delayed scale for step N is computed from amax(steps < N); the + forward must snapshot the scale before the ring update, so a changing + amax across steps does not leak the next-step scale into the output.""" torch.manual_seed(11) dev = torch.device("cuda") state = f8mod.fp8_state() @@ -153,8 +131,7 @@ def test_delayed_scaling_forward_uses_snapshot_scale(): m, n, k = 32, 16, 64 x1 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.5 # Smaller amax than x1: the delayed scale (amax(x1)/448) still covers - # x2 without fp8 saturation, while the next-step scale would differ — - # exactly the condition that exposed the overwrite bug. + # x2 without fp8 saturation, while the next-step scale would differ. x2 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.35 w = torch.randn(n, k, device=dev, dtype=torch.bfloat16) * 0.5 bias = torch.zeros(n, device=dev, dtype=torch.bfloat16) @@ -177,93 +154,57 @@ def test_delayed_scaling_forward_uses_snapshot_scale(): @skip_no_fp8 def test_fp8_linear_forward_and_backward(): + """The composed strategy path: forward quantize+GEMM+bias, backward + dX/dW GEMMs on transposed operands (E5M2 in hybrid).""" torch.manual_seed(7) m, n, k = 19, 13, 37 x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) - grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) - scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad) - out, x8, w8, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w) - grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( - grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3" - ) + state = f8mod.fp8_state() + state.reset() + state.default_recipe = DynamicScaling() + try: + out, _, _ = f8mod.fp8_linear_forward(x, weight, bias) - qx = _quantize(x, scale_x) - qw = _quantize(weight, scale_w) - qg = _quantize(grad, scale_g) - expected_out = (qx @ qw.t() * scale_x * scale_w + bias).to(torch.bfloat16) - expected_grad_x = (qg @ qw * scale_g * scale_w).to(torch.bfloat16) - expected_grad_w = (qg.t() @ qx * scale_g * scale_x).to(torch.bfloat16) + sx, sw = _scale(x), _scale(weight) + qx = _quantize(x, sx) + qw = _quantize(weight, sw) + expected_out = (qx @ qw.t() * sx * sw + bias).to(torch.bfloat16) + torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01) - torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01) - torch.testing.assert_close(grad_x, expected_grad_x, atol=0.125, rtol=0.01) - torch.testing.assert_close(grad_w, expected_grad_w, atol=0.125, rtol=0.01) - torch.testing.assert_close(grad_b, grad.sum(0).to(torch.bfloat16)) - torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1)) - torch.testing.assert_close(amax_w, weight.abs().amax().float().reshape(1)) - torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1)) + # backward through the aten::linear integration (hybrid E5M2). The + # incoming gradient is 2*out of the *fp8* forward (bf16-rounded), not + # 2*exact — derive the reference from the actual output. + xr = x.detach().clone().requires_grad_() + wr = weight.detach().clone().requires_grad_() + br = bias.detach().clone().requires_grad_() + with fp8_autocast(enabled=True): + loss = F.linear(xr, wr, br).float().pow(2).sum() + loss.backward() - -@skip_no_fp8 -def test_linear_backward_e5m2_gradients(): - """Hybrid backward: gradient GEMMs run in E5M2 (larger dynamic range).""" - torch.manual_seed(5) - m, n, k = 32, 16, 64 - x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 3.0 - weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) - grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) * 10.0 - sg = _scale(grad) * 0.5 - sw = _scale(weight) - sx = _scale(x) - - grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( - grad, x, weight, [1, 1, 1], sg, sw, sx, "e5m2" - ) - - def q5(t, s): - return (t.float() / s).to(torch.float8_e5m2).float() - - qg = q5(grad, sg) - qw = q5(weight, sw) - qx = q5(x, sx) - expected_grad_x = (qg @ qw * sg * sw).to(torch.bfloat16) - expected_grad_w = (qg.t() @ qx * sg * sx).to(torch.bfloat16) - torch.testing.assert_close(grad_x, expected_grad_x, atol=0.5, rtol=0.05) - torch.testing.assert_close(grad_w, expected_grad_w, atol=0.5, rtol=0.05) - torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1)) - - -@skip_no_fp8 -def test_fp8_linear_static_fp8_weight_and_bias(): - """Static fp8 inference: pre-quantized w8/b8 + their scales take the GEMM - directly (no weight quantize, amax_w = 0); the bias is fused in the - epilogue (bf16 and fp8 bias share the fused path).""" - torch.manual_seed(9) - m, n, k = 67, 45, 129 - x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.5 - bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) * 0.5 - sx, sw, sb = _scale(x), _scale(weight), _scale(bias) - - w8, _ = quantize_bf16(weight, sw, "e4m3") - b8, _ = quantize_bf16(bias, sb, "e4m3") - out, x8, w8_back, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb) - - qx = _quantize(x, sx) - qw = _quantize(weight, sw) - qb = _quantize(bias, sb) - expected = (qx @ qw.t() * sx * sw + qb * sb).to(torch.bfloat16) - torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) - torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1)) - assert amax_w.item() == 0.0 # nothing measured on the static path - assert w8_back is w8 # pre-quantized w handed straight back - - # bf16 bias stays bf16 on the same fused-epilogue path - out_bf16bias, *_ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3") - expected_b = (qx @ qw.t() * sx * sw + bias.float()).to(torch.bfloat16) - torch.testing.assert_close(out_bf16bias, expected_b, atol=0.125, rtol=0.01) + g = (2 * out.float()).to(torch.bfloat16).float() # actual grad wrt out + # the dynamic path measures current-step amax in the bwd fmt (E5M2); + # amax must be taken in fp32 — a bf16-rounded scale flips E5M2 + # boundary rounding (2-bit mantissa) and the reference drifts. + e5 = 57344.0 + sg = (g.abs().amax() / e5).clamp_min(1e-12) + sw5 = (weight.abs().amax().float() / e5).clamp_min(1e-12) + sx5 = (x.abs().amax().float() / e5).clamp_min(1e-12) + expected_grad_x = ( + _quantize(g, sg, "e5m2") @ _quantize(weight, sw5, "e5m2") * sg * sw5 + ).to(torch.bfloat16) + expected_grad_w = ( + _quantize(g, sg, "e5m2").t() @ _quantize(x, sx5, "e5m2") * sg * sx5 + ).to(torch.bfloat16) + torch.testing.assert_close(xr.grad, expected_grad_x, atol=0.5, rtol=0.05) + torch.testing.assert_close(wr.grad, expected_grad_w, atol=0.5, rtol=0.05) + torch.testing.assert_close( + br.grad, g.sum(0).to(torch.bfloat16), atol=0.5, rtol=0.05 + ) + finally: + state.reset() @skip_no_fp8 @@ -279,24 +220,24 @@ def test_fp8_linear_backward_outside_autocast(): bias = torch.randn(96, device="cuda", dtype=torch.bfloat16, requires_grad=True) xr, wr, br = (t.detach().clone().requires_grad_() for t in (x, weight, bias)) - calls = {"bwd": 0} - orig = f8mod.linear_backward_fp8 + calls = {"fwd": 0} + orig = f8mod.fp8_linear_forward def spy(*args, **kwargs): - calls["bwd"] += 1 + calls["fwd"] += 1 return orig(*args, **kwargs) - f8mod.linear_backward_fp8 = spy + f8mod.fp8_linear_forward = spy try: with fp8_autocast(enabled=True): out = F.linear(x, weight, bias) assert type(out.grad_fn).__name__ == "_LinearFp8Backward" out.float().pow(2).sum().backward() # outside the autocast region finally: - f8mod.linear_backward_fp8 = orig + f8mod.fp8_linear_forward = orig f8mod.fp8_state().reset() - assert calls["bwd"] == 1 # fp8 kernels, not the bf16 fallback + assert calls["fwd"] == 1 # fp8 kernels, not the bf16 fallback ref = F.linear(xr, wr, br) ref.float().pow(2).sum().backward() @@ -319,15 +260,15 @@ def test_mm_fp8_matches_scaled_mm(): m, n, k = 512, 4096, 4096 a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16) - sa = torch.tensor([2.5], device="cuda") - sb = torch.tensor([1.5], device="cuda") - a8, _ = quantize_bf16(a, sa, "e4m3") - b8, _ = quantize_bf16(b, sb, "e4m3") - out = mm_fp8(a8, b8, sa, sb) + sa = _scale(a) + sb = _scale(b) + a8, _ = quantize(a, sa.reciprocal(), "e4m3") + b8, _ = quantize(b, sb.reciprocal(), "e4m3") + out = mm_fp8(a8, b8, sa * sb) assert out.dtype == torch.bfloat16 assert out.shape == (m, n) - ref = (a8.float().double() @ b8.float().double() * 2.5 * 1.5).to(torch.bfloat16) + ref = (a8.float().double() @ b8.float().double() * sa * sb).to(torch.bfloat16) torch.testing.assert_close(out, ref, atol=6.0, rtol=0.05) try: @@ -342,30 +283,6 @@ def test_mm_fp8_matches_scaled_mm(): ) -@skip_no_fp8 -def test_mm_fp8_fp8_output(): - """mm_fp8 with out_dtype='e4m3' produces an FP8 output (layer-to-layer).""" - torch.manual_seed(12) - m, n, k = 256, 128, 64 - a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) - b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16) - sa = torch.tensor([2.0], device="cuda") - sb = torch.tensor([1.0], device="cuda") - os_ = torch.tensor([0.5], device="cuda") - a8, _ = quantize_bf16(a, sa, "e4m3") - b8, _ = quantize_bf16(b, sb, "e4m3") - out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_) - assert out8.dtype == torch.float8_e4m3fn - assert out8.shape == (m, n) - - ref = (a8.float().double() @ b8.float().double() * 2.0 * 1.0 * 0.5).to( - torch.bfloat16 - ) - torch.testing.assert_close( - out8.float().to(torch.bfloat16), ref, atol=6.0, rtol=0.05 - ) - - # -------------------------------------------------------------------------- # Policy-level (CPU-verifiable) # -------------------------------------------------------------------------- @@ -423,23 +340,26 @@ 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 | counter] packing: views alias the single state buffer. + # [hist | scale] packing: views alias the single state buffer. assert meta.w.state.numel() == 4 + 2 assert meta.w.hist.data_ptr() == meta.w.state.data_ptr() assert meta.w.scale.data_ptr() == meta.w.state[4:].data_ptr() - # counter slot stays int32-zero (float bits) between launches - assert meta.w.state[-1].view(torch.int32).item() == 0 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])) -def test_quantize_bf16_cpu_fallback(): + +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]) - x8, amax = quantize_bf16(x, scale, "e4m3") + 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) + ref = (x.float() * 0.5).to(torch.float8_e4m3fn) assert torch.equal(x8, ref) torch.testing.assert_close(amax, x.abs().amax().float().reshape(1)) @@ -447,26 +367,12 @@ def test_quantize_bf16_cpu_fallback(): 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) - sa = torch.tensor([2.0]) - sb = torch.tensor([0.5]) - out = mm_fp8(a8, b8, sa, sb) - ref = (a8.float() @ b8.float() * 2.0 * 0.5).to(torch.bfloat16) + 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) -def test_mm_fp8_fp8_output_cpu(): - """CPU fallback with an FP8 output (out_dtype='e4m3' + out_scale).""" - a8 = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn) - b8 = torch.tensor([[3.0], [4.0]], dtype=torch.float8_e4m3fn) - sa = torch.tensor([2.0]) - sb = torch.tensor([0.5]) - os_ = torch.tensor([0.25]) - out8 = mm_fp8(a8, b8, sa, sb, out_dtype="e4m3", out_scale=os_) - assert out8.dtype == torch.float8_e4m3fn - ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn) - assert torch.equal(out8, ref) - - # -------------------------------------------------------------------------- # torch-autocast parity: context semantics (nesting, thread locality, switch) # -------------------------------------------------------------------------- @@ -488,7 +394,7 @@ def test_nested_disabled_region_redispatches_bf16(): resumes when it exits.""" x, w = _linear() with fp8_autocast(enabled=True): - out_fp8 = F.linear(x, w) + F.linear(x, w) with fp8_autocast(enabled=False): out_bf16 = F.linear(x, w) assert type(out_bf16.grad_fn).__name__ != "_LinearFp8Backward"