ViperEkura
962c10c52b
- 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.
117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
"""FP8 CUDA kernel interface adapter (the only module touching the pybind).
|
|
|
|
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, 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`` 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.
|
|
"""
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
|
|
from astrai.extension.loader import get_module
|
|
|
|
# fmt string -> kernel int (0 = E4M3, 1 = E5M2)
|
|
_FMT_TO_INT = {"e4m3": 0, "e5m2": 1}
|
|
|
|
|
|
def _fmt_int(fmt: str) -> int:
|
|
try:
|
|
return _FMT_TO_INT[fmt]
|
|
except KeyError:
|
|
raise ValueError(f"unsupported fp8 format {fmt!r} (expected 'e4m3' or 'e5m2')")
|
|
|
|
|
|
def quantize(
|
|
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.
|
|
``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.
|
|
"""
|
|
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(
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
trans_a: bool = False,
|
|
trans_b: bool = False,
|
|
bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
"""Pre-quantized FP8 GEMM: ``a @ b * scale (+ bias)``.
|
|
|
|
``a``/``b`` must be FP8 tensors of the same format, 2D or 3D (batched,
|
|
matmul-style broadcast on the batch dim). Inner-transposed views (e.g.
|
|
``x.t()``) fold into the layout at zero copy. ``scale`` is their combined
|
|
dequantization scale. ``bias`` (CUDA bf16 1D of length n) adds inside the
|
|
kernel epilogue in fp32 — no separate elementwise pass. The result is
|
|
BF16; FP8 output is a separate quantize operation.
|
|
"""
|
|
return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
|