perf: fold the delayed-scaling ring update into the quantize kernel
- the kernel's last block folds amax into the history window and publishes the next scale in-kernel (atomicAdd ticket + fences), replacing the host update chain - quantize bindings split into quantize(transposed) / quantize_dual with fixed arities and a QuantLayout enum; the python adapter becomes a thin attention-style wrapper over pybind (Optional ring_state at the boundary, no torch.library custom_ops) - tests: in-kernel fold vs host reference (exact), dual/transposed orientation byte-equality Benchmark: L20 (sm_89), 1.2B model, full train step. Per-linear fixed overhead 28.8us -> 8.8us; fp8 vs bf16: M=512 77.5ms, M=2048 144.5ms (1.15x), M=8192 527.4ms (1.28x); losses bit-identical.
This commit is contained in:
+65
-200
@@ -1,14 +1,20 @@
|
||||
"""FP8 CUDA kernel interface adapter (the only module touching the pybind).
|
||||
|
||||
Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives:
|
||||
Attention-style thin wrappers: one Python entry per binding, called directly
|
||||
— no torch.library dispatch layer. Optional arguments (``ring_state``,
|
||||
``bias``) keep native Optional semantics at the pybind boundary, and
|
||||
in-place buffer updates (the delayed-scaling ring fold, like attention's
|
||||
KV-cache appends) happen on-stream without mutation declarations. CUDA-only:
|
||||
non-CUDA or unsupported inputs raise from the binding's TORCH_CHECKs.
|
||||
|
||||
- ``quantize(x, scale, fmt) -> (x8, amax)`` — BF16/FP16/FP32 → FP8 with fused amax
|
||||
- ``quantize(x, scale, fmt, transposed=False) -> (x8|x8T, amax)`` — BF16/FP16/FP32
|
||||
→ FP8 with fused amax (``transposed`` picks the orientation; arity is fixed)
|
||||
- ``quantize_dual(x, scale, fmt) -> (x8, x8T, amax)`` — both orientations, one read
|
||||
- ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output)
|
||||
|
||||
Scale semantics: scales are *quantization steps* — the value divided out when
|
||||
quantizing (``x8 = x / scale``). Every primitive computes its own inverse
|
||||
internally; callers never pass ``scale_inv``. ``amax`` values are *returned*,
|
||||
never passed as output arguments. ``fmt`` is ``"e4m3"`` or ``"e5m2"``.
|
||||
``scale`` is the quantization multiplier (device scalar); ``fmt`` is
|
||||
``"e4m3"`` or ``"e5m2"``. ``amax`` values are *returned*, never passed as
|
||||
output arguments.
|
||||
|
||||
Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
|
||||
this module is stateless.
|
||||
@@ -17,7 +23,6 @@ this module is stateless.
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
from torch.library import custom_op
|
||||
|
||||
from astrai.extension.loader import get_module
|
||||
|
||||
@@ -32,192 +37,63 @@ def _fmt_int(fmt: str) -> int:
|
||||
raise ValueError(f"unsupported fp8 format {fmt!r} (expected 'e4m3' or 'e5m2')")
|
||||
|
||||
|
||||
def _fmt_name(fmt: int) -> str:
|
||||
if fmt == 0:
|
||||
return "e4m3"
|
||||
if fmt == 1:
|
||||
return "e5m2"
|
||||
raise ValueError(f"unsupported quantization type {fmt!r}")
|
||||
|
||||
|
||||
def _fmt_dtype(fmt: str) -> torch.dtype:
|
||||
return torch.float8_e5m2 if _fmt_int(fmt) else torch.float8_e4m3fn
|
||||
|
||||
|
||||
@custom_op("custom::fp8_quantize", mutates_args=())
|
||||
def fp8_quantize(
|
||||
x: torch.Tensor, scale: torch.Tensor, fmt: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; ``scale`` is a multiplier."""
|
||||
|
||||
|
||||
@fp8_quantize.register_fake
|
||||
def _fp8_quantize_fake(x, scale, fmt):
|
||||
dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
|
||||
return (
|
||||
torch.empty(x.shape, device=x.device, dtype=dtype),
|
||||
torch.empty(1, device=x.device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
_QUANT_INPUT_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
|
||||
|
||||
|
||||
@custom_op("custom::fp8_quantize_t", mutates_args=())
|
||||
def fp8_quantize_t(
|
||||
x: torch.Tensor, scale: torch.Tensor, fmt: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Transposed-output variant of fp8_quantize: returns ``(x8T, amax)``
|
||||
where ``x8T`` is the [cols][rows] row-major transpose of the quantized
|
||||
input (the K-contiguous operand orientation for NT GEMMs)."""
|
||||
|
||||
|
||||
@fp8_quantize_t.register_fake
|
||||
def _fp8_quantize_t_fake(x, scale, fmt):
|
||||
dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
|
||||
rows, cols = x.shape[-2], x.shape[-1]
|
||||
return (
|
||||
torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
|
||||
torch.empty(1, device=x.device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
@fp8_quantize_t.register_kernel("cuda")
|
||||
def _fp8_quantize_t_cuda(x, scale, fmt):
|
||||
if x.dtype not in _QUANT_INPUT_DTYPES:
|
||||
raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
|
||||
return get_module("fp8_ops").quantize(x, scale, int(fmt), 1)
|
||||
|
||||
|
||||
@fp8_quantize_t.register_kernel("cpu")
|
||||
def _fp8_quantize_t_cpu(x, scale, fmt):
|
||||
x8, amax = _fp8_quantize_cpu(x, scale, fmt)
|
||||
return x8.transpose(-2, -1).contiguous(), amax
|
||||
|
||||
|
||||
@custom_op("custom::fp8_quantize_dual", mutates_args=())
|
||||
def fp8_quantize_dual(
|
||||
x: torch.Tensor, scale: torch.Tensor, fmt: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Dual-orientation quantize: one read of ``x`` produces both the
|
||||
row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
|
||||
consumed by GEMMs on both orientations (backward ``g``)."""
|
||||
|
||||
|
||||
@fp8_quantize_dual.register_fake
|
||||
def _fp8_quantize_dual_fake(x, scale, fmt):
|
||||
dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
|
||||
rows, cols = x.shape[-2], x.shape[-1]
|
||||
return (
|
||||
torch.empty(x.shape, device=x.device, dtype=dtype),
|
||||
torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
|
||||
torch.empty(1, device=x.device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
@fp8_quantize_dual.register_kernel("cuda")
|
||||
def _fp8_quantize_dual_cuda(x, scale, fmt):
|
||||
if x.dtype not in _QUANT_INPUT_DTYPES:
|
||||
raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
|
||||
return get_module("fp8_ops").quantize(x, scale, int(fmt), 2)
|
||||
|
||||
|
||||
@fp8_quantize_dual.register_kernel("cpu")
|
||||
def _fp8_quantize_dual_cpu(x, scale, fmt):
|
||||
x8, amax = _fp8_quantize_cpu(x, scale, fmt)
|
||||
return x8, x8.transpose(-2, -1).contiguous(), amax
|
||||
|
||||
|
||||
@fp8_quantize.register_kernel("cuda")
|
||||
def _fp8_quantize_cuda(x, scale, fmt):
|
||||
if x.dtype not in _QUANT_INPUT_DTYPES:
|
||||
raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
|
||||
return get_module("fp8_ops").quantize(x, scale, int(fmt))
|
||||
|
||||
|
||||
@fp8_quantize.register_kernel("cpu")
|
||||
def _fp8_quantize_cpu(x, scale, fmt):
|
||||
x8 = (x.float() * scale).to(_fmt_dtype(_fmt_name(fmt)))
|
||||
amax = x.abs().amax().float().reshape(1).clamp_min(1e-12)
|
||||
return x8, amax
|
||||
|
||||
|
||||
@custom_op("custom::fp8_gemm", mutates_args=())
|
||||
def fp8_gemm(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
trans_a: int = 0,
|
||||
trans_b: int = 0,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""FP8 GEMM: ``a @ b * scale (+ bias)`` with FP32 accumulation.
|
||||
|
||||
2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules).
|
||||
``bias`` (bf16, length n) fuses into the epilogue in fp32 before the
|
||||
single bf16 rounding. The result is always BF16; FP8 output is a
|
||||
separate quantize operation.
|
||||
"""
|
||||
|
||||
|
||||
@fp8_gemm.register_fake
|
||||
def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0, bias=None):
|
||||
dtype = torch.bfloat16
|
||||
rows = a.size(2) if trans_a else a.size(1)
|
||||
cols = b.size(1) if trans_b else b.size(2)
|
||||
batches = [t.size(0) for t in (a, b) if t.dim() == 3]
|
||||
shape = (max(batches), rows, cols) if batches else (rows, cols)
|
||||
return torch.empty(shape, device=a.device, dtype=dtype)
|
||||
|
||||
|
||||
@fp8_gemm.register_kernel("cuda")
|
||||
def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0, bias=None):
|
||||
if a.dtype != b.dtype or a.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||||
raise TypeError(
|
||||
f"fp8 GEMM requires matching fp8 inputs, got {a.dtype}/{b.dtype}"
|
||||
)
|
||||
return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
|
||||
|
||||
|
||||
@fp8_gemm.register_kernel("cpu")
|
||||
def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
|
||||
aa = a.float().transpose(-2, -1) if trans_a else a.float()
|
||||
bb = b.float().transpose(-2, -1) if trans_b else b.float()
|
||||
acc = aa @ bb * scale
|
||||
if bias is not None and bias.numel() > 0:
|
||||
acc = acc + bias.float()
|
||||
return acc.to(torch.bfloat16)
|
||||
|
||||
|
||||
def quantize(
|
||||
x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3", layout: int = 0
|
||||
) -> tuple:
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
fmt: str = "e4m3",
|
||||
transposed: bool = False,
|
||||
ring_state: Optional[torch.Tensor] = None,
|
||||
hist_idx: int = 0,
|
||||
fp8_max: float = 448.0,
|
||||
pow2_margin: float = 1.0,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Float (bf16/fp16/fp32) -> FP8 quantize with fused amax.
|
||||
|
||||
``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
|
||||
E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. ``layout``
|
||||
picks the output orientation: 0 = row-major ``(x8, amax)``; 1 =
|
||||
transposed ``[cols][rows]`` ``(x8T, amax)`` — the K-contiguous operand
|
||||
orientation NT GEMMs want; 2 = both from one read ``(x8, x8T, amax)``
|
||||
(for tensors consumed in both orientations, e.g. backward ``g``).
|
||||
E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
|
||||
``transposed=True`` swaps ``x8`` for ``x8T``, the ``[cols][rows]``
|
||||
row-major transpose of the quantized input — the K-contiguous operand
|
||||
orientation NT GEMMs want — at the same 2-tuple arity.
|
||||
|
||||
``ring_state`` (a 1D float32 CUDA buffer laid out
|
||||
``[hist n | scale | legacy | amax | done]``) switches on the in-kernel
|
||||
delayed-scaling fold: the kernel's last block folds the amax into
|
||||
``hist[hist_idx]`` and publishes the next scale as
|
||||
``max(hist) / fp8_max / pow2_margin`` — the returned ``amax`` is then the
|
||||
self-cleaned persistent slot (reads zero). None keeps the classic
|
||||
fresh-amax return.
|
||||
"""
|
||||
# Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a
|
||||
# 512-wide GEMM): real CUDA tensors of a supported dtype go straight to
|
||||
# the extension. Fake/subclass tensors and non-CUDA inputs keep the
|
||||
# custom_op route so torch.compile / meta / fake-tensor tracing and the
|
||||
# CPU fallback behave exactly as before.
|
||||
if (
|
||||
type(x) is torch.Tensor
|
||||
and x.is_cuda
|
||||
and x.dtype in _QUANT_INPUT_DTYPES
|
||||
and fmt in _FMT_TO_INT
|
||||
):
|
||||
return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt], layout)
|
||||
if layout == 0:
|
||||
return fp8_quantize(x, scale, _fmt_int(fmt))
|
||||
if layout == 1:
|
||||
return fp8_quantize_t(x, scale, _fmt_int(fmt))
|
||||
return fp8_quantize_dual(x, scale, _fmt_int(fmt))
|
||||
return get_module("fp8_ops").quantize(
|
||||
x,
|
||||
scale,
|
||||
_fmt_int(fmt),
|
||||
transposed,
|
||||
ring_state,
|
||||
hist_idx,
|
||||
fp8_max,
|
||||
pow2_margin,
|
||||
)
|
||||
|
||||
|
||||
def quantize_dual(
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
fmt: str = "e4m3",
|
||||
ring_state: Optional[torch.Tensor] = None,
|
||||
hist_idx: int = 0,
|
||||
fp8_max: float = 448.0,
|
||||
pow2_margin: float = 1.0,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Dual-orientation quantize: one read of ``x`` produces both the
|
||||
row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
|
||||
consumed by GEMMs in both orientations (backward ``g``).
|
||||
|
||||
``ring_state`` switches on the in-kernel delayed-scaling fold exactly as
|
||||
in :func:`quantize`.
|
||||
"""
|
||||
return get_module("fp8_ops").quantize_dual(
|
||||
x, scale, _fmt_int(fmt), ring_state, hist_idx, fp8_max, pow2_margin
|
||||
)
|
||||
|
||||
|
||||
def mm_fp8(
|
||||
@@ -237,15 +113,4 @@ def mm_fp8(
|
||||
kernel epilogue in fp32 — no separate elementwise pass. The result is
|
||||
BF16; FP8 output is a separate quantize operation.
|
||||
"""
|
||||
# Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep
|
||||
# validation identical on the direct route (bias may be None — the
|
||||
# binding resolves it to the no-bias path).
|
||||
if (
|
||||
type(a) is torch.Tensor
|
||||
and a.is_cuda
|
||||
and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
|
||||
):
|
||||
return get_module("fp8_ops").mm_fp8(
|
||||
a, b, scale, int(trans_a), int(trans_b), bias
|
||||
)
|
||||
return fp8_gemm(a, b, scale, trans_a, trans_b, bias)
|
||||
return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
|
||||
|
||||
Reference in New Issue
Block a user