diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index db1606d..591f51f 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -1,161 +1,214 @@ -"""FP8 training: scaling state and aten::linear dispatch. +"""FP8 training: scaling recipes, per-tensor state, and aten::linear dispatch. Layered (see also ``ops/fp8.py`` for the CUDA interface adapter): -1. Kernel interface: ``ops.fp8`` - the only module touching the pybind. -2. Training state (this module): per-tensor scales, amax history, delayed - scaling, and the ``fp8_autocast`` context (TE-style, like - ``torch.autocast``). +1. Kernel interface: ``ops.fp8`` — the only module touching the pybind. +2. Training state (this module): scaling *recipes* (TE-style delayed scaling + or dynamic current-amax scaling), per-tensor scales + amax history, and + the ``fp8_autocast`` context (like ``torch.autocast``). 3. aten::linear integration (this module): registers the CUDA impl and the - M/N alignment guard. + dtype guard. Usage:: from astrai.extension.fp8 import fp8_autocast - with fp8_autocast(enabled=True): + with fp8_autocast(enabled=True, fp8_format="hybrid"): logits = model(input_ids) loss.backward() Importing this module registers the aten::linear CUDA implementation. + +Format defaults follow the ecosystem consensus: E4M3 for the forward pass, +E5M2 for the backward (gradient) pass ("hybrid"); every operand's scale is a +quantization step derived from its amax history by the active recipe. """ from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from typing import Optional import torch from torch.library import Library from astrai.extension.ops.fp8 import ( - linear_backward_scaled, - linear_forward_scaled, + linear_backward_fp8, + linear_forward_fp8, + mm_fp8, + quantize_bf16, ) -E4M3_MAX = 448.0 +# Max representable value per FP8 format (E4M3: 448, E5M2: 57344). +FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0} +E4M3_MAX = FP8_MAX["e4m3"] # legacy alias -# --------------------------------------------------------------------------- -# Layer 2: training state (scales, amax history, delayed scaling, autocast) -# --------------------------------------------------------------------------- +class FP8Format(str, Enum): + """Per-direction FP8 format. HYBRID = E4M3 forward / E5M2 backward.""" + + E4M3 = "e4m3" + E5M2 = "e5m2" + HYBRID = "hybrid" + + def fwd(self) -> str: + return "e4m3" if self is FP8Format.HYBRID else self.value + + def bwd(self) -> str: + return "e5m2" if self is FP8Format.HYBRID else self.value + + +class FP8Recipe: + """Scale-from-amax policy; the scale computation is the injection point. + + ``scale_from_history`` receives the amax tensor for this operand (a ring + window for delayed scaling, the current amax for dynamic scaling) and + returns the quantization step: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``. + """ + + history_len: int = 16 + margin: int = 0 + + def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor: + raise NotImplementedError + + +@dataclass +class DelayedScaling(FP8Recipe): + """TE-style delayed scaling: max over the amax history window. + + The scale is computed from amax measured in *previous* steps (delayed one + step); the window length trades responsiveness against stability. + """ + + history_len: int = 16 + margin: int = 0 + + def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor: + peak = amax.max() + return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12) + + +@dataclass +class DynamicScaling(FP8Recipe): + """Current-amax scaling (torchao DYNAMIC): measure, then quantize. + + No history — the scale is derived from the amax of the tensor being + quantized in the same step, at the cost of an extra reduction pass. + """ + + history_len: int = 1 + margin: int = 0 + + def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor: + peak = amax.max() + return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12) class FP8TensorMeta: - """Scales + amax state for one weight tensor and its paired activations. + """Per-tensor scaling state: amax history rings + derived scales. - - weight: delayed scale from a 16-step amax history window (TE style) - - x/g: delayed one step, reuse the quantize kernel's free atomic amax + One ring per operand (weight / activation / gradient). Scales are derived + from the ring by the recipe; fused kernels record the amax while + quantizing, so the scale used at step N reflects amax from steps < N + (delayed one step). """ __slots__ = ( - "scale", - "scale_inv", - "amax_history", - "idx", - "x_scale", - "x_scale_inv", - "x_history", + "recipe", + "w_hist", + "x_hist", + "g_hist", + "w_idx", "x_idx", - "g_scale", - "g_scale_inv", - "g_history", "g_idx", + "w_scale", + "x_scale", + "g_scale", "w_init", "x_init", "g_init", ) - def __init__(self, device: torch.device, update_interval: int): - self.scale = torch.ones(1, device=device, dtype=torch.float32) - self.scale_inv = torch.ones(1, device=device, dtype=torch.float32) - self.amax_history = torch.ones( - update_interval, device=device, dtype=torch.float32 - ) - self.idx = 0 + def __init__(self, device: torch.device, recipe: FP8Recipe): + self.recipe = recipe + n = recipe.history_len + self.w_hist = torch.ones(n, device=device, dtype=torch.float32) + self.x_hist = torch.ones(n, device=device, dtype=torch.float32) + self.g_hist = torch.ones(n, device=device, dtype=torch.float32) + self.w_idx = self.x_idx = self.g_idx = 0 + self.w_scale = torch.ones(1, device=device, dtype=torch.float32) self.x_scale = torch.ones(1, device=device, dtype=torch.float32) - self.x_scale_inv = torch.ones(1, device=device, dtype=torch.float32) - self.x_history = torch.ones(update_interval, device=device, dtype=torch.float32) - self.x_idx = 0 self.g_scale = torch.ones(1, device=device, dtype=torch.float32) - self.g_scale_inv = torch.ones(1, device=device, dtype=torch.float32) - self.g_history = torch.ones(update_interval, device=device, dtype=torch.float32) - self.g_idx = 0 - self.w_init = False - self.x_init = False - self.g_init = False + self.w_init = self.x_init = self.g_init = False - def init_scale(self, t: torch.Tensor) -> None: - """Immediate scale from the current amax; used on the first call. + # -- ring helpers ------------------------------------------------------- - A scale of 1 would underflow small activations/gradients (e4m3 min - normal is 2^-6); initialize from the actual amax once, then delayed - updates take over. - """ + def _record(self, hist: torch.Tensor, idx: int, amax: torch.Tensor) -> int: + hist[idx] = amax.reshape(()) + return (idx + 1) % hist.numel() + + def _refresh(self, hist: torch.Tensor, scale: torch.Tensor, fmt: str) -> None: + scale.copy_(self.recipe.scale_from_history(hist, fmt)) + + def _seed( + self, hist: torch.Tensor, scale: torch.Tensor, t: torch.Tensor, fmt: str + ) -> None: amax = t.abs().amax().to(torch.float32).clamp_min(1e-12) - self.scale.copy_(amax / E4M3_MAX) - self.scale_inv.copy_(E4M3_MAX / amax) - self.record(amax) + hist.fill_(amax) + scale.copy_(self.recipe.scale_from_history(hist, fmt)) - def push_x_scale(self, amax: torch.Tensor) -> None: - """Window update for the activation scale (delayed, TE style).""" - self.x_history[self.x_idx] = amax.reshape(()) - self.x_idx = (self.x_idx + 1) % self.x_history.numel() - m = self.x_history.max() - self.x_scale.copy_(m / E4M3_MAX) - self.x_scale_inv.copy_(E4M3_MAX / m) + # -- per-operand updates (delayed: record now, refresh for next step) --- - def push_g_scale(self, amax: torch.Tensor) -> None: - """Window update for the gradient scale (delayed, TE style).""" - self.g_history[self.g_idx] = amax.reshape(()) - self.g_idx = (self.g_idx + 1) % self.g_history.numel() - m = self.g_history.max() - self.g_scale.copy_(m / E4M3_MAX) - self.g_scale_inv.copy_(E4M3_MAX / m) + def update_w(self, amax: torch.Tensor, fmt: str) -> None: + self.w_idx = self._record(self.w_hist, self.w_idx, amax) + self._refresh(self.w_hist, self.w_scale, fmt) - def record(self, amax: torch.Tensor) -> None: - """Push the latest amax into the ring buffer (device-side copy, no sync).""" - self.amax_history[self.idx] = amax.reshape(()) - self.idx = (self.idx + 1) % self.amax_history.numel() + def update_x(self, amax: torch.Tensor, fmt: str) -> None: + self.x_idx = self._record(self.x_hist, self.x_idx, amax) + self._refresh(self.x_hist, self.x_scale, fmt) - def refresh(self) -> None: - """Recompute scale from the amax history window (delayed scaling).""" - amax = self.amax_history.max() - if amax > 0: - self.scale.copy_(amax / E4M3_MAX) - self.scale_inv.copy_(E4M3_MAX / amax) + def update_g(self, amax: torch.Tensor, fmt: str) -> None: + self.g_idx = self._record(self.g_hist, self.g_idx, amax) + self._refresh(self.g_hist, self.g_scale, fmt) + + # -- first-use seeding -------------------------------------------------- + + def init_w(self, w: torch.Tensor, fmt: str) -> None: + self._seed(self.w_hist, self.w_scale, w, fmt) + self.w_init = True + + def init_x(self, x: torch.Tensor, fmt: str) -> None: + self._seed(self.x_hist, self.x_scale, x, fmt) + self.x_init = True + + def init_g(self, g: torch.Tensor, fmt: str) -> None: + self._seed(self.g_hist, self.g_scale, g, fmt) + self.g_init = True class FP8State: - """Global fp8 training state, TE-style.""" + """Global fp8 training state: active recipe + per-tensor metas.""" - def __init__(self, update_interval: int = 16): + def __init__(self): self.enabled = False - self.update_interval = update_interval - self.step_count = 0 + self.recipe: FP8Recipe = DelayedScaling() + self.fp8_format: FP8Format = FP8Format.HYBRID self._metas: dict[tuple, FP8TensorMeta] = {} - self._last_device: torch.device | None = None - - def _get_device(self, t: torch.Tensor) -> torch.device: - if self._last_device is None: - self._last_device = t.device - return t.device + self._last_device: Optional[torch.device] = None def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta: key = (w.data_ptr(), w.shape, w.dtype) meta = self._metas.get(key) if meta is None: - meta = FP8TensorMeta(self._get_device(w), self.update_interval) + if self._last_device is None: + self._last_device = w.device + meta = FP8TensorMeta(w.device, self.recipe) self._metas[key] = meta return meta - def step(self) -> None: - """Advance the counter and refresh all weight scales every N steps.""" - self.step_count += 1 - if self.step_count % self.update_interval == 0: - for meta in self._metas.values(): - meta.refresh() - def reset(self) -> None: self.enabled = False - self.step_count = 0 self._metas.clear() self._last_device = None @@ -171,100 +224,114 @@ def fp8_state() -> FP8State: @contextmanager -def fp8_autocast(enabled: bool = True, update_interval: int = 16): +def fp8_autocast( + enabled: bool = True, + update_interval: int = 16, + recipe: Optional[FP8Recipe] = None, + fp8_format: str = "hybrid", + margin: int = 0, +): """Autocast-style context: fp8 linear dispatch on this thread. Usage:: - with fp8_autocast(enabled=True): + with fp8_autocast(enabled=True, fp8_format="hybrid"): logits = model(input_ids) # aten::linear -> fp8 path loss.backward() - The scale-update counter advances once per ``enter`` (one training step), - refreshing weight scales from their amax history every ``update_interval``. + Args: + enabled: toggle fp8 dispatch for aten::linear. + update_interval: legacy alias for the delayed-scaling history window + (used only when ``recipe`` is not given). + recipe: scaling policy; defaults to ``DelayedScaling(update_interval)``. + fp8_format: ``"e4m3"`` / ``"e5m2"`` / ``"hybrid"`` (default) — hybrid + means E4M3 forward, E5M2 backward. + margin: scale headroom (``scale = (amax / FP8_MAX) / 2^margin``) used + with the default delayed recipe. """ state = fp8_state() - prev_enabled = state.enabled - prev_interval = state.update_interval + prev = (state.enabled, state.recipe, state.fp8_format) + if recipe is None: + recipe = DelayedScaling(history_len=update_interval, margin=margin) state.enabled = enabled - state.update_interval = update_interval + state.recipe = recipe + state.fp8_format = FP8Format(fp8_format) try: - if enabled: - state.step() yield finally: - state.enabled = prev_enabled - state.update_interval = prev_interval + state.enabled, state.recipe, state.fp8_format = prev + + +# --------------------------------------------------------------------------- +# Strategy-level forward / backward (called from the aten::linear impl) +# --------------------------------------------------------------------------- + + +def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor: + amax = t.abs().amax().to(torch.float32).clamp_min(1e-12) + return recipe.scale_from_history(amax, fmt) def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): - """TE-style scaled fp8 linear forward (called from the aten::linear impl). + """Scaled fp8 linear forward (called from the aten::linear impl). - x uses the delayed scale of its paired weight meta (amax from the previous - forward of this linear); the quantize kernel emits the current amax for the - next step. No extra abs/max reduce. + Delayed scaling uses the fused BF16->E4M3 GEMM (quantize + amax inside the + kernel); dynamic scaling measures the current amax first and runs the + pre-quantized path. """ if bias is None: bias = torch.empty(0, device=x.device, dtype=x.dtype) state = fp8_state() + fmt = state.fp8_format.fwd() meta = state.get_weight_meta(w) if not meta.w_init: - meta.init_scale(w) - meta.w_init = True + meta.init_w(w, fmt) + if isinstance(state.recipe, DynamicScaling): + x_2d = x.reshape(-1, w.size(1)) + sx = _dynamic_scale(x_2d, state.recipe, fmt) + sw = _dynamic_scale(w, state.recipe, fmt) + x8, _ = quantize_bf16(x_2d, sx, fmt) + w8, _ = quantize_bf16(w, sw, fmt) + out = mm_fp8(x8, w8, sx, sw) + out = out.reshape(*x.shape[:-1], w.size(0)) + if bias.numel(): + out = out + bias + return out if not meta.x_init: - amax = x.abs().amax().to(torch.float32).clamp_min(1e-12) - meta.x_history.fill_(amax) - meta.x_scale.copy_(amax / E4M3_MAX) - meta.x_scale_inv.copy_(E4M3_MAX / amax) - meta.x_init = True - amax_x = torch.empty(1, device=x.device, dtype=torch.float32) - amax_w = torch.empty(1, device=x.device, dtype=torch.float32) - out = linear_forward_scaled( - x, - w, - bias, - meta.x_scale, - meta.scale, - meta.x_scale_inv, - meta.scale_inv, - amax_x, - amax_w, - ) - meta.record(amax_w) - meta.push_x_scale(amax_x) + meta.init_x(x, fmt) + out, amax_x, amax_w = linear_forward_fp8(x, w, bias, meta.x_scale, meta.w_scale) + meta.update_x(amax_x, fmt) + meta.update_w(amax_w, fmt) return out -def fp8_linear_backward(g, x, w, masks): - """TE-style scaled fp8 linear backward (called from aten::linear_backward).""" +def fp8_linear_backward(g: torch.Tensor, x: torch.Tensor, w: torch.Tensor, masks): + """Scaled fp8 linear backward (called from aten::linear_backward). + + The gradient is quantized to the backward format (E5M2 in hybrid mode) + and the dX / dW GEMMs share that single quantization. + """ state = fp8_state() + fmt = state.fp8_format.bwd() meta = state.get_weight_meta(w) - if not meta.g_init: - amax = g.abs().amax().to(torch.float32).clamp_min(1e-12) - meta.g_history.fill_(amax) - meta.g_scale.copy_(amax / E4M3_MAX) - meta.g_scale_inv.copy_(E4M3_MAX / amax) - meta.g_init = True - amax_g = torch.empty(1, device=g.device, dtype=torch.float32) - out = linear_backward_scaled( - g, - x, - w, - masks, - meta.g_scale, - meta.scale, - meta.x_scale, - meta.g_scale_inv, - meta.scale_inv, - meta.x_scale_inv, - amax_g, + if isinstance(state.recipe, DynamicScaling): + sg = _dynamic_scale(g, state.recipe, fmt) + sw = _dynamic_scale(w, state.recipe, fmt) + sx = _dynamic_scale(x, state.recipe, fmt) + else: + if not meta.g_init: + meta.init_g(g, fmt) + sg, sw, sx = meta.g_scale, meta.w_scale, meta.x_scale + grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( + g, x, w, masks, sg, sw, sx, fmt ) - meta.push_g_scale(amax_g) - return out + if not isinstance(state.recipe, DynamicScaling): + meta.update_g(amax_g, fmt) + return grad_x, grad_w, grad_b # --------------------------------------------------------------------------- -# Layer 3: aten::linear integration +# aten::linear integration # --------------------------------------------------------------------------- @@ -279,9 +346,14 @@ def fp8_linear_enabled() -> bool: def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool: - """cuBLASLt fp8 requires M % 16 == 0 and N % 16 == 0 (K is padded).""" - m = x.numel() // x.size(-1) - return m % 16 == 0 and w.size(0) % 16 == 0 + """Shape guard for the fp8 linear path. + + Unlike a strict 16-alignment requirement, the fp8 kernels handle unaligned + M/N via boundary checks (slower but correct) — so no whole-call bf16 + fallback for small decode batches. Only the K-dimension contraction must + match, and the weight must be 2D. + """ + return x.dim() >= 2 and w.dim() == 2 and x.size(-1) == w.size(1) def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): diff --git a/astrai/extension/loader.py b/astrai/extension/loader.py index 7647c21..9dceefe 100644 --- a/astrai/extension/loader.py +++ b/astrai/extension/loader.py @@ -1,24 +1,35 @@ """Dynamic discovery and loading of compiled CUDA kernel modules. -Each kernel is registered in ``csrc/build.py`` and built into a ``.so`` placed -in this package directory. On import we try to load each one; kernels that -failed to build (or are running on a CPU-only machine) are marked unavailable -so the wrapper functions can fall back to ``torch`` SDPA. +Each kernel is built by the CMake build in ``csrc/CMakeLists.txt`` into a +``.so`` placed in ``astrai/extension/lib/`` — the module name equals the +``.so`` name equals the pybind name (e.g. ``attn_decode``, defined via +``TORCH_EXTENSION_NAME``). ``KERNEL_NAMES`` is discovered automatically from +the ``.so`` files present, so adding a kernel to the CMake ``KERNELS`` +registry needs no change here. On import we try to load each one; kernels +that failed to build (or are running on a CPU-only machine) are marked +unavailable so the wrapper functions can fall back to ``torch`` SDPA. """ +import glob import importlib import logging +import os logger = logging.getLogger(__name__) -KERNEL_NAMES = [ - "attn_decode", - "attn_prefill", - "attn_paged_decode", - "attn_paged_prefill", - "rotary_emb", - "fp8_mm", -] +_LIB_DIR = os.path.join(os.path.dirname(__file__), "lib") + + +def _discover_kernel_names() -> list[str]: + """Return the module names of the compiled kernel ``.so`` files in lib/.""" + names: list[str] = [] + for path in glob.glob(os.path.join(_LIB_DIR, "*.so")): + # strip the ".so" suffix, e.g. attn_decode.cpython-312-...so + names.append(os.path.basename(path).split(".", 1)[0]) + return sorted(names) + + +KERNEL_NAMES = _discover_kernel_names() _available: dict[str, bool] = {} _modules: dict[str, object] = {} diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index ccb2288..3ec0134 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -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)) diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 3aa63bf..5346253 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -48,10 +48,33 @@ set(TORCH_LIBS set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}") -set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb fp8_mm) +# Kernel registry — parallel lists of module names (.so / pybind names, +# globally unique across families) and their per-family source paths under +# kernels/. `loader.py` auto-discovers the .so files in astrai/extension/lib/, +# so this CMake registry is the single place to register a new kernel. +set(KERNEL_NAMES + attn_decode + attn_prefill + attn_paged_decode + attn_paged_prefill + rotary_emb + fp8_mm +) +set(KERNEL_SRCS + attention/decode.cu + attention/prefill.cu + attention/paged_decode.cu + attention/paged_prefill.cu + rotary/rotary_emb.cu + fp8/mm.cu +) -foreach(name ${KERNELS}) - add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu") +list(LENGTH KERNEL_NAMES _kernel_count) +math(EXPR _kernel_last "${_kernel_count} - 1") +foreach(i RANGE ${_kernel_last}) + list(GET KERNEL_NAMES ${i} name) + list(GET KERNEL_SRCS ${i} src) + add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${src}") target_compile_definitions(${name} PRIVATE TORCH_EXTENSION_NAME=${name}) diff --git a/csrc/__init__.py b/csrc/__init__.py index bcae4f2..face7cf 100644 --- a/csrc/__init__.py +++ b/csrc/__init__.py @@ -1,2 +1,2 @@ # Source directory for CUDA kernels — build-time only. -# Compiled .so files live in astrAI/_ext/. +# Compiled .so files live in astrai/extension/lib/ (see csrc/CMakeLists.txt). diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attention/common.h similarity index 96% rename from csrc/kernels/attn_common.h rename to csrc/kernels/attention/common.h index 92f66fe..fe773fe 100644 --- a/csrc/kernels/attn_common.h +++ b/csrc/kernels/attention/common.h @@ -13,7 +13,7 @@ enum TensorLayout : int { // - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v). // - Paged (SGLang-style): flat pool [size, kv_head, head_dim] + req_to_token. // Each kernel selects the addressing via a KVSource policy (see -// attn_layout_policies.cuh); a given call only touches the fields of one mode, so +// layout_policies.cuh); a given call only touches the fields of one mode, so // this is a POD shared by both paths rather than two parallel structs that // drift out of sync. template diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attention/decode.cu similarity index 97% rename from csrc/kernels/attn_decode.cu rename to csrc/kernels/attention/decode.cu index 598af72..6d48f26 100644 --- a/csrc/kernels/attn_decode.cu +++ b/csrc/kernels/attention/decode.cu @@ -1,5 +1,5 @@ -#include "attn_dispatchers.cuh" -#include "attn_entry_utils.cuh" +#include "dispatchers.cuh" +#include "entry_utils.cuh" torch::Tensor attn_decode( torch::Tensor q, diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attention/decode_split_kv.cuh similarity index 98% rename from csrc/kernels/attn_decode_split_kv.cuh rename to csrc/kernels/attention/decode_split_kv.cuh index 337f5f7..564fe30 100644 --- a/csrc/kernels/attn_decode_split_kv.cuh +++ b/csrc/kernels/attention/decode_split_kv.cuh @@ -1,9 +1,9 @@ #pragma once #include #include -#include "attn_common.h" -#include "attn_layout_policies.cuh" -#include "attn_warp_utils.cuh" +#include "common.h" +#include "layout_policies.cuh" +#include "warp_utils.cuh" constexpr int DC_CHUNK = 64; // Scalar split-KV decode (fallback for sm < 80, no tensor cores), unified diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attention/decode_split_kv_mma.cuh similarity index 98% rename from csrc/kernels/attn_decode_split_kv_mma.cuh rename to csrc/kernels/attention/decode_split_kv_mma.cuh index 6a66da6..a454e7c 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attention/decode_split_kv_mma.cuh @@ -1,10 +1,10 @@ #pragma once #include #include -#include "attn_common.h" -#include "attn_layout_policies.cuh" -#include "attn_mma_utils.cuh" -#include "attn_warp_utils.cuh" +#include "common.h" +#include "layout_policies.cuh" +#include "mma_utils.cuh" +#include "warp_utils.cuh" // Split-K (FlashDecoding) tensor-core decode via GQA head-packing, unified // across contiguous and paged (SGLang flat-pool) K/V via the KV template diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attention/dispatchers.cuh similarity index 97% rename from csrc/kernels/attn_dispatchers.cuh rename to csrc/kernels/attention/dispatchers.cuh index 84ad175..1fc412f 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attention/dispatchers.cuh @@ -3,20 +3,20 @@ // No torch dependency; pure CUDA. // // The paged and contiguous kernels are unified by the KVSource policy -// (ContigKV / PagedKV from attn_layout_policies.cuh), so each launcher struct +// (ContigKV / PagedKV from layout_policies.cuh), so each launcher struct // below is templated on KV and the paged dispatch is just the same launcher // instantiated with PagedKV. Only the grid/split math differs, and that is // covered by KV::host_q_len / KV::host_kv_len. #include #include -#include "attn_warp_utils.cuh" -#include "attn_layout_policies.cuh" -#include "attn_prefill_split_q.cuh" -#include "attn_decode_split_kv.cuh" +#include "warp_utils.cuh" +#include "layout_policies.cuh" +#include "prefill_split_q.cuh" +#include "decode_split_kv.cuh" #ifndef ASTRAI_NO_MMA -#include "attn_prefill_split_q_mma.cuh" -#include "attn_decode_split_kv_mma.cuh" +#include "prefill_split_q_mma.cuh" +#include "decode_split_kv_mma.cuh" #endif // Split-KV: compute number of splits to fill all SMs for small-batch decode. diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attention/entry_utils.cuh similarity index 99% rename from csrc/kernels/attn_entry_utils.cuh rename to csrc/kernels/attention/entry_utils.cuh index 309db1b..7f55705 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attention/entry_utils.cuh @@ -2,8 +2,8 @@ #include #include #include -#include "attn_common.h" -#include "attn_warp_utils.cuh" +#include "common.h" +#include "warp_utils.cuh" using bf16 = __nv_bfloat16; diff --git a/csrc/kernels/attn_layout_policies.cuh b/csrc/kernels/attention/layout_policies.cuh similarity index 99% rename from csrc/kernels/attn_layout_policies.cuh rename to csrc/kernels/attention/layout_policies.cuh index 84ebaa3..c8ec715 100644 --- a/csrc/kernels/attn_layout_policies.cuh +++ b/csrc/kernels/attention/layout_policies.cuh @@ -1,6 +1,6 @@ #pragma once #include -#include "attn_common.h" +#include "common.h" // ============================================================================ // Attention layout policies keep Q scheduling independent from K/V storage. diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attention/mma_utils.cuh similarity index 84% rename from csrc/kernels/attn_mma_utils.cuh rename to csrc/kernels/attention/mma_utils.cuh index d5ac83e..06aaebe 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attention/mma_utils.cuh @@ -3,6 +3,8 @@ #include #include +#include "../common/mma.cuh" + // Predicated cp.async (4-operand form) requires CUDA 11.2+. // bf16 mma.sync requires sm_80+ (guarded at build time by ASTRAI_NO_MMA). #if CUDART_VERSION < 11020 @@ -24,10 +26,10 @@ struct KernelTraits { static constexpr int BR = 16; // Q rows per warp (mma M=16) - // Derived: mma.sync.m16n8k16 tile counts - static constexpr int KD = HEAD_DIM / 16; // Q/K k-slides + // Derived: mma tile counts from the shared mma_shape (m16n8k16 for bf16) + static constexpr int KD = HEAD_DIM / astrai::mma_shape::k; // Q/K k-slides static constexpr int NC8 = BC / 8; // S n-tiles (N=8) - static constexpr int KT2 = BC / 16; // P k-tiles (K=16) + static constexpr int KT2 = BC / astrai::mma_shape::k; // P k-tiles (K=16) static constexpr int DN8 = HEAD_DIM / 8; // O n-tiles (N=8) static constexpr int LD = HEAD_DIM; // smem leading dim @@ -43,16 +45,7 @@ struct KernelTraits { // ---- PTX wrappers ---- using bf16 = __nv_bfloat16; - -__device__ __forceinline__ void mma16816(float* d, const unsigned* a, - const unsigned* b, const float* c) { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};" - : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); -} +// bf16 mma.sync lives in the shared astrai::mma_sync template (common/mma.cuh). // read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high) __device__ __forceinline__ unsigned ld2(const bf16* p) { @@ -73,26 +66,9 @@ __device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) { return *reinterpret_cast(&v); } -// ldmatrix: cooperatively load mma fragments from smem (one instruction per -// 16x16 / 16x8 tile) with the exact register layout mma expects. -__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) { - unsigned a = __cvta_generic_to_shared(p); - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" - : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) - : "r"(a)); -} -__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) { - unsigned a = __cvta_generic_to_shared(p); - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];" - : "=r"(r[0]), "=r"(r[1]) - : "r"(a)); -} -__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) { - unsigned a = __cvta_generic_to_shared(p); - asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];" - : "=r"(r[0]), "=r"(r[1]) - : "r"(a)); -} +// ldmatrix lives in the shared template (common/mma.cuh): +// `astrai::ldmatrix_x2` / `` load the K/V +// fragments with the exact register layout mma expects. // XOR swizzle for shared-memory column at 8-bf16 chunk granularity. __device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) { @@ -180,9 +156,9 @@ __device__ inline void mma_compute_scores( #pragma unroll for (int kt = 0; kt < Traits::KD; kt++) { unsigned b[2]; - ldmatrix_x2(b, &sK[krow_l * Traits::LD + astrai::ldmatrix_x2(b, &sK[krow_l * Traits::LD + swiz_col(kt * 16 + kcol_h, krow_l, Traits::SWIZ_MASK)]); - mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); + astrai::mma_sync(Sacc[n8], Qa[kt], b, Sacc[n8]); } } } @@ -290,9 +266,9 @@ __device__ inline void mma_pv_accumulate( #pragma unroll for (int dn8 = 0; dn8 < Traits::DN8; dn8++) { unsigned b[2]; - ldmatrix_x2_trans(b, &sV[vrow_l * Traits::LD + astrai::ldmatrix_x2(b, &sV[vrow_l * Traits::LD + swiz_col(dn8 * 8, vrow_l, Traits::SWIZ_MASK)]); - mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); + astrai::mma_sync(Oacc[dn8], Pa, b, Oacc[dn8]); } } } diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attention/paged_decode.cu similarity index 98% rename from csrc/kernels/attn_paged_decode.cu rename to csrc/kernels/attention/paged_decode.cu index 4cca860..bff3662 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/csrc/kernels/attention/paged_decode.cu @@ -1,5 +1,5 @@ -#include "attn_dispatchers.cuh" -#include "attn_entry_utils.cuh" +#include "dispatchers.cuh" +#include "entry_utils.cuh" torch::Tensor attn_paged_decode( torch::Tensor q, diff --git a/csrc/kernels/attn_paged_prefill.cu b/csrc/kernels/attention/paged_prefill.cu similarity index 96% rename from csrc/kernels/attn_paged_prefill.cu rename to csrc/kernels/attention/paged_prefill.cu index b30f89b..157401f 100644 --- a/csrc/kernels/attn_paged_prefill.cu +++ b/csrc/kernels/attention/paged_prefill.cu @@ -1,5 +1,5 @@ -#include "attn_dispatchers.cuh" -#include "attn_entry_utils.cuh" +#include "dispatchers.cuh" +#include "entry_utils.cuh" torch::Tensor attn_paged_prefill( torch::Tensor q, diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attention/prefill.cu similarity index 94% rename from csrc/kernels/attn_prefill.cu rename to csrc/kernels/attention/prefill.cu index 3588e74..47a8e8a 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attention/prefill.cu @@ -1,5 +1,5 @@ -#include "attn_dispatchers.cuh" -#include "attn_entry_utils.cuh" +#include "dispatchers.cuh" +#include "entry_utils.cuh" torch::Tensor attn_prefill( torch::Tensor q, diff --git a/csrc/kernels/attn_prefill_split_q.cuh b/csrc/kernels/attention/prefill_split_q.cuh similarity index 98% rename from csrc/kernels/attn_prefill_split_q.cuh rename to csrc/kernels/attention/prefill_split_q.cuh index 3157512..acdc2b1 100644 --- a/csrc/kernels/attn_prefill_split_q.cuh +++ b/csrc/kernels/attention/prefill_split_q.cuh @@ -1,8 +1,8 @@ #pragma once #include #include -#include "attn_common.h" -#include "attn_layout_policies.cuh" +#include "common.h" +#include "layout_policies.cuh" using bf16 = __nv_bfloat16; diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attention/prefill_split_q_mma.cuh similarity index 98% rename from csrc/kernels/attn_prefill_split_q_mma.cuh rename to csrc/kernels/attention/prefill_split_q_mma.cuh index c3c5122..2dc616b 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attention/prefill_split_q_mma.cuh @@ -1,9 +1,9 @@ #pragma once #include #include -#include "attn_common.h" -#include "attn_layout_policies.cuh" -#include "attn_mma_utils.cuh" +#include "common.h" +#include "layout_policies.cuh" +#include "mma_utils.cuh" // Tensor-core prefill flash attention (raw mma.sync PTX), unified across // contiguous and paged (SGLang flat-pool) K/V via the KV template parameter. diff --git a/csrc/kernels/attn_warp_utils.cuh b/csrc/kernels/attention/warp_utils.cuh similarity index 100% rename from csrc/kernels/attn_warp_utils.cuh rename to csrc/kernels/attention/warp_utils.cuh diff --git a/csrc/kernels/common/device.cuh b/csrc/kernels/common/device.cuh new file mode 100644 index 0000000..7f32dca --- /dev/null +++ b/csrc/kernels/common/device.cuh @@ -0,0 +1,23 @@ +// Pure-CUDA device helpers shared across kernel families (no torch). +// +// Family-local headers under kernels// own their POD params and +// strategy traits; anything cross-cutting (compute-capability checks, device +// constants) lives here. + +#pragma once + +namespace astrai { + +// Compute-capability comparison: is the device at least (major, minor)? +inline bool sm_at_least(int device_major, int device_minor, int major, + int minor) { + return device_major > major || + (device_major == major && device_minor >= minor); +} + +// FP8 tensor-core MMA (`mma.sync.aligned.m16n8k32` with fp8 inputs) exists on +// Ada (sm_89) and Hopper (sm_90+); sm_80 has no fp8 instructions. +inline constexpr int kMinSmForFp8Major = 8; +inline constexpr int kMinSmForFp8Minor = 9; + +} // namespace astrai diff --git a/csrc/kernels/common/mma.cuh b/csrc/kernels/common/mma.cuh new file mode 100644 index 0000000..afbcfd8 --- /dev/null +++ b/csrc/kernels/common/mma.cuh @@ -0,0 +1,143 @@ +// Shared mma.sync wrappers — pure CUDA, no torch. +// +// One template for every tensor-core MMA used by the kernel families. The +// instruction shape follows from the input element type: +// __nv_bfloat16 -> mma.sync.aligned.m16n8k16 (sm_80+), A = 4x b32, B = 2x b32 +// __nv_fp8_e4m3/e5m2 -> mma.sync.aligned.m16n8k32 (sm_89+), A = 4x b32, B = 2x b32 +// All variants accumulate into fp32: d = a*b + c, with the PTX mnemonic and +// the K dimension differing per type. `d` may alias `c` (in-place accumulate, +// as the FP8 GEMM does). + +#pragma once + +#include +#include +#include +#include + +namespace astrai { + +// Compute capability of the current compilation pass: 0 in the host pass, +// the numeric CC (e.g. 890) in device passes where __CUDA_ARCH__ is defined. +// Defined() cannot appear in expressions, so this macro lets mma_sync use +// the arch in a static_assert instead of per-branch #if guards. +#ifndef __CUDA_ARCH__ +#define ASTRAI_DEVICE_ARCH 0 +#else +#define ASTRAI_DEVICE_ARCH __CUDA_ARCH__ +#endif + +// Compile-time shape of the MMA instruction for an input element type. +// `min_arch` is the numeric compute capability the instruction requires — +// the single place that encodes the hardware floor for each type. +template +struct mma_shape { + static constexpr int k = 16; // m16n8k16 + static constexpr int a_regs = 4; // A fragment: 4x b32 + static constexpr int b_regs = 2; // B fragment: 2x b32 + static constexpr int min_arch = 800; // bf16 mma.sync, sm_80+ +}; + +template <> +struct mma_shape<__nv_fp8_e4m3> { + static constexpr int k = 32; // m16n8k32 + static constexpr int a_regs = 4; + static constexpr int b_regs = 2; + static constexpr int min_arch = 890; // fp8 mma.sync, sm_89+ (Ada/Hopper) +}; + +template <> +struct mma_shape<__nv_fp8_e5m2> { + static constexpr int k = 32; + static constexpr int a_regs = 4; + static constexpr int b_regs = 2; + static constexpr int min_arch = 890; +}; + +// d[4] = a[4] x b[2] + c[4], row-major A, col-major B, fp32 accumulator. +// The PTX mnemonic is selected from InT. Building for a compute capability +// below `mma_shape::min_arch` is a **compile error** — the instruction +// does not exist there, and a silent no-op would produce wrong results. +template +__device__ __forceinline__ void mma_sync(float d[4], const unsigned a[4], + const unsigned b[2], + const float c[4]) { + static_assert(ASTRAI_DEVICE_ARCH == 0 || + ASTRAI_DEVICE_ARCH >= mma_shape::min_arch, + "mma_sync: this MMA shape requires a newer compute " + "capability than the build target"); + if constexpr (std::is_same_v) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same_v) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } +} + +#undef ASTRAI_DEVICE_ARCH + +// --------------------------------------------------------------------------- +// ldmatrix — cooperatively load 8x8 b16 matrices from smem into registers. +// +// The instruction is identical for every 16-bit-storage element type: bf16 +// maps 1:1 onto b16 slots; fp8 is stored packed two-per-slot (see +// fp8/gemm.cuh), so one b16 slot holds two fp8 values. `T` is the element +// type and only serves as a semantic tag. +// +// x2 (single address): matrix0 = p (8 rows), matrix1 = p + 8*16 bytes +// x4: four matrices at p, +128, +256, +384 bytes +// Trans: transpose variant (V fragments of attention) +// +// ldmatrix takes a *single* smem address per thread, but the addresses of +// the 32 lanes are *not* all the same: lane i supplies the start address of +// matrix-row i (modulo 8) for matrix (i/8) — lanes 0-7 feed matrix 0's rows, +// lanes 8-15 matrix 1's rows (x2/x4), lanes 16-23 / 24-31 matrix 2 / 3's rows +// (x4 only; their addresses are ignored by x2). Each matrix is 8 rows x 16 +// bytes, and consecutive matrices of one instruction are contiguous at +// 128-byte strides. fp8 fragment layouts in fp8/gemm.cuh are arranged around +// this constraint. +// --------------------------------------------------------------------------- + +template +__device__ __forceinline__ void ldmatrix_x2(unsigned r[2], const T* p) { + const unsigned a = __cvta_generic_to_shared(p); + if constexpr (Trans) { + asm volatile( + "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];" + : "=r"(r[0]), "=r"(r[1]) + : "r"(a)); + } else { + asm volatile( + "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];" + : "=r"(r[0]), "=r"(r[1]) + : "r"(a)); + } +} + +// Four matrices at p, p+128, p+256, p+384 bytes (16-byte row stride). +template +__device__ __forceinline__ void ldmatrix_x4(unsigned r[4], const T* p) { + const unsigned a = __cvta_generic_to_shared(p); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} + +} // namespace astrai diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h new file mode 100644 index 0000000..a6d7214 --- /dev/null +++ b/csrc/kernels/fp8/common.h @@ -0,0 +1,62 @@ +#pragma once + +#include +#include +#include +#include + +// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or +// E5M2 (gradient / large dynamic range, max 57344). +enum class FP8Format : int { + E4M3 = 0, + E5M2 = 1, +}; + +// Compile-time tile configuration, mirroring KernelTraits in the attention kernels. `Fmt` selects the FP8 conversion +// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile and +// the cp.async pipeline depth. +template +struct Fp8GemmTraits { + static constexpr FP8Format kFormat = Fmt; + static constexpr int kBlockM = BlockM; + static constexpr int kBlockN = BlockN; + static constexpr int kK = K; + static constexpr int kStages = Stages; + static constexpr bool kIsE5M2 = (Fmt == FP8Format::E5M2); + static constexpr __nv_fp8_interpretation_t kNvFormat = + kIsE5M2 ? __NV_E5M2 : __NV_E4M3; + static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f; + + // Saturated float -> FP8 conversion for this format. + __device__ __forceinline__ static unsigned char cvt(float f) { + return static_cast( + __nv_cvt_float_to_fp8(f, __NV_SATFINITE, kNvFormat)); + } +}; + +// Unified GEMM parameter POD, mirroring AttentionParams: one struct flows +// through quantize / fused / pre-quantized kernels. Each kernel touches only +// the fields it needs; buffers are raw pointers packed by the torch binding. +struct FP8Params { + // Inputs: a/b are BF16 for the fused (quantize-in-GEMM) path, FP8 for + // the pre-quantized path. Scales are quantization steps (device scalars). + const void* __restrict__ a_ptr; + const void* __restrict__ b_ptr; + const float* __restrict__ scale_a; + const float* __restrict__ scale_b; + + // Output: BF16 or FP8 (E4M3). out_scale is the output quantization step + // (FP8 output only). + void* __restrict__ out_ptr; + const float* __restrict__ out_scale; + + // Fused forward extras: bias (may be null) and amax slots (may be null). + const __nv_bfloat16* __restrict__ bias; + float* __restrict__ amax_a; + float* __restrict__ amax_b; + + // Shapes. total is only used by the elementwise quantize kernel. + int64_t m, n, k; + int64_t total; +}; diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh new file mode 100644 index 0000000..6398f33 --- /dev/null +++ b/csrc/kernels/fp8/gemm.cuh @@ -0,0 +1,668 @@ +#pragma once +// 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. + +#include +#include +#include +#include + +#include "common.h" +#include "../common/mma.cuh" + +namespace fp8 { + +// m16n8k32 (see astrai::mma_shape::k in common/mma.cuh) +constexpr int kMmaK = 32; +constexpr int kWarps = 8; // 128x64 CTA = 8 warps + +// Map the FP8Format enum to the CUDA fp8 element type consumed by mma_sync. +template +struct fp8_input { + using type = __nv_fp8_e4m3; +}; +template <> +struct fp8_input { + using type = __nv_fp8_e5m2; +}; + +// --------------------------------------------------------------------------- +// Shared device helpers +// --------------------------------------------------------------------------- + +__device__ __forceinline__ unsigned pack_fp8x4_vector( + float x0, float x1, float x2, float x3, + __nv_fp8_interpretation_t fmt = __NV_E4M3) { + const auto low = __nv_cvt_float2_to_fp8x2(make_float2(x0, x1), + __NV_SATFINITE, fmt); + const auto high = __nv_cvt_float2_to_fp8x2(make_float2(x2, x3), + __NV_SATFINITE, fmt); + return static_cast(low) | (static_cast(high) << 16); +} + +// FP8 MMA lives in the shared astrai::mma_sync template (common/mma.cuh); +// instantiate it with fp8_input::type. Accumulates in-place: callers +// pass the same accumulator array as both `d` and `c`. + +__device__ __forceinline__ void atomic_max_float(float* destination, + float value) { + if (destination) + atomicMax(reinterpret_cast(destination), + __float_as_uint(value)); +} + +__device__ __forceinline__ float warp_reduce_max(float value) { +#pragma unroll + for (int offset = 16; offset; offset >>= 1) { + value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset)); + } + return value; +} + +// Block-wide max reduction of a per-warp tracked value, then an atomic +// update of the global amax slot when `track` is set. +template +__device__ __forceinline__ void block_reduce_amax(float& local, float* slots, + int warp, int lane, + bool track, float* global) { + local = warp_reduce_max(local); + if (lane == 0) slots[warp] = local; + __syncthreads(); + if (warp == 0) { + float value = lane < NWarps ? slots[lane] : 0.0f; + value = warp_reduce_max(value); + if (lane == 0 && track && global) atomic_max_float(global, value); + } +} + +// One thread moves eight BF16 values (16 bytes) via cp.async; the uint4 +// shape keeps source and destination naturally 128-bit aligned. +__device__ __forceinline__ void cp_async_bf16_8( + __nv_bfloat16* destination, const __nv_bfloat16* source, bool valid) { + const unsigned shared_address = __cvta_generic_to_shared(destination); + const uint4* source_vec = reinterpret_cast(source); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" + :: "r"(shared_address), "l"(source_vec), + "r"(valid ? 16 : 0)); +} + +// One thread moves sixteen FP8 values (16 bytes) via cp.async. +template +__device__ __forceinline__ void cp_async_16b(T* destination, + const T* source, bool valid) { + const unsigned shared_address = __cvta_generic_to_shared(destination); + const uint4* source_vec = reinterpret_cast(source); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" + :: "r"(shared_address), "l"(source_vec), + "r"(valid ? 16 : 0)); +} + +// Convert four BF16 values to one 4xFP8 pack, tracking the raw (pre-scale) +// amax — scaling first would saturate amax at the FP8 max and collapse the +// scale. Format comes from Traits. +template +__device__ __forceinline__ unsigned load_fp8x4_from_bf16( + const __nv_bfloat16* source, float scale_inv, float& amax, + bool track_amax = true) { + float x0 = __bfloat162float(source[0]); + float x1 = __bfloat162float(source[1]); + float x2 = __bfloat162float(source[2]); + float x3 = __bfloat162float(source[3]); + if constexpr (TrackAmax) { + if (track_amax) { + amax = fmaxf(amax, fmaxf(fabsf(x0), fmaxf(fabsf(x1), + fmaxf(fabsf(x2), fabsf(x3))))); + } + } + return pack_fp8x4_vector(x0 * scale_inv, x1 * scale_inv, + x2 * scale_inv, x3 * scale_inv, + Traits::kNvFormat); +} + +// --------------------------------------------------------------------------- +// Quantize kernel: BF16 -> FP8 (E4M3 or E5M2), fused amax over raw values. +// --------------------------------------------------------------------------- + +template +__global__ void fp8_quantize_kernel(FP8Params p) { + const float inv = 1.0f / *p.scale_a; + const auto* x = reinterpret_cast(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; + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < p.total; + i += stride) { + const float f = __bfloat162float(x[i]); + local_amax = fmaxf(local_amax, fabsf(f)); + const float q = f * inv; + if constexpr (Fmt == FP8Format::E5M2) { + reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] = __nv_fp8_e5m2(q); + } else { + reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] = __nv_fp8_e4m3(q); + } + } + 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); + } + } +} + +// --------------------------------------------------------------------------- +// Fused kernel: BF16 A/B -> inline E4M3 quantize -> ldmatrix fragments -> +// MMA -> BF16 out. 128x64 CTA / 64x16 warp tile / cp.async pipeline. +// The quantized FP8 tiles live in a separate smem region laid out around +// ldmatrix's single-address, 128-byte-strided matrices (16-byte rows): +// A8: [M/16 block][4 sub-blocks of 8 rows x 16 fp8][...] where sub-block +// order is (h0,m0-7), (h0,m8-15), (h1,m0-7), (h1,m8-15) — one +// ldmatrix.x4 emits the whole m16n8k32 A fragment (regs 0..3 match). +// B8: [N/8 block][2 sub-blocks of 8 rows x 16 fp8][...] with h0 then h1 — +// one ldmatrix.x2 emits the m16n8k32 B fragment (regs 0,1). +// --------------------------------------------------------------------------- + +template +__global__ void fp8_fused_gemm_kernel(FP8Params p) { + using T8 = __nv_fp8_e4m3; // fused forward always quantizes to E4M3 + constexpr int kBlockM = Traits::kBlockM; + constexpr int kBlockN = Traits::kBlockN; + constexpr int kK = Traits::kK; + constexpr int kStages = Traits::kStages; + constexpr int kWarpM = 64; // warp tile rows (BlockM / 2) + constexpr int kWarpN = 16; // warp tile cols (BlockN / 4) + constexpr int a_stride = kBlockM * kK; // bf16 elements per A stage + constexpr int b_stride = kBlockN * kK; // bf16 elements per B stage + // A8 block layout: (M/16) blocks x 4 sub-blocks x 128 B = BlockM*32 B. + // B8 block layout: (N/8) blocks x 2 sub-blocks x 128 B = BlockN*32 B. + constexpr int a8_bytes = kBlockM * 32; + constexpr int b8_bytes = kBlockN * 32; + // smem layout: [A bf16 stages][B bf16 stages][A8 fp8 tiles][B8 fp8 tiles] + constexpr int bf16_bytes = kStages * (a_stride + b_stride) * 2; + extern __shared__ char smem[]; + auto* a_bf16 = reinterpret_cast<__nv_bfloat16*>(smem); + auto* b_bf16 = reinterpret_cast<__nv_bfloat16*>(smem + kStages * a_stride * 2); + auto* a8 = reinterpret_cast(smem + bf16_bytes); + auto* b8 = reinterpret_cast(smem + bf16_bytes + a8_bytes); + __shared__ float warp_amax_a[kWarps]; + __shared__ float warp_amax_b[kWarps]; + + const auto* a = reinterpret_cast(p.a_ptr); + const auto* b = reinterpret_cast(p.b_ptr); + auto* out = reinterpret_cast<__nv_bfloat16*>(p.out_ptr); + const auto* bias = p.bias; + const float* scale_a = p.scale_a; + const float* scale_b = p.scale_b; + float* amax_a = p.amax_a; + float* amax_b = p.amax_b; + const int64_t m = p.m, n = p.n, k = p.k; + + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int group = lane >> 2; + const int thread_in_group = lane & 3; + constexpr int warps_n = kBlockN / 16; + const int warp_m = warp / warps_n; + const int warp_n = warp % warps_n; + const int64_t row_base = blockIdx.y * kBlockM + warp_m * kWarpM + group; + const int64_t output_col = + blockIdx.x * kBlockN + warp_n * 16 + thread_in_group * 2; + const float sa = *scale_a; + const float sb = *scale_b; + const float inv_a = 1.0f / sa; + const float inv_b = 1.0f / sb; + float local_amax_a = 0.0f; + float local_amax_b = 0.0f; + float acc[4 * 4 * 2] = {}; + + const bool track_amax_a = TrackAmax && blockIdx.x == 0; + const bool track_amax_b = TrackAmax && blockIdx.y == 0; + // Each thread issues 8 A chunks and 4 B chunks of 8 BF16 (16B) per stage. + auto load_tile = [&](int stage, int64_t k_base) { + const int r0 = tid >> 2; + const int c0 = (tid & 3) * 8; +#pragma unroll + for (int j = 0; j < kK / 32; ++j) { + const int col = c0 + 32 * j; + const bool full_chunk = k_base + col + 7 < k; + const int64_t a_row = blockIdx.y * kBlockM + r0; + const int64_t b_row = blockIdx.x * kBlockN + r0; + auto* a_dst = &a_bf16[stage * a_stride + r0 * kK + col]; + auto* b_dst = &b_bf16[stage * b_stride + r0 * kK + col]; + const auto* a_ptr = a + a_row * k + k_base + col; + const auto* b_ptr = b + b_row * k + k_base + col; + const bool full_a = a_row < m && full_chunk; + const bool full_b = b_row < n && full_chunk; + const bool aligned_a = + (reinterpret_cast(a_ptr) & 15) == 0; + const bool aligned_b = + (reinterpret_cast(b_ptr) & 15) == 0; + if (full_a && aligned_a) { + cp_async_bf16_8(a_dst, a_ptr, true); + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) { + a_dst[i] = a_row < m && k_base + col + i < k + ? a_ptr[i] + : __float2bfloat16(0.0f); + } + } + if (full_b && aligned_b) { + cp_async_bf16_8(b_dst, b_ptr, true); + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) { + b_dst[i] = b_row < n && k_base + col + i < k + ? b_ptr[i] + : __float2bfloat16(0.0f); + } + } + if (r0 + kWarpM < kBlockM) { + const int64_t a_row_hi = blockIdx.y * kBlockM + r0 + kWarpM; + auto* a_dst_hi = + &a_bf16[stage * a_stride + (r0 + kWarpM) * kK + col]; + const auto* a_ptr_hi = a + a_row_hi * k + k_base + col; + const bool full_a_hi = a_row_hi < m && full_chunk; + const bool aligned_a_hi = + (reinterpret_cast(a_ptr_hi) & 15) == 0; + if (full_a_hi && aligned_a_hi) { + cp_async_bf16_8(a_dst_hi, a_ptr_hi, true); + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) { + a_dst_hi[i] = a_row_hi < m && k_base + col + i < k + ? a_ptr_hi[i] + : __float2bfloat16(0.0f); + } + } + } + } + }; + + // Quantize the BF16 staging area into the ldmatrix-friendly FP8 tiles. + // A8 sub-block for global row `row` and K half `h`: + // (row>>4)*512 + ((h<<1)|((row>>3)&1))*128 + (row&7)*16 + // B8 sub-block: (row>>3)*256 + h*128 + (row&7)*16. + // Each thread emits one 4-FP8 pack at a time (256 threads, kK/4 = 8 packs + // per row). + auto quantize_tile = [&](int stage) { + constexpr int kA_packs = kBlockM * kK / 4; + constexpr int kB_packs = kBlockN * kK / 4; +#pragma unroll + for (int i = tid; i < kA_packs; i += 256) { + const int row = i >> 3; // 8 packs per row + const int k4 = (i & 7) * 4; + const int half = k4 >> 4; // 0: k 0-15, 1: k 16-31 + const int k16 = k4 & 15; + const int a8_idx = + (row >> 4) * 512 + (((half << 1) | ((row >> 3) & 1)) * 128) + + (row & 7) * 16 + k16; + auto* src = &a_bf16[stage * a_stride + row * kK + k4]; + auto* dst = reinterpret_cast(&a8[a8_idx]); + *dst = load_fp8x4_from_bf16( + src, inv_a, local_amax_a, track_amax_a); + } +#pragma unroll + for (int i = tid; i < kB_packs; i += 256) { + const int row = i >> 3; + const int k4 = (i & 7) * 4; + const int half = k4 >> 4; + const int k16 = k4 & 15; + const int b8_idx = + (row >> 3) * 256 + half * 128 + (row & 7) * 16 + k16; + auto* src = &b_bf16[stage * b_stride + row * kK + k4]; + auto* dst = reinterpret_cast(&b8[b8_idx]); + *dst = load_fp8x4_from_bf16( + src, inv_b, local_amax_b, track_amax_b); + } + }; + + const int64_t tile_count = (k + kK - 1) / kK; + load_tile(0, 0); + asm volatile("cp.async.commit_group;"); + if (tile_count > 1) { + load_tile(1, kK); + asm volatile("cp.async.commit_group;"); + } + if (tile_count > 2) { + load_tile(2, 2 * kK); + asm volatile("cp.async.commit_group;"); + } + for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) { + const int stage = static_cast(tile_index % kStages); + // 3-stage pipeline: at most 2 groups in flight; the tail of the K + // loop waits for everything. + const int64_t remaining = tile_count - tile_index - 1; + if (remaining >= 2) { + asm volatile("cp.async.wait_group 2;"); + } else if (remaining == 1) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + // wait_group only waits for this thread's async copies. All threads + // must finish loading before the tile is read by the CTA. + __syncthreads(); + quantize_tile(stage); + __syncthreads(); + + // kK == kMmaK, so one m16n8k32 MMA segment per K stage; fragments + // come from the fp8 tiles via ldmatrix. +#pragma unroll + for (int k_seg = 0; k_seg < kK / kMmaK; ++k_seg) { +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const int b_row0 = warp_n * 16 + nt * 8; + unsigned b_frag[2]; + // B8 block = (b_row0>>3), sub-blocks h0 then h1 at +0/+128. + // ldmatrix: each thread supplies one matrix-row address — + // threads 0-7 feed matrix 0 (h0) rows, 8-15 matrix 1 (h1); + // the remaining threads' addresses are ignored. + const int b8_base = (b_row0 >> 3) * 256; + astrai::ldmatrix_x2( + b_frag, + &b8[b8_base + ((lane / 8) & 1) * 128 + (lane % 8) * 16]); +#pragma unroll + for (int mt = 0; mt < 4; ++mt) { + const int a_row0 = warp_m * kWarpM + mt * 16; + unsigned a_frag[4]; + // A8 block = (a_row0>>4); one x4 emits regs 0..3 in the + // exact mma A-operand order: h0m0-7, h0m8-15, h1m0-7, + // h1m8-15. Each thread supplies matrix (tid/8) row + // (tid%8) — all 32 addresses are used by x4. + const int a8_base = (a_row0 >> 4) * 512; + astrai::ldmatrix_x4( + a_frag, + &a8[a8_base + (lane / 8) * 128 + (lane % 8) * 16]); + astrai::mma_sync::type>( + acc + (nt * 4 + mt) * 4, + a_frag, b_frag, acc + (nt * 4 + mt) * 4); + } + } + } + __syncthreads(); + if (tile_index + 3 < tile_count) { + load_tile(stage, (tile_index + 3) * kK); + asm volatile("cp.async.commit_group;"); + } + } + + if constexpr (TrackAmax) { + block_reduce_amax(local_amax_a, warp_amax_a, warp, lane, + track_amax_a, amax_a); + block_reduce_amax(local_amax_b, warp_amax_b, warp, lane, + track_amax_b, amax_b); + } + + const float output_scale = sa * sb; +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const int64_t col = output_col + nt * 8; +#pragma unroll + for (int mt = 0; mt < 4; ++mt) { + const int64_t row0 = row_base + mt * 16; + const int64_t row1 = row0 + 8; + float* tile_acc = acc + (nt * 4 + mt) * 4; + if (col < n) { + float bias0 = 0.0f; + float bias1 = 0.0f; + if constexpr (AddBias) { + bias0 = __bfloat162float(bias[col]); + if (col + 1 < n) + bias1 = __bfloat162float(bias[col + 1]); + } + if (row0 < m) { + out[row0 * n + col] = + __float2bfloat16(tile_acc[0] * output_scale + bias0); + if (col + 1 < n) + out[row0 * n + col + 1] = __float2bfloat16( + tile_acc[1] * output_scale + bias1); + } + if (row1 < m) { + out[row1 * n + col] = + __float2bfloat16(tile_acc[2] * output_scale + bias0); + if (col + 1 < n) + out[row1 * n + col + 1] = __float2bfloat16( + tile_acc[3] * output_scale + bias1); + } + } + } + } +} + +// --------------------------------------------------------------------------- +// Pre-quantized kernel: FP8 A/B read straight into shared memory, FP32 +// accumulation, BF16 or FP8 output. The input format follows Traits; the +// tile is compact (row = kK bytes) so MMA fragments read directly. +// --------------------------------------------------------------------------- + +template +__global__ void fp8_pq_gemm_kernel(FP8Params p) { + using T8 = std::conditional_t; + constexpr int kBlockM = Traits::kBlockM; + constexpr int kBlockN = Traits::kBlockN; + constexpr int kK = Traits::kK; + constexpr int kStages = Traits::kStages; + __shared__ __align__(16) T8 a_tile[kStages][kBlockM][kK]; + __shared__ __align__(16) T8 b_tile[kStages][kBlockN][kK]; + + const auto* a = reinterpret_cast(p.a_ptr); + const auto* b = reinterpret_cast(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 int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int group = lane >> 2; + const int thread_in_group = lane & 3; + constexpr int warps_n = kBlockN / 16; + const int warp_m = warp / warps_n; + const int warp_n = warp % warps_n; + const int64_t row_base = blockIdx.y * kBlockM + warp_m * 64 + group; + const int64_t output_col = + blockIdx.x * kBlockN + warp_n * 16 + thread_in_group * 2; + const float sa = *p.scale_a; + const float sb = *p.scale_b; + float acc[4 * 4 * 2] = {}; + + // One A chunk (16 FP8) per thread covers the 128x32 tile; the first 128 + // threads issue the 64x32 B chunks. + auto load_tile = [&](int stage, int64_t k_base) { + const int r0 = tid >> 1; + const int c0 = (tid & 1) * 16; + const bool full_chunk = k_base + c0 + 15 < k; + const int64_t a_row = blockIdx.y * kBlockM + r0; + auto* a_dst = &a_tile[stage][r0][c0]; + const auto* a_ptr = a + a_row * k + k_base + c0; + const bool full_a = a_row < m && full_chunk; + const bool aligned_a = + (reinterpret_cast(a_ptr) & 15) == 0; + if (full_a && aligned_a) { + cp_async_16b(a_dst, a_ptr, true); + } else { +#pragma unroll + for (int i = 0; i < 16; ++i) { + a_dst[i] = a_row < m && k_base + c0 + i < k + ? a_ptr[i] + : T8(0.0f); + } + } + if (tid < 128) { + const int64_t b_row = blockIdx.x * kBlockN + r0; + auto* b_dst = &b_tile[stage][r0][c0]; + const auto* b_ptr = b + b_row * k + k_base + c0; + const bool full_b = b_row < n && full_chunk; + const bool aligned_b = + (reinterpret_cast(b_ptr) & 15) == 0; + if (full_b && aligned_b) { + cp_async_16b(b_dst, b_ptr, true); + } else { +#pragma unroll + for (int i = 0; i < 16; ++i) { + b_dst[i] = b_row < n && k_base + c0 + i < k + ? b_ptr[i] + : T8(0.0f); + } + } + } + }; + + const int64_t tile_count = (k + kK - 1) / kK; + load_tile(0, 0); + asm volatile("cp.async.commit_group;"); + if (tile_count > 1) { + load_tile(1, kK); + asm volatile("cp.async.commit_group;"); + } + if (tile_count > 2) { + load_tile(2, 2 * kK); + asm volatile("cp.async.commit_group;"); + } + for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) { + const int stage = static_cast(tile_index % kStages); + const int64_t remaining = tile_count - tile_index - 1; + if (remaining >= 2) { + asm volatile("cp.async.wait_group 2;"); + } else if (remaining == 1) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + // Barrier 1: every thread's cp.async for this stage is complete + // before any thread reads tiles written by other threads. + __syncthreads(); + +#pragma unroll + for (int k_seg = 0; k_seg < kK / kMmaK; ++k_seg) { + const int frag_col = thread_in_group * 4 + k_seg * 32; +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const int b_row = warp_n * 16 + nt * 8 + group; + unsigned b_frag[2]; + b_frag[0] = *reinterpret_cast( + &b_tile[stage][b_row][frag_col]); + b_frag[1] = *reinterpret_cast( + &b_tile[stage][b_row][frag_col + 16]); +#pragma unroll + for (int mt = 0; mt < 4; ++mt) { + const int a_row0 = warp_m * 64 + mt * 16 + group; + unsigned a_frag[4]; + a_frag[0] = *reinterpret_cast( + &a_tile[stage][a_row0][frag_col]); + a_frag[1] = *reinterpret_cast( + &a_tile[stage][a_row0 + 8][frag_col]); + a_frag[2] = *reinterpret_cast( + &a_tile[stage][a_row0][frag_col + 16]); + a_frag[3] = *reinterpret_cast( + &a_tile[stage][a_row0 + 8][frag_col + 16]); + astrai::mma_sync::type>( + acc + (nt * 4 + mt) * 4, + a_frag, b_frag, acc + (nt * 4 + mt) * 4); + } + } + } + // Barrier 2: every thread finished reading this stage's tiles before + // the prefetch for the (i+3)-th tile overwrites them. + __syncthreads(); + if (tile_index + 3 < tile_count) { + load_tile(stage, (tile_index + 3) * kK); + asm volatile("cp.async.commit_group;"); + } + } + + const float output_scale = sa * sb; + const float o8_scale = OutFp8 ? output_scale * *p.out_scale : 0.0f; +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const int64_t col = output_col + nt * 8; + // Per-row store: FP8 packs two adjacent columns into one 16-bit + // write; the BF16 path writes two scalars. Boundary columns fall + // back to a scalar convert so the pack never crosses the row edge. + auto store_out = [&](int64_t row, float v0, float v1) { + if (row >= m) return; + if constexpr (OutFp8) { + if (col + 1 < n) { + *reinterpret_cast(out_fp8 + row * n + col) = + static_cast(__nv_cvt_float2_to_fp8x2( + make_float2(v0 * o8_scale, v1 * o8_scale), + __NV_SATFINITE, __NV_E4M3)); + } else { + out_fp8[row * n + col] = __nv_fp8_e4m3(v0 * o8_scale); + } + } else { + out_bf16[row * n + col] = __float2bfloat16(v0 * output_scale); + if (col + 1 < n) + out_bf16[row * n + col + 1] = + __float2bfloat16(v1 * output_scale); + } + }; +#pragma unroll + for (int mt = 0; mt < 4; ++mt) { + const int64_t row0 = row_base + mt * 16; + float* tile_acc = acc + (nt * 4 + mt) * 4; + if (col < n) { + store_out(row0, tile_acc[0], tile_acc[1]); + store_out(row0 + 8, tile_acc[2], tile_acc[3]); + } + } + } +} + +// --------------------------------------------------------------------------- +// Launchers — pure CUDA (no torch), usable from the binding and pure C tests. +// --------------------------------------------------------------------------- + +// Fused forward tile config: 128x64 CTA, K=32, 3-stage cp.async pipeline, +// plus the fp8 ldmatrix tile region (A8[2][BlockM][16] + B8[2][BlockN][16]). +using FusedTraits = Fp8GemmTraits; +// Pre-quantized tile config: 128x64 CTA, K=32, 3-stage pipeline. +template +using PqTraits = Fp8GemmTraits; + +template +void launch_fp8_quantize(const FP8Params& p, cudaStream_t stream) { + constexpr int kThreads = 256; + const int64_t blocks = (p.total + kThreads - 1) / kThreads; + fp8_quantize_kernel<<>>(p); +} + +template +void launch_fp8_fused(const FP8Params& p, cudaStream_t stream) { + // bf16 staging (3 stages) + fp8 ldmatrix tiles (A8[2][M][16] + B8[2][N][16]) + constexpr int kSmemBytes = + FusedTraits::kStages * + (FusedTraits::kBlockM * FusedTraits::kK + + FusedTraits::kBlockN * FusedTraits::kK) * + 2 + + 2 * FusedTraits::kBlockM * 16 + 2 * FusedTraits::kBlockN * 16; + dim3 grid((p.n + FusedTraits::kBlockN - 1) / FusedTraits::kBlockN, + (p.m + FusedTraits::kBlockM - 1) / FusedTraits::kBlockM); + auto kernel = fp8_fused_gemm_kernel; + static bool attribute_set = false; + if (!attribute_set) { + cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes); + attribute_set = true; + } + kernel<<>>(p); +} + +template +void launch_fp8_pq(const FP8Params& p, cudaStream_t stream) { + using Traits = PqTraits; + dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN, + (p.m + Traits::kBlockM - 1) / Traits::kBlockM); + fp8_pq_gemm_kernel<<>>(p); +} + +} // namespace fp8 diff --git a/csrc/kernels/fp8/mm.cu b/csrc/kernels/fp8/mm.cu new file mode 100644 index 0000000..06af3ed --- /dev/null +++ b/csrc/kernels/fp8/mm.cu @@ -0,0 +1,369 @@ +// 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. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gemm.cuh" +#include "../common/device.cuh" + +namespace { + +// FP8Format / FP8Params live in the global namespace (common.h); the +// launchers live in fp8:: (gemm.cuh). + +void check_fp8_device(const torch::Tensor& tensor) { + static std::mutex mutex; + static std::unordered_map supported; + const int device = tensor.device().index(); + { + std::lock_guard 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"); + return; + } + } + + const auto* properties = at::cuda::getDeviceProperties(device); + const bool is_supported = + astrai::sm_at_least(properties->major, properties->minor, + astrai::kMinSmForFp8Major, + astrai::kMinSmForFp8Minor); + { + std::lock_guard lock(mutex); + supported.emplace(device, is_supported); + } + TORCH_CHECK(is_supported, + "fused FP8 MMA requires compute capability 8.9 or newer"); +} + +void check_scale(const torch::Tensor& scale, const torch::Tensor& input, + const char* name) { + 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"); +} + +// ---- FP8Params packing (mirrors attention/entry_utils.cuh pack_* helpers) ---- + +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, int64_t m, int64_t n, + int64_t k) { + p.a_ptr = a; + p.b_ptr = b; + p.out_ptr = out; + p.scale_a = sa.data_ptr(); + p.scale_b = sb.data_ptr(); + p.out_scale = out_scale ? out_scale->data_ptr() : nullptr; + p.bias = nullptr; + p.amax_a = nullptr; + p.amax_b = nullptr; + p.m = m; + p.n = n; + p.k = k; + p.total = 0; +} + +void pack_quantize_params(FP8Params& p, const void* x, void* x8, + const torch::Tensor& scale, torch::Tensor* amax, + int64_t total) { + p.a_ptr = x; + p.b_ptr = nullptr; + p.out_ptr = x8; + p.scale_a = scale.data_ptr(); + p.scale_b = nullptr; + p.out_scale = nullptr; + p.bias = nullptr; + p.amax_a = amax ? amax->data_ptr() : nullptr; + p.amax_b = nullptr; + p.m = p.n = p.k = 0; + p.total = total; +} + +} // namespace + +// --------------------------------------------------------------------------- +// Entry points +// --------------------------------------------------------------------------- + +std::tuple 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"); + 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)); + auto amax = torch::zeros({1}, x_c.options().dtype(torch::kFloat32)); + FP8Params p; + pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax, + x_c.numel()); + if (fmt) { + fp8::launch_fp8_quantize(p, stream.stream()); + } else { + fp8::launch_fp8_quantize(p, stream.stream()); + } + C10_CUDA_CHECK(cudaGetLastError()); + return {x8, amax}; +} + +torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa, + torch::Tensor sb, int64_t out_dtype, + c10::optional out_scale) { + // Pre-quantized FP8 GEMM: out = a @ b^T * (sa * sb), FP32 accumulation. + // out_dtype: 0 = BF16 (default), 1 = FP8 E4M3 (requires out_scale, the + // quantization step for the output — mirrors torch._scaled_mm's + // out_dtype / scale_result). Both operands share the same FP8 format. + 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"); + 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"); + TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch"); + check_scale(sa, a, "sa"); + check_scale(sb, a, "sb"); + 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(); + int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0); + 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)); + FP8Params p; + pack_gemm_params(p, a_c.data_ptr(), b_c.data_ptr(), out.data_ptr(), sa, sb, + out_fp8 ? &os : nullptr, m, n, k); + if (a.scalar_type() == torch::kFloat8_e4m3fn) { + if (out_fp8) { + fp8::launch_fp8_pq(p, stream.stream()); + } else { + fp8::launch_fp8_pq(p, stream.stream()); + } + } else { + if (out_fp8) { + fp8::launch_fp8_pq(p, stream.stream()); + } else { + fp8::launch_fp8_pq(p, stream.stream()); + } + } + C10_CUDA_CHECK(cudaGetLastError()); + return out; +} + +std::tuple linear_forward_fp8( + torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx, + torch::Tensor sw) { + // Fused BF16 -> E4M3 -> MMA -> BF16 linear forward. amax_x / amax_w are + // zero-initialized here and returned (caller does not clear them). + TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(x.scalar_type() == torch::kBFloat16 && + w.scalar_type() == torch::kBFloat16, + "x and w must be bf16"); + 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(); + auto w_c = w.contiguous(); + 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"); + // amax slots are zero-initialized here; the kernel atomically maxes in. + auto amax_x = torch::zeros({1}, x.options().dtype(torch::kFloat32)); + auto amax_w = torch::zeros({1}, x.options().dtype(torch::kFloat32)); + auto out = torch::empty({m, n}, x_c.options()); + const bool has_bias = bias.defined() && bias.numel() > 0; + if (has_bias) { + TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() && + bias.scalar_type() == torch::kBFloat16 && + bias.numel() == n, + "bias must be CUDA bf16 with shape [N]"); + } + FP8Params p; + pack_gemm_params(p, x_c.data_ptr(), w_c.data_ptr(), out.data_ptr(), sx, sw, + nullptr, m, n, k); + p.bias = has_bias ? reinterpret_cast(bias.data_ptr()) + : nullptr; + p.amax_a = amax_x.data_ptr(); + p.amax_b = amax_w.data_ptr(); + if (has_bias) { + fp8::launch_fp8_fused(p, stream.stream()); + } else { + fp8::launch_fp8_fused(p, stream.stream()); + } + C10_CUDA_CHECK(cudaGetLastError()); + + std::vector shape(x.sizes().begin(), x.sizes().end() - 1); + shape.push_back(n); + return {out.reshape(shape), amax_x, amax_w}; +} + +std::tuple +linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w, + std::vector masks, torch::Tensor sg, + torch::Tensor sw, torch::Tensor sx, int64_t fmt) { + // 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). + TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(g.scalar_type() == torch::kBFloat16 && + x.scalar_type() == torch::kBFloat16 && + w.scalar_type() == torch::kBFloat16, + "g, x, and w must be bf16"); + 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); + auto grad_weight = torch::empty_like(w); + auto grad_bias = torch::empty({0}, g.options()); + auto amax_g = torch::zeros({1}, g.options().dtype(torch::kFloat32)); + auto f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2) + : g.options().dtype(torch::kFloat8_e4m3fn); + const auto q_fmt = fmt ? FP8Format::E5M2 : FP8Format::E4M3; + + auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst, + const torch::Tensor& scale, torch::Tensor* amax) { + FP8Params qp; + pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax, + src.numel()); + if (fmt) { + fp8::launch_fp8_quantize(qp, stream.stream()); + } else { + fp8::launch_fp8_quantize(qp, stream.stream()); + } + }; + auto pq = [&](const torch::Tensor& a8, const torch::Tensor& b8, + torch::Tensor& out, const torch::Tensor& sa, + const torch::Tensor& sb, int64_t mm, int64_t nn, int64_t kk) { + FP8Params gp; + pack_gemm_params(gp, a8.data_ptr(), b8.data_ptr(), out.data_ptr(), sa, + sb, nullptr, mm, nn, kk); + if (fmt) { + fp8::launch_fp8_pq(gp, stream.stream()); + } else { + fp8::launch_fp8_pq(gp, stream.stream()); + } + }; + + // dX = g @ W: quantize g once, then g8 @ w8^T. + if (masks[0]) { + auto g8 = torch::empty({m, n}, f8opt); + quantize(g_c, g8, sg, &amax_g); + auto w_t = w_c.transpose(0, 1).contiguous(); // [K, N] + auto w8_t = torch::empty({k, n}, f8opt); + quantize(w_t, w8_t, sw, nullptr); + auto grad_input_2d = grad_input.reshape({m, k}); + pq(g8, w8_t, grad_input_2d, sg, sw, m, k, n); + } + // dW = g^T @ x: transposed layouts for both operands. + if (masks[1]) { + auto g_t = g_c.transpose(0, 1).contiguous(); // [N, M] + auto x_t = x_c.transpose(0, 1).contiguous(); // [K, M] + auto g8_t = torch::empty({n, m}, f8opt); + auto x8_t = torch::empty({k, m}, f8opt); + quantize(g_t, g8_t, sg, nullptr); + quantize(x_t, x8_t, sx, nullptr); + pq(g8_t, x8_t, grad_weight, sg, sx, n, k, m); + } + if (!masks[0] && !masks[1]) { + amax_g.copy_(g_c.abs().amax().to(torch::kFloat32)); + } + C10_CUDA_CHECK(cudaGetLastError()); + if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type()); + return {grad_input, grad_weight, grad_bias, amax_g}; +} + +torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor sx, + torch::Tensor sw) { + // BF16-in fused FP8 GEMM primitive (no bias, no amax): a @ b^T. + TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(a.scalar_type() == torch::kBFloat16 && + b.scalar_type() == torch::kBFloat16, + "a and b must be bf16"); + 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"); + TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch"); + check_scale(sx, a, "sx"); + check_scale(sw, a, "sw"); + 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(); + int64_t m = a_c.size(0), n = b_c.size(0), k = a_c.size(1); + auto out = torch::empty({m, n}, a_c.options()); + FP8Params p; + pack_gemm_params(p, a_c.data_ptr(), b_c.data_ptr(), out.data_ptr(), sx, sw, + nullptr, m, n, k); + fp8::launch_fp8_fused(p, stream.stream()); + C10_CUDA_CHECK(cudaGetLastError()); + return out; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), py::arg("sx"), + py::arg("sw"), + "Fused BF16 input, E4M3 MMA, FP32 accumulation, BF16 output GEMM"); + 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(), + "Pre-quantized FP8 GEMM: a @ b^T * (sa * sb); out_dtype 0=bf16, " + "1=fp8 e4m3 (requires out_scale)"); + m.def("linear_forward_fp8", &linear_forward_fp8, py::arg("x"), + py::arg("w"), py::arg("bias"), py::arg("sx"), py::arg("sw"), + "Fused BF16-to-FP8 linear forward; returns (out, amax_x, amax_w)"); + 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"), + "FP8 linear backward; returns (grad_input, grad_weight, grad_bias, amax_g)"); +} diff --git a/csrc/kernels/fp8_mm.cu b/csrc/kernels/fp8_mm.cu deleted file mode 100644 index b628aa2..0000000 --- a/csrc/kernels/fp8_mm.cu +++ /dev/null @@ -1,817 +0,0 @@ -// Fused BF16 -> E4M3 MMA -> BF16 matrix multiplication - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -constexpr int kMmaK = 32; -constexpr int kWarps = 8; -// Fast forward path: 128x64 CTA, 64x16 warp tile, 2-stage pipeline, dynamic -// shared memory. Mirrors the CUTLASS 58_ada_fp8_gemm threadblock geometry -// while keeping the fused BF16->FP8 quantize path. The FP8 tile overwrites -// the BF16 staging area in place. L20 opts in to only 101376 B shared per -// block; K=32 keeps the footprint at 24576 B so four CTAs/SM stay resident. -constexpr int kFastBlockM = 128; -constexpr int kFastBlockN = 64; -constexpr int kFastK = 32; -constexpr int kFastStages = 2; -constexpr int kFastSmemBytes = - kFastStages * (kFastBlockM * kFastK * 2 + kFastBlockN * kFastK * 2); - -__device__ __forceinline__ unsigned pack_fp8x4_vector(float x0, float x1, - float x2, float x3) { - const auto low = __nv_cvt_float2_to_fp8x2( - make_float2(x0, x1), __NV_SATFINITE, __NV_E4M3); - const auto high = __nv_cvt_float2_to_fp8x2( - make_float2(x2, x3), __NV_SATFINITE, __NV_E4M3); - return static_cast(low) | (static_cast(high) << 16); -} - - -__device__ __forceinline__ void mma_fp8_16832(float d[4], - const unsigned a[4], - const unsigned b[2]) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890 - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), - "r"(b[0]), "r"(b[1])); -#endif -} - -__device__ __forceinline__ void atomic_max_float(float* destination, - float value) { - if (destination) - atomicMax(reinterpret_cast(destination), __float_as_uint(value)); -} - -__device__ __forceinline__ float warp_reduce_max(float value) { -#pragma unroll - for (int offset = 16; offset; offset >>= 1) { - value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset)); - } - return value; -} - -// Block-wide max reduction of a per-warp tracked value, then an atomic -// update of the global amax slot when `track` is set. -template -__device__ __forceinline__ void block_reduce_amax(float& local, float* slots, - int warp, int lane, - bool track, float* global) { - local = warp_reduce_max(local); - if (lane == 0) slots[warp] = local; - __syncthreads(); - if (warp == 0) { - float value = lane < NWarps ? slots[lane] : 0.0f; - value = warp_reduce_max(value); - if (lane == 0 && track && global) atomic_max_float(global, value); - } -} - -// One thread moves eight BF16 values (16 bytes). The async copy is issued -// through a uint4-shaped pointer so the source and destination are both -// naturally 128-bit aligned for contiguous forward GEMMs. -__device__ __forceinline__ void cp_async_bf16_8( - __nv_bfloat16* destination, const __nv_bfloat16* source, bool valid) { - const unsigned shared_address = __cvta_generic_to_shared(destination); - const uint4* source_vec = reinterpret_cast(source); - asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" - :: "r"(shared_address), "l"(source_vec), - "r"(valid ? 16 : 0)); -} - -template -__device__ __forceinline__ unsigned load_fp8x4_from_bf16( - const __nv_bfloat16* source, float scale_inv, float& amax, - bool track_amax = true) { - float x0 = __bfloat162float(source[0]); - float x1 = __bfloat162float(source[1]); - float x2 = __bfloat162float(source[2]); - float x3 = __bfloat162float(source[3]); - if constexpr (TrackAmax) { - if (track_amax) { - amax = fmaxf(amax, fmaxf(fabsf(x0), fmaxf(fabsf(x1), - fmaxf(fabsf(x2), fabsf(x3))))); - } - } - return pack_fp8x4_vector(x0 * scale_inv, x1 * scale_inv, - x2 * scale_inv, x3 * scale_inv); -} - - -template -__global__ void fused_fp8_gemm_fast_kernel( - const __nv_bfloat16* __restrict__ a, - const __nv_bfloat16* __restrict__ b, - __nv_bfloat16* __restrict__ out, - const __nv_bfloat16* __restrict__ bias, - const float* __restrict__ scale_a, - const float* __restrict__ scale_b, - float* __restrict__ amax_a, - float* __restrict__ amax_b, - int64_t m, int64_t n, int64_t k) { - extern __shared__ char smem[]; - constexpr int a_stride = kFastBlockM * kFastK; - constexpr int b_stride = kFastBlockN * kFastK; - constexpr int b_bf16_offset = kFastStages * a_stride; - auto* a_bf16 = reinterpret_cast<__nv_bfloat16*>(smem); - auto* b_bf16 = - reinterpret_cast<__nv_bfloat16*>(smem + b_bf16_offset * 2); - __shared__ float warp_amax_a[kWarps]; - __shared__ float warp_amax_b[kWarps]; - - const int tid = threadIdx.x; - const int warp = tid >> 5; - const int lane = tid & 31; - const int group = lane >> 2; - const int thread_in_group = lane & 3; - constexpr int warps_n = kFastBlockN / 16; - const int warp_m = warp / warps_n; - const int warp_n = warp % warps_n; - const int64_t row_base = - blockIdx.y * kFastBlockM + warp_m * 64 + group; - const int64_t output_col = - blockIdx.x * kFastBlockN + warp_n * 16 + thread_in_group * 2; - const float sa = *scale_a; - const float sb = *scale_b; - const float inv_a = 1.0f / sa; - const float inv_b = 1.0f / sb; - float local_amax_a = 0.0f; - float local_amax_b = 0.0f; - float acc[4 * 4 * 2] = {}; - - const bool track_amax_a = TrackAmax && blockIdx.x == 0; - const bool track_amax_b = TrackAmax && blockIdx.y == 0; - // Each thread issues 8 A chunks and 4 B chunks of 8 BF16 (16B) per stage. - auto load_tile = [&](int stage, int64_t k_base) { - const int r0 = tid >> 2; - const int c0 = (tid & 3) * 8; -#pragma unroll - for (int j = 0; j < kFastK / 32; ++j) { - const int col = c0 + 32 * j; - const bool full_chunk = k_base + col + 7 < k; - const int64_t a_row = blockIdx.y * kFastBlockM + r0; - const int64_t b_row = blockIdx.x * kFastBlockN + r0; - auto* a_dst = &a_bf16[stage * a_stride + r0 * kFastK + col]; - auto* b_dst = &b_bf16[stage * b_stride + r0 * kFastK + col]; - const auto* a_ptr = a + a_row * k + k_base + col; - const auto* b_ptr = b + b_row * k + k_base + col; - const bool full_a = a_row < m && full_chunk; - const bool full_b = b_row < n && full_chunk; - const bool aligned_a = - (reinterpret_cast(a_ptr) & 15) == 0; - const bool aligned_b = - (reinterpret_cast(b_ptr) & 15) == 0; - if (full_a && aligned_a) { - cp_async_bf16_8(a_dst, a_ptr, true); - } else { -#pragma unroll - for (int i = 0; i < 8; ++i) { - a_dst[i] = a_row < m && k_base + col + i < k - ? a_ptr[i] - : __float2bfloat16(0.0f); - } - } - if (full_b && aligned_b) { - cp_async_bf16_8(b_dst, b_ptr, true); - } else { -#pragma unroll - for (int i = 0; i < 8; ++i) { - b_dst[i] = b_row < n && k_base + col + i < k - ? b_ptr[i] - : __float2bfloat16(0.0f); - } - } - if (r0 + 64 < kFastBlockM) { - const int64_t a_row_hi = blockIdx.y * kFastBlockM + r0 + 64; - auto* a_dst_hi = - &a_bf16[stage * a_stride + (r0 + 64) * kFastK + col]; - const auto* a_ptr_hi = a + a_row_hi * k + k_base + col; - const bool full_a_hi = a_row_hi < m && full_chunk; - const bool aligned_a_hi = - (reinterpret_cast(a_ptr_hi) & 15) == 0; - if (full_a_hi && aligned_a_hi) { - cp_async_bf16_8(a_dst_hi, a_ptr_hi, true); - } else { -#pragma unroll - for (int i = 0; i < 8; ++i) { - a_dst_hi[i] = a_row_hi < m && k_base + col + i < k - ? a_ptr_hi[i] - : __float2bfloat16(0.0f); - } - } - } - } - }; - - // Quantize must place each 4-BP8 group at the byte offset the MMA - // fragment reads: 8*(lane&3) + 64*k_seg for a row. With in-place storage - // (fp8 element k lives at byte 2k), the BF16 column of a group is - // 4*(tid&7) + 32*j, so partition by 4-element groups instead of the - // 8-element cp.async chunks. - auto quantize_tile = [&](int stage) { - const int r0 = tid >> 3; - const int c0 = (tid & 7) * 4; -#pragma unroll - for (int s = 0; s < 4; ++s) { - const int row = r0 + 32 * s; - auto* a_src = &a_bf16[stage * a_stride + row * kFastK + c0]; - auto* a_dst = reinterpret_cast(a_src); -#pragma unroll - for (int j = 0; j < kFastK / 32; ++j) { - a_dst[16 * j] = load_fp8x4_from_bf16( - a_src + 32 * j, inv_a, local_amax_a, track_amax_a); - } - } -#pragma unroll - for (int s = 0; s < 2; ++s) { - const int row = r0 + 32 * s; - auto* b_src = &b_bf16[stage * b_stride + row * kFastK + c0]; - auto* b_dst = reinterpret_cast(b_src); -#pragma unroll - for (int j = 0; j < kFastK / 32; ++j) { - b_dst[16 * j] = load_fp8x4_from_bf16( - b_src + 32 * j, inv_b, local_amax_b, track_amax_b); - } - } - }; - - const int64_t tile_count = (k + kFastK - 1) / kFastK; - load_tile(0, 0); - asm volatile("cp.async.commit_group;"); - if (tile_count > 1) { - load_tile(1, kFastK); - asm volatile("cp.async.commit_group;"); - } - for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) { - const int stage = static_cast(tile_index % kFastStages); - if (tile_index + 1 == tile_count) { - asm volatile("cp.async.wait_group 0;"); - } else { - asm volatile("cp.async.wait_group 1;"); - } - // wait_group only waits for this thread's async copies. All threads - // must finish loading before the tile is read by the CTA. - __syncthreads(); - quantize_tile(stage); - __syncthreads(); - - // Four m16n8k32 MMA segments per 128-K stage. -#pragma unroll - for (int k_seg = 0; k_seg < kFastK / kMmaK; ++k_seg) { - const int frag_col = thread_in_group * 4 + k_seg * 32; -#pragma unroll - for (int nt = 0; nt < 2; ++nt) { - const int b_row = warp_n * 16 + nt * 8 + group; - unsigned b_frag[2]; - b_frag[0] = *reinterpret_cast( - &b_bf16[stage * b_stride + b_row * kFastK + frag_col]); - b_frag[1] = *reinterpret_cast( - &b_bf16[stage * b_stride + b_row * kFastK + frag_col + 16]); -#pragma unroll - for (int mt = 0; mt < 4; ++mt) { - const int a_row0 = warp_m * 64 + mt * 16 + group; - unsigned a_frag[4]; - a_frag[0] = *reinterpret_cast( - &a_bf16[stage * a_stride + a_row0 * kFastK + frag_col]); - a_frag[1] = *reinterpret_cast( - &a_bf16[stage * a_stride + (a_row0 + 8) * kFastK + frag_col]); - a_frag[2] = *reinterpret_cast( - &a_bf16[stage * a_stride + a_row0 * kFastK + frag_col + 16]); - a_frag[3] = *reinterpret_cast( - &a_bf16[stage * a_stride + (a_row0 + 8) * kFastK + frag_col + 16]); - mma_fp8_16832(acc + (nt * 4 + mt) * 4, a_frag, b_frag); - } - } - } - __syncthreads(); - if (tile_index + 2 < tile_count) { - load_tile(stage, (tile_index + 2) * kFastK); - asm volatile("cp.async.commit_group;"); - } - } - - if constexpr (TrackAmax) { - block_reduce_amax(local_amax_a, warp_amax_a, warp, lane, - track_amax_a, amax_a); - block_reduce_amax(local_amax_b, warp_amax_b, warp, lane, - track_amax_b, amax_b); - } - - const float output_scale = sa * sb; -#pragma unroll - for (int nt = 0; nt < 2; ++nt) { - const int64_t col = output_col + nt * 8; -#pragma unroll - for (int mt = 0; mt < 4; ++mt) { - const int64_t row0 = row_base + mt * 16; - const int64_t row1 = row0 + 8; - float* tile_acc = acc + (nt * 4 + mt) * 4; - if (col < n) { - float bias0 = 0.0f; - float bias1 = 0.0f; - if constexpr (AddBias) { - bias0 = __bfloat162float(bias[col]); - if (col + 1 < n) - bias1 = __bfloat162float(bias[col + 1]); - } - if (row0 < m) { - out[row0 * n + col] = - __float2bfloat16(tile_acc[0] * output_scale + bias0); - if (col + 1 < n) - out[row0 * n + col + 1] = __float2bfloat16( - tile_acc[1] * output_scale + bias1); - } - if (row1 < m) { - out[row1 * n + col] = - __float2bfloat16(tile_acc[2] * output_scale + bias0); - if (col + 1 < n) - out[row1 * n + col + 1] = __float2bfloat16( - tile_acc[3] * output_scale + bias1); - } - } - } - } -} - -// Pre-quantized FP8-in path: FP8 A/B read straight into shared memory (no -// BF16 staging, no inline quantization), FP32 accumulation, BF16 output. -// Same 128x64 CTA / 64x16 warp tile geometry as the fused kernel; the fp8 -// tile is compact (row = kFastK bytes) so MMA fragments read directly. -constexpr int kPqBlockM = 128; -constexpr int kPqBlockN = 64; -constexpr int kPqK = 32; -constexpr int kPqStages = 3; - -template -__device__ __forceinline__ void cp_async_16b(T* destination, - const T* source, bool valid) { - const unsigned shared_address = __cvta_generic_to_shared(destination); - const uint4* source_vec = reinterpret_cast(source); - asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" - :: "r"(shared_address), "l"(source_vec), - "r"(valid ? 16 : 0)); -} - -template -__global__ void fp8_mm_pq_kernel( - const __nv_fp8_e4m3* __restrict__ a, - const __nv_fp8_e4m3* __restrict__ b, - __nv_bfloat16* __restrict__ out_bf16, - __nv_fp8_e4m3* __restrict__ out_fp8, - const float scale, const float out_scale, - int64_t m, int64_t n, int64_t k) { - __shared__ __align__(16) __nv_fp8_e4m3 a_tile[kPqStages][kPqBlockM][kPqK]; - __shared__ __align__(16) __nv_fp8_e4m3 b_tile[kPqStages][kPqBlockN][kPqK]; - - const int tid = threadIdx.x; - const int warp = tid >> 5; - const int lane = tid & 31; - const int group = lane >> 2; - const int thread_in_group = lane & 3; - constexpr int warps_n = kPqBlockN / 16; - const int warp_m = warp / warps_n; - const int warp_n = warp % warps_n; - const int64_t row_base = blockIdx.y * kPqBlockM + warp_m * 64 + group; - const int64_t output_col = - blockIdx.x * kPqBlockN + warp_n * 16 + thread_in_group * 2; - float acc[4 * 4 * 2] = {}; - - // One A chunk (16 FP8) per thread covers the 128x32 tile; the first 128 - // threads issue the 64x32 B chunks. - auto load_tile = [&](int stage, int64_t k_base) { - const int r0 = tid >> 1; - const int c0 = (tid & 1) * 16; - const bool full_chunk = k_base + c0 + 15 < k; - const int64_t a_row = blockIdx.y * kPqBlockM + r0; - auto* a_dst = &a_tile[stage][r0][c0]; - const auto* a_ptr = a + a_row * k + k_base + c0; - const bool full_a = a_row < m && full_chunk; - const bool aligned_a = - (reinterpret_cast(a_ptr) & 15) == 0; - if (full_a && aligned_a) { - cp_async_16b(a_dst, a_ptr, true); - } else { -#pragma unroll - for (int i = 0; i < 16; ++i) { - a_dst[i] = a_row < m && k_base + c0 + i < k - ? a_ptr[i] - : __nv_fp8_e4m3(0.0f); - } - } - if (tid < 128) { - const int64_t b_row = blockIdx.x * kPqBlockN + r0; - auto* b_dst = &b_tile[stage][r0][c0]; - const auto* b_ptr = b + b_row * k + k_base + c0; - const bool full_b = b_row < n && full_chunk; - const bool aligned_b = - (reinterpret_cast(b_ptr) & 15) == 0; - if (full_b && aligned_b) { - cp_async_16b(b_dst, b_ptr, true); - } else { -#pragma unroll - for (int i = 0; i < 16; ++i) { - b_dst[i] = b_row < n && k_base + c0 + i < k - ? b_ptr[i] - : __nv_fp8_e4m3(0.0f); - } - } - } - }; - - const int64_t tile_count = (k + kPqK - 1) / kPqK; - load_tile(0, 0); - asm volatile("cp.async.commit_group;"); - if (tile_count > 1) { - load_tile(1, kPqK); - asm volatile("cp.async.commit_group;"); - } - if (tile_count > 2) { - load_tile(2, 2 * kPqK); - asm volatile("cp.async.commit_group;"); - } - for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) { - const int stage = static_cast(tile_index % kPqStages); - const int64_t remaining = tile_count - tile_index - 1; - if (remaining >= 2) { - asm volatile("cp.async.wait_group 2;"); - } else if (remaining == 1) { - asm volatile("cp.async.wait_group 1;"); - } else { - asm volatile("cp.async.wait_group 0;"); - } - // Barrier 1: every thread's cp.async for this stage is complete - // before any thread reads tiles written by other threads. - __syncthreads(); - -#pragma unroll - for (int k_seg = 0; k_seg < kPqK / kMmaK; ++k_seg) { - const int frag_col = thread_in_group * 4 + k_seg * 32; -#pragma unroll - for (int nt = 0; nt < 2; ++nt) { - const int b_row = warp_n * 16 + nt * 8 + group; - unsigned b_frag[2]; - b_frag[0] = *reinterpret_cast( - &b_tile[stage][b_row][frag_col]); - b_frag[1] = *reinterpret_cast( - &b_tile[stage][b_row][frag_col + 16]); -#pragma unroll - for (int mt = 0; mt < 4; ++mt) { - const int a_row0 = warp_m * 64 + mt * 16 + group; - unsigned a_frag[4]; - a_frag[0] = *reinterpret_cast( - &a_tile[stage][a_row0][frag_col]); - a_frag[1] = *reinterpret_cast( - &a_tile[stage][a_row0 + 8][frag_col]); - a_frag[2] = *reinterpret_cast( - &a_tile[stage][a_row0][frag_col + 16]); - a_frag[3] = *reinterpret_cast( - &a_tile[stage][a_row0 + 8][frag_col + 16]); - mma_fp8_16832(acc + (nt * 4 + mt) * 4, a_frag, b_frag); - } - } - } - // Barrier 2: every thread finished reading this stage's tiles before - // the prefetch for the (i+3)-th tile overwrites them. - __syncthreads(); - if (tile_index + 3 < tile_count) { - load_tile(stage, (tile_index + 3) * kPqK); - asm volatile("cp.async.commit_group;"); - } - } - - const float output_scale = scale * out_scale; -#pragma unroll - for (int nt = 0; nt < 2; ++nt) { - const int64_t col = output_col + nt * 8; - // Per-row store: FP8 packs two adjacent columns into one 16-bit - // write; the BF16 path writes two scalars. Boundary columns fall - // back to a scalar convert so the pack never crosses the row edge. - auto store_out = [&](int64_t row, float v0, float v1) { - if (row >= m) return; - if constexpr (OutFp8) { - if (col + 1 < n) { - *reinterpret_cast( - out_fp8 + row * n + col) = - static_cast(__nv_cvt_float2_to_fp8x2( - make_float2(v0 * output_scale, v1 * output_scale), - __NV_SATFINITE, __NV_E4M3)); - } else { - out_fp8[row * n + col] = __nv_fp8_e4m3(v0 * output_scale); - } - } else { - out_bf16[row * n + col] = __float2bfloat16(v0 * scale); - if (col + 1 < n) - out_bf16[row * n + col + 1] = __float2bfloat16(v1 * scale); - } - }; -#pragma unroll - for (int mt = 0; mt < 4; ++mt) { - const int64_t row0 = row_base + mt * 16; - float* tile_acc = acc + (nt * 4 + mt) * 4; - if (col < n) { - store_out(row0, tile_acc[0], tile_acc[1]); - store_out(row0 + 8, tile_acc[2], tile_acc[3]); - } - } - } -} - -template -void launch_fused_fp8_gemm_fast( - const torch::Tensor& a, const torch::Tensor& b, torch::Tensor& out, - const torch::Tensor& bias, const torch::Tensor& scale_a, - const torch::Tensor& scale_b, torch::Tensor* amax_a, - torch::Tensor* amax_b, int64_t m, int64_t n, int64_t k, - cudaStream_t stream) { - dim3 grid((n + kFastBlockN - 1) / kFastBlockN, - (m + kFastBlockM - 1) / kFastBlockM); - const auto* bias_ptr = AddBias - ? reinterpret_cast(bias.data_ptr()) - : nullptr; - auto kernel = fused_fp8_gemm_fast_kernel; - static bool attribute_set = false; - if (!attribute_set) { - C10_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - kFastSmemBytes)); - attribute_set = true; - } - kernel<<>>( - reinterpret_cast(a.data_ptr()), - reinterpret_cast(b.data_ptr()), - reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), bias_ptr, - scale_a.data_ptr(), scale_b.data_ptr(), - amax_a ? amax_a->data_ptr() : nullptr, - amax_b ? amax_b->data_ptr() : nullptr, m, n, k); -} - -void check_fp8_device(const torch::Tensor& tensor) { - static std::mutex mutex; - static std::unordered_map supported; - const int device = tensor.device().index(); - { - std::lock_guard 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"); - return; - } - } - - const auto* properties = at::cuda::getDeviceProperties(device); - const bool is_supported = properties->major > 8 || - (properties->major == 8 && properties->minor >= 9); - { - std::lock_guard lock(mutex); - supported.emplace(device, is_supported); - } - TORCH_CHECK(is_supported, - "fused FP8 MMA requires compute capability 8.9 or newer"); -} - -void check_scale(const torch::Tensor& scale, const torch::Tensor& input, - const char* name) { - 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"); -} - -} // namespace - -torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor sx, - torch::Tensor sw) { - TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(a.scalar_type() == torch::kBFloat16 && - b.scalar_type() == torch::kBFloat16, - "a and b must be bf16"); - 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"); - TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch"); - check_scale(sx, a, "sx"); - check_scale(sw, a, "sw"); - 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(); - auto out = torch::empty({a_c.size(0), b_c.size(0)}, a_c.options()); - torch::Tensor no_bias; - launch_fused_fp8_gemm_fast( - a_c, b_c, out, no_bias, sx, sw, nullptr, nullptr, - a_c.size(0), b_c.size(0), a_c.size(1), stream.stream()); - C10_CUDA_CHECK(cudaGetLastError()); - return out; -} - -torch::Tensor fp8_linear_forward_scaled( - torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx, - torch::Tensor sw, torch::Tensor sx_inv, torch::Tensor sw_inv, - torch::Tensor amax_x, torch::Tensor amax_w) { - TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(x.scalar_type() == torch::kBFloat16 && - w.scalar_type() == torch::kBFloat16, - "x and w must be bf16"); - 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(); - auto w_c = w.contiguous(); - 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"); - C10_CUDA_CHECK(cudaMemsetAsync(amax_x.data_ptr(), 0, sizeof(float), - stream.stream())); - C10_CUDA_CHECK(cudaMemsetAsync(amax_w.data_ptr(), 0, sizeof(float), - stream.stream())); - auto out = torch::empty({m, n}, x_c.options()); - if (bias.defined() && bias.numel() > 0) { - TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() && - bias.scalar_type() == torch::kBFloat16 && - bias.numel() == n, - "bias must be CUDA bf16 with shape [N]"); - launch_fused_fp8_gemm_fast( - x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w, - m, n, k, stream.stream()); - } else { - launch_fused_fp8_gemm_fast( - x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w, - m, n, k, stream.stream()); - } - C10_CUDA_CHECK(cudaGetLastError()); - - (void)sx_inv; - (void)sw_inv; - std::vector shape(x.sizes().begin(), x.sizes().end() - 1); - shape.push_back(n); - return out.reshape(shape); -} - -std::tuple fp8_linear_backward_scaled( - torch::Tensor g, torch::Tensor x, torch::Tensor w, - std::vector masks, torch::Tensor sg, torch::Tensor sw, - torch::Tensor sx, torch::Tensor sg_inv, torch::Tensor sw_inv, - torch::Tensor sx_inv, torch::Tensor amax_g) { - TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(g.scalar_type() == torch::kBFloat16 && - x.scalar_type() == torch::kBFloat16 && - w.scalar_type() == torch::kBFloat16, - "g, x, and w must be bf16"); - 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(); - auto x_c = x.reshape({-1, x.size(-1)}).contiguous(); - auto w_c = w.contiguous(); - 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); - auto grad_weight = torch::empty_like(w); - auto grad_bias = torch::empty({0}, g.options()); - C10_CUDA_CHECK(cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), - stream.stream())); - torch::Tensor no_bias; - bool recorded_amax = false; - if (masks[0]) { - auto grad_input_2d = grad_input.reshape({m, k}); - // The fast kernel computes A @ B^T. A contiguous W^T makes dX use - // the same coalesced forward tile path instead of scalar fragments. - auto w_t = w_c.transpose(0, 1).contiguous(); - launch_fused_fp8_gemm_fast( - g_c, w_t, grad_input_2d, no_bias, sg, sw, &amax_g, nullptr, - m, k, n, stream.stream()); - recorded_amax = true; - } - if (masks[1]) { - // dW = G^T @ X, expressed as (G^T) @ (X^T)^T for the same kernel. - auto g_t = g_c.transpose(0, 1).contiguous(); - auto x_t = x_c.transpose(0, 1).contiguous(); - if (recorded_amax) { - launch_fused_fp8_gemm_fast( - g_t, x_t, grad_weight, no_bias, sg, sx, nullptr, nullptr, - n, k, m, stream.stream()); - } else { - launch_fused_fp8_gemm_fast( - g_t, x_t, grad_weight, no_bias, sg, sx, &amax_g, nullptr, - n, k, m, stream.stream()); - } - recorded_amax = true; - } - if (!recorded_amax) { - amax_g.copy_(g_c.abs().amax().to(torch::kFloat32)); - } - C10_CUDA_CHECK(cudaGetLastError()); - if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type()); - - (void)sg_inv; - (void)sw_inv; - (void)sx_inv; - return {grad_input, grad_weight, grad_bias}; -} - -torch::Tensor fp8_mm_prequant(torch::Tensor a, torch::Tensor b, - torch::Tensor scale) { - TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn && - b.scalar_type() == torch::kFloat8_e4m3fn, - "a and b must be fp8_e4m3fn"); - 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"); - TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch"); - check_scale(scale, a, "scale"); - 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(); - int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0); - auto out = torch::empty({m, n}, - a_c.options().dtype(torch::kBFloat16)); - const float scale_value = scale.item(); - dim3 grid((n + kPqBlockN - 1) / kPqBlockN, - (m + kPqBlockM - 1) / kPqBlockM); - fp8_mm_pq_kernel<<>>( - reinterpret_cast(a_c.data_ptr()), - reinterpret_cast(b_c.data_ptr()), - reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), nullptr, - scale_value, 1.0f, m, n, k); - C10_CUDA_CHECK(cudaGetLastError()); - return out; -} - -torch::Tensor fp8_mm_prequant_fp8(torch::Tensor a, torch::Tensor b, - torch::Tensor scale, - torch::Tensor out_scale) { - TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn && - b.scalar_type() == torch::kFloat8_e4m3fn, - "a and b must be fp8_e4m3fn"); - 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"); - TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch"); - check_scale(scale, a, "scale"); - check_scale(out_scale, a, "out_scale"); - 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(); - int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0); - auto out = torch::empty({m, n}, a_c.options()); - const float scale_value = scale.item(); - const float out_scale_value = out_scale.item(); - dim3 grid((n + kPqBlockN - 1) / kPqBlockN, - (m + kPqBlockM - 1) / kPqBlockM); - fp8_mm_pq_kernel<<>>( - reinterpret_cast(a_c.data_ptr()), - reinterpret_cast(b_c.data_ptr()), nullptr, - reinterpret_cast<__nv_fp8_e4m3*>(out.data_ptr()), - scale_value, out_scale_value, m, n, k); - C10_CUDA_CHECK(cudaGetLastError()); - return out; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), py::arg("sx"), - py::arg("sw"), - "Fused BF16 input, E4M3 MMA, FP32 accumulation, BF16 output GEMM"); - m.def("fp8_mm_prequant", &fp8_mm_prequant, py::arg("a"), py::arg("b"), - py::arg("scale"), - "Pre-quantized FP8 GEMM with FP32 accumulation and BF16 output"); - m.def("fp8_mm_prequant_fp8", &fp8_mm_prequant_fp8, py::arg("a"), - py::arg("b"), py::arg("scale"), py::arg("out_scale"), - "Pre-quantized FP8 GEMM with FP32 accumulation and FP8 output"); - m.def("fp8_linear_forward_scaled", &fp8_linear_forward_scaled, - py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("sx"), - py::arg("sw"), py::arg("sx_inv"), py::arg("sw_inv"), - py::arg("amax_x"), py::arg("amax_w"), - "Fused BF16-to-FP8 linear forward with FP32 accumulation"); - m.def("fp8_linear_backward_scaled", &fp8_linear_backward_scaled, - py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"), - py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("sg_inv"), - py::arg("sw_inv"), py::arg("sx_inv"), py::arg("amax_g"), - "Fused BF16-to-FP8 linear backward with FP32 accumulation"); -} diff --git a/csrc/kernels/rotary_emb.cu b/csrc/kernels/rotary/rotary_emb.cu similarity index 100% rename from csrc/kernels/rotary_emb.cu rename to csrc/kernels/rotary/rotary_emb.cu diff --git a/csrc/tests/attn_paged_test.cu b/csrc/tests/attn_paged_test.cu index 25c7f7c..9010526 100644 --- a/csrc/tests/attn_paged_test.cu +++ b/csrc/tests/attn_paged_test.cu @@ -7,7 +7,7 @@ #include #include #include "test_utils.cuh" -#include "../kernels/attn_dispatchers.cuh" +#include "../kernels/attention/dispatchers.cuh" struct PagedDecodeDispatch { AttentionParams& p; template void operator()() { dispatch_paged_decode(p, 0); } }; struct PagedPrefillDispatch { AttentionParams& p; template void operator()() { dispatch_paged_prefill(p, 0); } }; diff --git a/csrc/tests/attn_test.cu b/csrc/tests/attn_test.cu index c5b01a8..b42c2ec 100644 --- a/csrc/tests/attn_test.cu +++ b/csrc/tests/attn_test.cu @@ -7,7 +7,7 @@ nvcc -I csrc -arch=sm_89 -O3 \ */ #include "test_utils.cuh" -#include "../kernels/attn_dispatchers.cuh" +#include "../kernels/attention/dispatchers.cuh" struct DecodeDispatch { AttentionParams& p; template void operator()() { dispatch_decode(p, 0); } }; struct PrefillDispatch { AttentionParams& p; template void operator()() { dispatch_prefill(p, 0); } }; diff --git a/csrc/tests/fp8_mma_test.cu b/csrc/tests/fp8_mma_test.cu index c949258..334c9c2 100644 --- a/csrc/tests/fp8_mma_test.cu +++ b/csrc/tests/fp8_mma_test.cu @@ -10,6 +10,8 @@ nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \ #include +#include "../kernels/common/mma.cuh" + #include #include @@ -37,17 +39,6 @@ __device__ __forceinline__ unsigned load_quantize_fp8x4( __bfloat162float(src[3]) * scale_inv); } -__device__ __forceinline__ void mma_fp8_16832(float d[4], - const unsigned a[4], - const unsigned b[2]) { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), - "r"(b[0]), "r"(b[1])); -} - __global__ void fused_bf16_fp8_mma_kernel( const bf16* __restrict__ a, const bf16* __restrict__ b, bf16* __restrict__ out, float scale_a, float scale_b) { @@ -71,7 +62,7 @@ __global__ void fused_bf16_fp8_mma_kernel( b_frag[1] = load_quantize_fp8x4(&b[group * K + k0 + 16], 1.0f / scale_b); float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; - mma_fp8_16832(acc, a_frag, b_frag); + astrai::mma_sync<__nv_fp8_e4m3>(acc, a_frag, b_frag, acc); const int col = thread_in_group * 2; const float output_scale = scale_a * scale_b; diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index 99ec6cb..511700b 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -1,11 +1,30 @@ -"""Fused BF16-boundary FP8 MMA kernel tests.""" +"""FP8 primitives: kernel-level (CUDA) and policy-level (CPU-verifiable) tests. + +The kernel-level tests exercise the fused and pre-quantized CUDA paths; the +policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks +of the custom ops) run without a GPU. +""" import pytest import torch +from astrai.extension.fp8 import ( + DelayedScaling, + DynamicScaling, + FP8Format, + FP8TensorMeta, + fp8_autocast, + fp8_state, +) from astrai.extension.loader import get_module, is_available +from astrai.extension.ops.fp8 import ( + linear_backward_fp8, + linear_forward_fp8, + mm_fp8, + quantize_bf16, +) -pytestmark = pytest.mark.skipif( +_GPU = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9) or not is_available("fp8_mm"), @@ -21,6 +40,12 @@ def _quantize(tensor, scale): return (tensor.float() / scale).to(torch.float8_e4m3fn).float() +# -------------------------------------------------------------------------- +# Kernel-level (CUDA) +# -------------------------------------------------------------------------- + + +@_GPU @pytest.mark.parametrize( ("m", "n", "k"), [(16, 8, 32), (17, 9, 33), (31, 15, 64), (32, 48, 96)], @@ -42,6 +67,31 @@ def test_fused_fp8_mma_matches_explicit_quantization(m, n, k): torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) +@_GPU +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).""" + torch.manual_seed(3) + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor([0.5], device="cuda") + x8, amax = quantize_bf16(x, scale, "e4m3") + assert x8.dtype == torch.float8_e4m3fn + 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) + assert torch.equal(x8, ref) + + +@_GPU +def test_quantize_bf16_e5m2_format(): + x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + x8, amax = quantize_bf16(x, torch.tensor([0.1], device="cuda"), "e5m2") + assert x8.dtype == torch.float8_e5m2 + torch.testing.assert_close(amax, x.abs().amax().float().reshape(1)) + + +@_GPU def test_fused_fp8_linear_forward_and_backward(): torch.manual_seed(7) m, n, k = 19, 13, 37 @@ -50,34 +100,10 @@ def test_fused_fp8_linear_forward_and_backward(): 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) - amax_x = torch.empty(1, device="cuda", dtype=torch.float32) - amax_w = torch.empty(1, device="cuda", dtype=torch.float32) - amax_g = torch.empty(1, device="cuda", dtype=torch.float32) - module = get_module("fp8_mm") - out = module.fp8_linear_forward_scaled( - x, - weight, - bias, - scale_x, - scale_w, - scale_x.reciprocal(), - scale_w.reciprocal(), - amax_x, - amax_w, - ) - grad_x, grad_w, grad_b = module.fp8_linear_backward_scaled( - grad, - x, - weight, - [1, 1, 1], - scale_g, - scale_w, - scale_x, - scale_g.reciprocal(), - scale_w.reciprocal(), - scale_x.reciprocal(), - amax_g, + out, 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" ) qx = _quantize(x, scale_x) @@ -96,61 +122,178 @@ def test_fused_fp8_linear_forward_and_backward(): torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1)) -def test_fp8_mm_prequant_matches_scaled_mm(): +@_GPU +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)) + + +@_GPU +def test_mm_fp8_matches_scaled_mm(): torch.manual_seed(11) m, n, k = 512, 4096, 4096 a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) - a8 = a.to(torch.float8_e4m3fn) - w8 = weight.to(torch.float8_e4m3fn) - scale = torch.tensor([2.5], device="cuda") - - out = get_module("fp8_mm").fp8_mm_prequant(a8, w8, scale) - - # Reference via fp64: FP8 quantization error is dominated by the 3-bit - # mantissa, so the tolerance must track the input quantization scale. - ref = (a8.float().double() @ w8.float().double().t() * 2.5).to(torch.bfloat16) + b = torch.randn(n, k, 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) assert out.dtype == torch.bfloat16 assert out.shape == (m, n) + + ref = (a8.float().double() @ b8.float().double().t() * 2.5 * 1.5).to(torch.bfloat16) torch.testing.assert_close(out, ref, atol=6.0, rtol=0.05) - # Cross-check against torch's native FP8 GEMM on identical inputs. try: - torch._scaled_mm( - a8, - w8.t(), - torch.full((m, 1), 2.5, device="cuda"), - torch.ones((1, n), device="cuda"), - out_dtype=torch.bfloat16, - ) + torch._scaled_mm(a8, b8.t(), sa, sb, out_dtype=torch.bfloat16) except (RuntimeError, NotImplementedError): return torch.testing.assert_close( out, - torch._scaled_mm( - a8, - w8.t(), - torch.full((m, 1), 2.5, device="cuda"), - torch.ones((1, n), device="cuda"), - out_dtype=torch.bfloat16, - ), + torch._scaled_mm(a8, b8.t(), sa, sb, out_dtype=torch.bfloat16), atol=2.0, rtol=0.01, ) -def test_fp8_mm_prequant_fp8_output(): - torch.manual_seed(13) - m, n, k = 512, 4096, 4096 +@_GPU +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) - weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) - a8 = a.to(torch.float8_e4m3fn) - w8 = weight.to(torch.float8_e4m3fn) - scale = torch.tensor([2.5], device="cuda") - out_scale = torch.tensor([0.1], device="cuda") + b = torch.randn(n, k, 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) - out = get_module("fp8_mm").fp8_mm_prequant_fp8(a8, w8, scale, out_scale) - assert out.dtype == torch.float8_e4m3fn - assert out.shape == (m, n) - ref = (a8.float().double() @ w8.float().double().t() * 2.5 * 0.1).to(torch.bfloat16) - torch.testing.assert_close(out.float().to(torch.bfloat16), ref, atol=1.0, rtol=0.05) + ref = (a8.float().double() @ b8.float().double().t() * 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) +# -------------------------------------------------------------------------- + + +def test_recipe_scale_from_history(): + """Delayed: max over the window + margin; dynamic: current amax.""" + hist = torch.tensor([1.0, 2.0, 0.5]) + d = DelayedScaling(history_len=3, margin=0) + assert torch.allclose(d.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0)) + d_m = DelayedScaling(history_len=3, margin=2) + assert torch.allclose( + d_m.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0 / 4.0) + ) + dyn = DynamicScaling() + amax = torch.tensor([0.25]) + assert torch.allclose( + dyn.scale_from_history(amax, "e4m3"), torch.tensor(0.25 / 448.0) + ) + assert torch.allclose( + dyn.scale_from_history(amax, "e5m2"), torch.tensor(0.25 / 57344.0) + ) + + +def test_fp8_format_enum(): + assert FP8Format.HYBRID.fwd() == "e4m3" + assert FP8Format.HYBRID.bwd() == "e5m2" + assert FP8Format.E4M3.fwd() == FP8Format.E4M3.bwd() == "e4m3" + assert FP8Format.E5M2.fwd() == FP8Format.E5M2.bwd() == "e5m2" + + +def test_fp8_autocast_context(): + """fp8_autocast sets and restores recipe + format on the global state.""" + state = fp8_state() + prev = (state.enabled, state.recipe, state.fp8_format) + try: + with fp8_autocast(enabled=True, fp8_format="hybrid", update_interval=8): + assert state.enabled + assert isinstance(state.recipe, DelayedScaling) + assert state.recipe.history_len == 8 + assert state.fp8_format is FP8Format.HYBRID + with fp8_autocast(enabled=True, recipe=DynamicScaling(), fp8_format="e4m3"): + assert isinstance(state.recipe, DynamicScaling) + assert state.fp8_format is FP8Format.E4M3 + assert state.fp8_format is FP8Format.HYBRID # restored on exit + assert not state.enabled + finally: + state.enabled, state.recipe, state.fp8_format = prev + + +def test_fp8_tensor_meta_delayed_update(): + """Meta seeds from data and refreshes the scale from the amax ring.""" + meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0)) + w = torch.randn(8, 8) + meta.init_w(w, "e4m3") + assert meta.w_init + torch.testing.assert_close(meta.w_scale, (w.abs().amax() / 448.0).reshape(1)) + meta.update_w(torch.tensor([4.0]), "e4m3") + torch.testing.assert_close(meta.w_scale, torch.tensor(4.0 / 448.0).reshape(1)) + + +def test_quantize_bf16_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") + assert x8.dtype == 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)) + + +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().t() * 2.0 * 0.5).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().t() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn) + assert torch.equal(out8, ref)