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