refactor: split quantize into multi-type primitive
- split quantize into quantize.cuh, templated on input type (bf16/fp16/fp32) - rename pybind entry quantize_bf16 to quantize; validate the fmt enum - fix fp8x2 packing: one 32-bit word packs two pairs (halves were dropped) - drop the dead OutFp8 template param; GEMM output is always bf16 - fp8_state.reset() restores recipe/format defaults too (test state leak) - rewrite tests for the two-primitive API with fp32-domain amax references
This commit is contained in:
+53
-67
@@ -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
|
||||
|
||||
|
||||
|
||||
+50
-149
@@ -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)
|
||||
|
||||
+11
-35
@@ -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;
|
||||
|
||||
+21
-174
@@ -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 <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
@@ -46,128 +47,11 @@ struct fp8_input<FP8Format::E5M2> {
|
||||
// FP8 MMA lives in the shared astrai::mma_sync template (common/mma.cuh);
|
||||
// instantiate it with fp8_input<Fmt>::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 <FP8Format Fmt>
|
||||
__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<unsigned>(__nv_cvt_float2_to_fp8x2(
|
||||
make_float2(lo * inv, hi * inv), __NV_SATFINITE, kFmt));
|
||||
}
|
||||
|
||||
template <FP8Format Fmt>
|
||||
__global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
|
||||
const float inv = 1.0f / *p.scale_a;
|
||||
const auto* x = reinterpret_cast<const __nv_bfloat16*>(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<uintptr_t>(x) | reinterpret_cast<uintptr_t>(x8)) & 15) ==
|
||||
0;
|
||||
const int64_t total_vec = aligned ? p.total / 8 : 0;
|
||||
const uint4* xv = reinterpret_cast<const uint4*>(x);
|
||||
uint2* o8 = reinterpret_cast<uint2*>(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<Fmt>(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<int*>(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 <typename Traits, bool OutFp8 = false,
|
||||
typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
|
||||
template <typename Traits, typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
|
||||
bool kBStaged = true>
|
||||
__global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB,
|
||||
@@ -495,7 +378,6 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
const auto* a = reinterpret_cast<const T8*>(p.a_ptr);
|
||||
const auto* b = reinterpret_cast<const T8*>(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<const __nv_bfloat16*>(p.bias);
|
||||
const auto* bias8 = static_cast<const T8*>(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<unsigned short*>(out_fp8 + row * n + col) =
|
||||
static_cast<unsigned short>(__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<uintptr_t>(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<uintptr_t>(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 <FP8Format Fmt>
|
||||
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<Fmt><<<blocks, kThreads, 0, stream>>>(p);
|
||||
}
|
||||
|
||||
// Launch one kernel instantiation with its shared-memory budget: stages live
|
||||
// in dynamic smem, so budgets beyond the 48KB static limit opt in once per
|
||||
// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
|
||||
@@ -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 <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor,
|
||||
template <FP8Format Fmt, typename LayoutA = RowMajor,
|
||||
typename LayoutB = RowMajor, int kK = 64, int Stages = 2,
|
||||
bool GroupRaster = std::is_same_v<LayoutA, ColMajor> || std::is_same_v<LayoutB, ColMajor>>
|
||||
void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
|
||||
@@ -860,24 +707,24 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
|
||||
if (p.m <= 64) {
|
||||
using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
|
||||
if (b_staged)
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, true>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
else
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, false>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
} else {
|
||||
using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
|
||||
if (b_staged)
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, true>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
else
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB,
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, false>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
|
||||
+94
-433
@@ -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 <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <mutex>
|
||||
#include <tuple>
|
||||
#include <unordered_map>
|
||||
|
||||
#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<int, bool> supported;
|
||||
const int device = tensor.device().index();
|
||||
{
|
||||
std::lock_guard<std::mutex> 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<std::mutex> 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<float>();
|
||||
p.amax = amax.data_ptr<float>();
|
||||
p.total = static_cast<int>(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<float>();
|
||||
p.scale_b = sb.data_ptr<float>();
|
||||
p.out_scale = out_scale ? out_scale->data_ptr<float>() : nullptr;
|
||||
p.bias = bias;
|
||||
p.bias_scale = bias_scale ? bias_scale->data_ptr<float>() : nullptr;
|
||||
p.out_ptr = output;
|
||||
p.scale = scale.data_ptr<float>();
|
||||
p.m = static_cast<int>(m);
|
||||
p.n = static_cast<int>(n);
|
||||
p.k = static_cast<int>(k);
|
||||
@@ -78,443 +70,112 @@ void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
|
||||
p.b_ld = static_cast<int>(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<float>();
|
||||
p.amax_a = amax ? amax->data_ptr<float>() : 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<float>();
|
||||
p.ring_len = static_cast<int>(ring->numel() - 2);
|
||||
p.ring_idx = static_cast<int>(ring_idx);
|
||||
p.ring_margin = static_cast<int>(ring_margin);
|
||||
}
|
||||
p.total = static_cast<int>(total);
|
||||
}
|
||||
|
||||
// ---- GEMM launch dispatch (runtime flags -> compile-time kernel variants) ----
|
||||
|
||||
template <FP8Format Fmt, int Variant>
|
||||
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<Fmt, out_fp8, LayoutA, LayoutB>(p, stream);
|
||||
launch_fp8_gemm<Fmt, LayoutA, LayoutB>(p, stream);
|
||||
}
|
||||
|
||||
template <FP8Format Fmt>
|
||||
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<int>(out_fp8) << 2) |
|
||||
(static_cast<int>(trans_a) << 1) |
|
||||
void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a,
|
||||
bool trans_b) {
|
||||
const int variant = (static_cast<int>(trans_a) << 1) |
|
||||
static_cast<int>(trans_b);
|
||||
switch (variant) {
|
||||
case 0: launch_gemm_variant<Fmt, 0>(p, stream); break;
|
||||
case 1: launch_gemm_variant<Fmt, 1>(p, stream); break;
|
||||
case 2: launch_gemm_variant<Fmt, 2>(p, stream); break;
|
||||
case 3: launch_gemm_variant<Fmt, 3>(p, stream); break;
|
||||
case 4: launch_gemm_variant<Fmt, 4>(p, stream); break;
|
||||
case 5: launch_gemm_variant<Fmt, 5>(p, stream); break;
|
||||
case 6: launch_gemm_variant<Fmt, 6>(p, stream); break;
|
||||
case 7: launch_gemm_variant<Fmt, 7>(p, stream); break;
|
||||
case 0: launch_variant<Fmt, 0>(p, stream); break;
|
||||
case 1: launch_variant<Fmt, 1>(p, stream); break;
|
||||
case 2: launch_variant<Fmt, 2>(p, stream); break;
|
||||
case 3: launch_variant<Fmt, 3>(p, stream); break;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Entry points
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor> 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<int64_t>(FP8Format::E4M3) ||
|
||||
fmt == static_cast<int64_t>(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<FP8Format::E5M2>(p, stream.stream());
|
||||
pack_quantize(p, input.data_ptr(), output.data_ptr(), scale, amax,
|
||||
input.numel());
|
||||
const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
|
||||
if (x.scalar_type() == torch::kHalf) {
|
||||
if (e5m2)
|
||||
launch_fp8_quantize<FP8Format::E5M2, __half>(p, stream.stream());
|
||||
else
|
||||
launch_fp8_quantize<FP8Format::E4M3, __half>(p, stream.stream());
|
||||
} else if (x.scalar_type() == torch::kFloat32) {
|
||||
if (e5m2)
|
||||
launch_fp8_quantize<FP8Format::E5M2, float>(p, stream.stream());
|
||||
else
|
||||
launch_fp8_quantize<FP8Format::E4M3, float>(p, stream.stream());
|
||||
} else {
|
||||
launch_fp8_quantize<FP8Format::E4M3>(p, stream.stream());
|
||||
if (e5m2)
|
||||
launch_fp8_quantize<FP8Format::E5M2, __nv_bfloat16>(
|
||||
p, stream.stream());
|
||||
else
|
||||
launch_fp8_quantize<FP8Format::E4M3, __nv_bfloat16>(
|
||||
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<torch::Tensor> 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<FP8Format::E4M3>(p, stream.stream(), out_fp8, ta, tb);
|
||||
dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), ta, tb);
|
||||
else
|
||||
dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), out_fp8, ta, tb);
|
||||
dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), ta, tb);
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
return out;
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
|
||||
torch::Tensor>
|
||||
linear_forward_fp8(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
|
||||
torch::Tensor sx, torch::Tensor sw, int64_t fmt,
|
||||
c10::optional<torch::Tensor> bias_scale,
|
||||
c10::optional<torch::Tensor> x_ring, int64_t x_ring_idx,
|
||||
int64_t x_ring_margin, c10::optional<torch::Tensor> w_ring,
|
||||
int64_t w_ring_idx, int64_t w_ring_margin) {
|
||||
// Pure FP8 forward: quantize x/w (fmt: 0 = E4M3, 1 = E5M2), then the
|
||||
// pre-quantized GEMM; the dequantized BF16 output gets the bias added.
|
||||
// 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<torch::Tensor>& 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<FP8Format::E5M2>(qp, stream.stream());
|
||||
} else {
|
||||
launch_fp8_quantize<FP8Format::E4M3>(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<FP8Format::E5M2, false, RowMajor, ColMajor>(
|
||||
p, stream.stream());
|
||||
} else {
|
||||
launch_fp8_gemm<FP8Format::E4M3, false, RowMajor, ColMajor>(
|
||||
p, stream.stream());
|
||||
}
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
|
||||
shape.push_back(n);
|
||||
return {out.reshape(shape), x8, w8, amax_x, amax_w};
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
|
||||
linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
|
||||
std::vector<int64_t> masks, torch::Tensor sg,
|
||||
torch::Tensor sw, torch::Tensor sx, int64_t fmt,
|
||||
c10::optional<torch::Tensor> 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<torch::Tensor>& 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<FP8Format::E5M2>(qp, stream.stream());
|
||||
} else {
|
||||
launch_fp8_quantize<FP8Format::E4M3>(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<FP8Format::E5M2>(gp, stream.stream(), false, trans_a,
|
||||
trans_b);
|
||||
else
|
||||
dispatch_gemm<FP8Format::E4M3>(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);
|
||||
}
|
||||
|
||||
@@ -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 <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
|
||||
#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 <typename InT>
|
||||
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<const __half2*>(&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<float> {
|
||||
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 <FP8Format Fmt>
|
||||
__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<unsigned>(__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 <FP8Format Fmt, typename InT>
|
||||
__global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
|
||||
const float mult = *p.scale;
|
||||
const auto* x = static_cast<const InT*>(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<InT>::kVecElems;
|
||||
const bool aligned =
|
||||
((reinterpret_cast<uintptr_t>(x) |
|
||||
reinterpret_cast<uintptr_t>(x8)) &
|
||||
15) == 0;
|
||||
const int64_t total_vec = aligned ? p.total / kVecElems : 0;
|
||||
const uint4* xv = reinterpret_cast<const uint4*>(x);
|
||||
for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total_vec;
|
||||
i += stride) {
|
||||
float f[kVecElems];
|
||||
quant_in_traits<InT>::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<Fmt>(f[4 * j] * mult, f[4 * j + 1] * mult);
|
||||
const unsigned hi =
|
||||
cvt_fp8x2<Fmt>(f[4 * j + 2] * mult, f[4 * j + 3] * mult);
|
||||
packed[j] = (lo & 0xffffu) | (hi << 16);
|
||||
}
|
||||
if constexpr (kVecElems == 8)
|
||||
reinterpret_cast<uint2*>(x8)[i] =
|
||||
make_uint2(packed[0], packed[1]);
|
||||
else
|
||||
reinterpret_cast<unsigned*>(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<InT>::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 <FP8Format Fmt, typename InT>
|
||||
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<InT>::kVecElems;
|
||||
int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
|
||||
if (blocks < 1) blocks = 1;
|
||||
fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
|
||||
}
|
||||
|
||||
} // namespace fp8
|
||||
} // namespace astrai
|
||||
+6
-10
@@ -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<unsigned char> 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<FP8Format::E4M3, false, LA, LB, kK, Stages>(p, 0);
|
||||
launch_fp8_gemm<FP8Format::E4M3, LA, LB, kK, Stages>(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;
|
||||
}
|
||||
|
||||
|
||||
@@ -47,16 +47,17 @@ style as attention, but split into **three** files:
|
||||
| File | Role |
|
||||
|------|------|
|
||||
| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `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<Fmt, InT>` (bf16/fp16/fp32 → FP8 + amax, `quant_in_traits<InT>` 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.
|
||||
|
||||
+115
-209
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user