refactor: reorganize CUDA kernels into per-family directories
- move attention kernels to csrc/kernels/attention/ and rotary to rotary/ - add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks) - split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu - fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads - update extension ops, loader, and kernel tests
This commit is contained in:
+151
-94
@@ -1,9 +1,16 @@
|
||||
"""FP8 CUDA kernel interface adapter (the only module touching the pybind.
|
||||
"""FP8 CUDA kernel interface adapter (the only module touching the pybind).
|
||||
|
||||
Isolates the ``fp8_mm`` CUDA extension behind stable Python functions:
|
||||
- availability / dtype checks and clear errors
|
||||
- torch.library ``custom::fp8_mm`` registration (meta + CPU fallback)
|
||||
- quantize-in-GEMM primitives used by ``fp8.py`` training state
|
||||
Isolates the ``fp8_mm`` CUDA extension behind stable Python primitives:
|
||||
|
||||
- ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → 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, 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
|
||||
internally; callers never pass ``scale_inv``. ``amax`` values are *returned*,
|
||||
never passed as output arguments. ``fmt`` is ``"e4m3"`` or ``"e5m2"``.
|
||||
|
||||
Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
|
||||
this module is stateless.
|
||||
@@ -14,106 +21,158 @@ from torch.library import custom_op
|
||||
|
||||
from astrai.extension.loader import get_module, is_available
|
||||
|
||||
# fmt string -> kernel int (0 = E4M3, 1 = E5M2)
|
||||
_FMT_TO_INT = {"e4m3": 0, "e5m2": 1}
|
||||
|
||||
def _mod():
|
||||
if not is_available("fp8_mm"):
|
||||
raise RuntimeError(
|
||||
"CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true."
|
||||
)
|
||||
return get_module("fp8_mm")
|
||||
# The pybind module is loaded once at first use and cached: the loader
|
||||
# resolves modules at import time and never reloads them, so every call
|
||||
# after the first is a single None check.
|
||||
_MOD: object | None = None
|
||||
|
||||
|
||||
@custom_op("custom::fp8_mm", mutates_args=())
|
||||
def fp8_mm(
|
||||
a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""BF16 inputs, fused FP8 GEMM with FP32 accumulation and BF16 output."""
|
||||
def _mod() -> object:
|
||||
global _MOD
|
||||
if _MOD is None:
|
||||
if not is_available("fp8_mm"):
|
||||
raise RuntimeError(
|
||||
"CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true."
|
||||
)
|
||||
_MOD = get_module("fp8_mm")
|
||||
return _MOD
|
||||
|
||||
|
||||
@fp8_mm.register_fake
|
||||
def _fp8_mm_fake(a, b, sx, sw):
|
||||
return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
|
||||
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')")
|
||||
|
||||
|
||||
@fp8_mm.register_kernel("cuda")
|
||||
def _fp8_mm_cuda(a, b, sx, sw):
|
||||
if not (a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16):
|
||||
raise TypeError(f"bf16 GEMM requires bf16 inputs, got {a.dtype}/{b.dtype}")
|
||||
return _mod().fp8_mm(a, b, sx, sw)
|
||||
def _fmt_dtype(fmt: str) -> torch.dtype:
|
||||
return torch.float8_e5m2 if _fmt_int(fmt) else torch.float8_e4m3fn
|
||||
|
||||
|
||||
@fp8_mm.register_kernel("cpu")
|
||||
def _fp8_mm_cpu(a, b, sx, sw):
|
||||
return torch.mm(a.float(), b.float().t()).to(torch.bfloat16)
|
||||
@custom_op("custom::fp8_quantize", mutates_args=())
|
||||
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)``."""
|
||||
|
||||
|
||||
@custom_op("custom::fp8_mm_prequant", mutates_args=())
|
||||
def fp8_mm_prequant(
|
||||
a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""Pre-quantized FP8 inputs, fused FP8 GEMM, FP32 accumulation, BF16 out."""
|
||||
|
||||
|
||||
@fp8_mm_prequant.register_fake
|
||||
def _fp8_mm_prequant_fake(a, b, scale):
|
||||
return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
|
||||
|
||||
|
||||
@fp8_mm_prequant.register_kernel("cuda")
|
||||
def _fp8_mm_prequant_cuda(a, b, scale):
|
||||
if not (a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn):
|
||||
raise TypeError(
|
||||
f"pre-quantized FP8 GEMM requires fp8 inputs, got {a.dtype}/{b.dtype}"
|
||||
)
|
||||
return _mod().fp8_mm_prequant(a, b, scale)
|
||||
|
||||
|
||||
@fp8_mm_prequant.register_kernel("cpu")
|
||||
def _fp8_mm_prequant_cpu(a, b, scale):
|
||||
return (a.float() @ b.float().t() * scale).to(torch.bfloat16)
|
||||
|
||||
|
||||
@custom_op("custom::fp8_mm_prequant_fp8", mutates_args=())
|
||||
def fp8_mm_prequant_fp8(
|
||||
a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor, out_scale: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""FP8 inputs and FP8 output: fused FP8 GEMM with FP32 accumulation."""
|
||||
|
||||
|
||||
@fp8_mm_prequant_fp8.register_fake
|
||||
def _fp8_mm_prequant_fp8_fake(a, b, scale, out_scale):
|
||||
return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=a.dtype)
|
||||
|
||||
|
||||
@fp8_mm_prequant_fp8.register_kernel("cuda")
|
||||
def _fp8_mm_prequant_fp8_cuda(a, b, scale, out_scale):
|
||||
if not (a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn):
|
||||
raise TypeError(
|
||||
f"pre-quantized FP8 GEMM requires fp8 inputs, got {a.dtype}/{b.dtype}"
|
||||
)
|
||||
return _mod().fp8_mm_prequant_fp8(a, b, scale, out_scale)
|
||||
|
||||
|
||||
@fp8_mm_prequant_fp8.register_kernel("cpu")
|
||||
def _fp8_mm_prequant_fp8_cpu(a, b, scale, out_scale):
|
||||
return (a.float() @ b.float().t() * scale * out_scale).to(torch.float8_e4m3fn)
|
||||
|
||||
|
||||
def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w):
|
||||
"""Quantize BF16 inputs to FP8, accumulate in FP32, and return BF16.
|
||||
|
||||
x/w: [..., K] / [N, K] bf16; sx/sw and their inverses control the fused
|
||||
E4M3 conversion; amax_x/amax_w receive the input max-abs values.
|
||||
"""
|
||||
if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16):
|
||||
raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")
|
||||
return _mod().fp8_linear_forward_scaled(
|
||||
x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w
|
||||
@fp8_quantize.register_fake
|
||||
def _fp8_quantize_fake(x, scale, fmt):
|
||||
dtype = torch.float8_e5m2 if fmt else torch.float8_e4m3fn
|
||||
return (
|
||||
torch.empty(x.shape, device=x.device, dtype=dtype),
|
||||
torch.empty(1, device=x.device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
def linear_backward_scaled(g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g):
|
||||
"""dX = g @ W, dW = g^T @ X, dB = sum(g) with per-tensor scales."""
|
||||
@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 _mod().quantize_bf16(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"))
|
||||
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,
|
||||
sa: torch.Tensor,
|
||||
sb: torch.Tensor,
|
||||
out_dtype: int = 0,
|
||||
out_scale: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""FP8 GEMM: ``a @ b^T * (sa * sb)`` with FP32 accumulation.
|
||||
|
||||
``out_dtype``: 0 = BF16 (default), 1 = FP8 E4M3 (requires ``out_scale``,
|
||||
the quantization step for the output — mirrors ``torch._scaled_mm``).
|
||||
"""
|
||||
|
||||
|
||||
@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(0)), device=a.device, dtype=dtype)
|
||||
|
||||
|
||||
@fp8_gemm.register_kernel("cuda")
|
||||
def _fp8_gemm_cuda(a, b, sa, sb, out_dtype=0, out_scale=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 _mod().mm_fp8(a, b, sa, sb, int(out_dtype), out_scale)
|
||||
|
||||
|
||||
@fp8_gemm.register_kernel("cpu")
|
||||
def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None):
|
||||
acc = a.float() @ b.float().t() * sa * sb
|
||||
if out_dtype:
|
||||
os_ = 1.0 if out_scale is None else out_scale
|
||||
return (acc * os_).to(torch.float8_e4m3fn)
|
||||
return acc.to(torch.bfloat16)
|
||||
|
||||
|
||||
def quantize_bf16(x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"):
|
||||
"""BF16 -> 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.
|
||||
"""
|
||||
return fp8_quantize(x, scale, _fmt_int(fmt))
|
||||
|
||||
|
||||
def mm_fp8(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
sa: torch.Tensor,
|
||||
sb: torch.Tensor,
|
||||
out_dtype: str = "bf16",
|
||||
out_scale: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Pre-quantized FP8 GEMM: ``a @ b^T * (sa * sb)``.
|
||||
|
||||
``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).
|
||||
"""
|
||||
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, w, bias, sx, sw):
|
||||
"""BF16 linear forward, quantizing x/w to E4M3 inside the GEMM.
|
||||
|
||||
Returns ``(out, amax_x, amax_w)``. ``bias`` may be ``None``.
|
||||
"""
|
||||
if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16):
|
||||
raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")
|
||||
if bias is None:
|
||||
bias = torch.empty(0, device=x.device, dtype=x.dtype)
|
||||
return _mod().linear_forward_fp8(x, w, bias, sx, sw)
|
||||
|
||||
|
||||
def linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt: str = "e5m2"):
|
||||
"""FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``.
|
||||
|
||||
The gradient (and the transposed w/x operands) are quantized to ``fmt``
|
||||
(default E5M2 — larger dynamic range for gradients) and the two GEMMs run
|
||||
as FP8 tensor-core products sharing a single gradient quantization.
|
||||
"""
|
||||
if not (
|
||||
g.dtype == torch.bfloat16
|
||||
and x.dtype == torch.bfloat16
|
||||
@@ -122,6 +181,4 @@ def linear_backward_scaled(g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, a
|
||||
raise TypeError(
|
||||
f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
|
||||
)
|
||||
return _mod().fp8_linear_backward_scaled(
|
||||
g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g
|
||||
)
|
||||
return _mod().linear_backward_fp8(g, x, w, list(masks), sg, sw, sx, _fmt_int(fmt))
|
||||
|
||||
Reference in New Issue
Block a user