ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: fast interior loop on the big cta and fused epilogue bias

Browse Source 
- re-enable kFastLoop on the 128x128 CTA for congruous layouts: the base-pair fragment addressing freed the registers the old offset tables spilled, and the predication-free interior loop now wins across the band (fast body 142 SASS instr with zero predicated fallback vs 719/136 generic; 128 regs, no spill)
- move the big/small CTA dispatch boundary from 3/4 to 5/8 wave: with the fast big-CTA loop the crossover sits between 49 and 63 tiles (63-tile rect +8%, 1024^3 now takes the big CTA)
- fuse the linear bias into the GEMM epilogue: FP8Params.bias_ptr adds in fp32 before the single bf16 rounding, replacing the separate out + bias elementwise pass; guarded loads keep N tails exact and batch broadcast falls out of the row-major layout
- resolve Python None bias in the pybind layer (py::object + cast) so ops/fp8.py and fp8.py pass the argument through untouched; drop the _empty_bias sentinel machinery
- add fused-bias tests covering odd N tails, no-bias parity and batched broadcast

Benchmark: L20 (sm_89), CUDA-graph e2e. Big-CTA fast loop + dispatch: 1024^3 102.6->106.3T, 1152^3 128.5->133.3T, 2048^3 173.8->178.2T, 3072^3 180.2->185.3T, 8192^3 196.2->197.7T. Bias fusion (with-bias GEMM vs unfused out + bias): 1024^3 90.5->106.1T (+17%), 2048^3 162.2->178.3T (+10%), 4096^3 178.2->191.1T (+7%). Fused bias differs from the split path by <=1 bf16 ulp and is closer to the fp64 reference. 596 tests pass.
This commit is contained in:
ViperEkura
2026-08-26 14:52:06 +08:00
parent f7d96455a5
commit a92bf79295
6 changed files with 140 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/fp8.py
+12 -26
View File
@@ -315,27 +315,14 @@ def _is_fp8(dtype: torch.dtype) -> bool:
return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
_zero_bias: Dict[Optional[int], torch.Tensor] = {}
def _empty_bias(x: torch.Tensor) -> torch.Tensor:
"""Per-device cached 0-element bf16 bias (the binding only checks numel —
never mutated), saving a CUDA allocation per bias-less linear."""
key = x.device.index
t = _zero_bias.get(key)
if t is None:
t = torch.empty(0, device=x.device, dtype=torch.bfloat16)
_zero_bias[key] = t
return t
def fp8_linear_forward(
x: torch.Tensor, w: torch.Tensor, bias=None, cfg: Optional[_ActiveConfig] = None
):
"""Scaled fp8 linear forward (called from the aten::linear impl).
Composed from the two stateless primitives: quantize x/w with the active
scales, run the pre-quantized GEMM, add the bias. Delayed scaling folds
scales, run the pre-quantized GEMM with the bias fused into its epilogue.
Delayed scaling folds
the returned amax into the history ring and publishes the next scale;
dynamic scaling measures the current amax itself. Training quantizes the
weight every step (the optimizer bumps its version, so there is no cast
@@ -345,17 +332,18 @@ def fp8_linear_forward(
if cfg is None:
cfg = _current_config()
fmt = cfg.fp8_format.fwd()
if bias is None:
bias = _empty_bias(x)
if isinstance(cfg.recipe, DynamicScaling):
sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
sw = _dynamic_scale(w, cfg.recipe, fmt)
x8, _ = quantize(x, sx.reciprocal(), fmt)
w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape(
*x.shape[:-1], w.size(0)
)
return (out + bias if bias.numel() else out), sx, sw
# Bias fuses into the GEMM epilogue (fp32 add before the single bf16
# rounding — one rounding fewer than the separate out + bias pass);
# None passes through to the kernel's no-bias path.
out = mm_fp8(
x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
).reshape(*x.shape[:-1], w.size(0))
return out, sx, sw
meta = state.get_weight_meta(w)
if not meta.w.initialized:
@@ -368,11 +356,9 @@ def fp8_linear_forward(
w8, amax_w = w, None
else:
w8, amax_w = quantize(w, sw.reciprocal(), fmt)
out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape(
*x.shape[:-1], w.size(0)
)
if bias.numel():
out = out + bias
out = mm_fp8(
x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
).reshape(*x.shape[:-1], w.size(0))
meta.x.update(amax_x, fmt)
if amax_w is not None:
meta.w.update(amax_w, fmt)
astrai/extension/ops/fp8.py
+23 -13
View File
@@ -14,7 +14,7 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
this module is stateless.
"""
from typing import Tuple
from typing import Optional, Tuple
import torch
from torch.library import custom_op
@@ -84,16 +84,19 @@ def fp8_gemm(
scale: torch.Tensor,
trans_a: int = 0,
trans_b: int = 0,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""FP8 GEMM: ``a @ b * scale`` with FP32 accumulation.
"""FP8 GEMM: ``a @ b * scale (+ bias)`` with FP32 accumulation.
2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules).
The result is always BF16; FP8 output is a separate quantize operation.
``bias`` (bf16, length n) fuses into the epilogue in fp32 before the
single bf16 rounding. The result is always BF16; FP8 output is a
separate quantize operation.
"""
@fp8_gemm.register_fake
def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0):
def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0, bias=None):
dtype = torch.bfloat16
rows = a.size(2) if trans_a else a.size(1)
cols = b.size(1) if trans_b else b.size(2)
@@ -103,19 +106,21 @@ def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0):
@fp8_gemm.register_kernel("cuda")
def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0):
def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0, bias=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 get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b)
return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
@fp8_gemm.register_kernel("cpu")
def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0):
def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
aa = a.float().transpose(-2, -1) if trans_a else a.float()
bb = b.float().transpose(-2, -1) if trans_b else b.float()
acc = aa @ bb * scale
if bias is not None and bias.numel() > 0:
acc = acc + bias.float()
return acc.to(torch.bfloat16)
@@ -149,21 +154,26 @@ def mm_fp8(
scale: torch.Tensor,
trans_a: bool = False,
trans_b: bool = False,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Pre-quantized FP8 GEMM: ``a @ b * scale``.
"""Pre-quantized FP8 GEMM: ``a @ b * scale (+ bias)``.
``a``/``b`` must be FP8 tensors of the same format, 2D or 3D (batched,
matmul-style broadcast on the batch dim). Inner-transposed views (e.g.
``x.t()``) fold into the layout at zero copy. ``scale`` is their combined
dequantization scale. The result is BF16; FP8 output is a separate
quantize operation.
dequantization scale. ``bias`` (CUDA bf16 1D of length n) adds inside the
kernel epilogue in fp32 — no separate elementwise pass. The result is
BF16; FP8 output is a separate quantize operation.
"""
# Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep
# validation identical on the direct route.
# validation identical on the direct route (bias may be None — the
# binding resolves it to the no-bias path).
if (
type(a) is torch.Tensor
and a.is_cuda
and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
):
return get_module("fp8_ops").mm_fp8(a, b, scale, int(trans_a), int(trans_b))
return fp8_gemm(a, b, scale, trans_a, trans_b)
return get_module("fp8_ops").mm_fp8(
a, b, scale, int(trans_a), int(trans_b), bias
)
return fp8_gemm(a, b, scale, trans_a, trans_b, bias)