ViperEkura
7540acb43e
- replace the per-shape auto tables in the linear backend with an M-banded rule (M in [2,4] on compute capability 8.0+) that measured at the HBM bandwidth floor across every family, and fold the capability check into the capable guard - drop the unreachable swiglu auto shape-table machinery so both backends share one env-mode ladder via the new dispatch.env_mode helper - add __all__ across extension modules, name the rotary registration records, and unify typing to the typing-module style - rewrite test_linear_dispatch.py around behavioral routing assertions and document the M-banded policy in the developer docs - Benchmark: L20 SM89, Python dispatch overhead 2.9us to 1.5us, auto now covers every projection shape at M in [2,4].
120 lines
4.3 KiB
Python
120 lines
4.3 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)
|
|
|
|
|
|
__all__ = ["mm_fp8", "quantize", "quantize_dual"]
|