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

perf: rebuild decode gemm dispatch around shape-driven tile configs

Browse Source 
- split-K removed entirely: tiled kernel walks K in one pass, no partials/semas workspace, no memset, single launch per call
- skinny GEMM (M<=8) dispatch table replaces the hand-written switch
- shape-driven four-family table replaces plan_gemm: wide-N (n>=4096) default {16,64,64,3,128} with BM=32 at M>16; narrow-N deep-K rings {16,32,256,2,64} while the grid fits one wave, {16,32,128,2,64} past it
- narrow-N is K-serial: widening the grid measurably does nothing (BN 64->32 ties, doubled m_tiles tie, kv at 4 blocks ties q/o at 24); deeper K chunks win until 72KB smem forces one CTA per SM and past one wave the 2-wave quantization loses to BK=128
- launch-check macros in common/launch.cuh; smem opt-in for the 72KB/60KB rings
- rename kernels/bf16_*.cu to gemm.cu/swiglu.cu; module names unchanged
- Python gate: lm_head (N>32768) falls back to cuBLAS, band narrows to M<=32
- drop the stale per-op benchmark narratives; fold the live numbers into cuda_kernels.md

Benchmark: NVIDIA L20 (sm_89, 92 SMs), CUDA 12.8, bf16, L2-thrash weight rotation, per-call medians at M=16: q/o 9.5us, kv 8.6us, gate/up 33.3us, down 33.7us (down -29% vs prior default). End-to-end 1B decode (gen 128, 3 trials, tokens/s vs cuBLAS): B=1 260 vs 252, B=8 1660 vs 1446, B=16 2464 vs 2437, B=32 3620 vs 3690. Prior split-K dispatch measured B=16 2243 / B=32 3393.
This commit is contained in:
ViperEkura
2026-09-04 22:41:39 +08:00
parent 8e39d9d8c9
commit 1798474316
21 changed files with 1098 additions and 697 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/__init__.py
+3 -3
View File
@@ -5,7 +5,7 @@ Public API:
families with safe torch fallbacks (see ``astrai.extension.backend``)
- ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` /
``attn_paged_prefill`` — direct attention kernel wrappers
- ``bf16_gemv`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
- ``bf16_gemm`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
- ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` /
``FlashAttnBackend`` — attention backend strategies
- ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared
@@ -53,7 +53,7 @@ from astrai.extension.ops import (
attn_decode,
attn_paged_decode,
attn_prefill,
bf16_gemv,
bf16_gemm,
bf16_swiglu,
)
@@ -73,7 +73,7 @@ __all__ = [
"attn_decode",
"attn_paged_decode",
"attn_prefill",
"bf16_gemv",
"bf16_gemm",
"bf16_swiglu",
"is_available",
"KERNEL_NAMES",
astrai/extension/backend/linear.py
+53 -48
View File
@@ -1,17 +1,15 @@
"""Inference-only dispatch for AstrAI linear layers.
The CUDA GEMV path is narrow by construction rather than by a measured
shape table: the kernel streams each weight exactly once, so automatic
selection is keyed on the decode batch size alone (M in [2, 4], where it
sits at the HBM bandwidth floor and beat the cuBLAS small-M path on every
measured family). Every training, prefill-sized, out-of-band, or
unsupported call falls back to PyTorch.
The CUDA GEMM path is sized by the decode batch M. M in [1, 8] uses the
register-resident GEMV kernel (any K); M in (8, 64] uses the tiled
kernel (K % 8 == 0, 16-byte-aligned tensors). Automatic mode
selects GEMM for M in [1, 64] where the primitive is capable; every
training, prefill-sized, or unsupported call falls back to PyTorch.
The family stays registered with the shared operator dispatcher, so
``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
``explain`` keep working like for attention and rotary. The per-layer
hot path only consults the dispatcher when one of those selections is
active, keeping it free of axes dictionaries and record sorting.
``explain`` keep working. The per-layer hot path only consults the
dispatcher when one of those selections is active.
"""
from typing import Any, Dict, List, Optional
@@ -32,33 +30,25 @@ from astrai.extension.dispatch import (
tensor_axes,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.gemv import bf16_gemv
# M=1 keeps cuBLAS (its GEMV path is already at the bandwidth floor; only
# OPT 1.3B shapes ever passed the full gate). M >= 5 approaches the cuBLAS
# tensor-core crossover (M=8 regressed at wrapper level on every measured
# family, and cuBLAS clearly wins from M ~ 12).
_AUTO_GEMV_M = frozenset({2, 3, 4})
from astrai.extension.ops.gemm import bf16_gemm
def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
return F.linear(x, weight, bias)
def _inference_bf16_gemv(
def _inference_bf16_gemm(
x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
) -> Tensor:
# Model parameters retain requires_grad=True after eval(). Dispatch is
# already restricted to no-grad, so detached views preserve storage and
# layout while satisfying the primitive's explicit autograd guard.
return bf16_gemv(
return bf16_gemm(
x.detach(),
weight.detach(),
bias.detach() if bias is not None else None,
)
def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
def _gemm_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
"""Check whether bf16_gemm can safely handle the call."""
if (
torch.is_grad_enabled()
or not x.is_cuda
@@ -66,13 +56,32 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
or weight.dtype != torch.bfloat16
or weight.ndim != 2
or x.ndim not in (1, 2)
or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
or x.shape[-1] != weight.shape[1]
or x.device != weight.device
or not x.is_contiguous()
or not weight.is_contiguous()
or torch.cuda.get_device_capability(x.get_device()) < (8, 0)
or not is_available("bf16_gemv")
or not is_available("bf16_gemm")
):
return False
m = 1 if x.ndim == 1 else x.shape[0]
# M <= 32 is where the kernel wins: L2-rotation measurements on L20
# show every production shape at M=24-32 winning or tying, while
# M=48-64 loses the long-K down_proj by 8-10% (cuBLAS switches to a
# wider tile there). The kernel itself still accepts M <= 64 when
# called directly through astrai.extension.ops.gemm.
if not (1 <= m <= 32):
return False
# Vocabulary-sized lm_head weights (N in the tens of thousands+) stream
# better through cuBLAS: our skinny path ties it at M<=8 and the tiled
# path loses ~4% at M=9-16 (L2-rotation measurements on L20). Gate the
# whole shape family out instead of splitting hairs per M band.
if weight.shape[0] > 32768:
return False
# M > 8 tiled path requires K % 8 == 0 and 16-byte alignment.
k = x.shape[-1]
if m > 8 and (
k % 8 != 0 or (x.data_ptr() & 15) != 0 or (weight.data_ptr() & 15) != 0
):
return False
return bias is None or (
@@ -87,7 +96,7 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
weight_shape = tuple(weight.shape)
m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
supported_m = m is not None and 1 <= m <= 8
supported_m = m is not None and 1 <= m <= 64
shape_matches = (
weight.ndim == 2
and x.ndim in (1, 2)
@@ -107,10 +116,10 @@ def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str,
capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None
return tensor_axes(
x,
mode=env_mode("ASTRAI_GEMV"),
mode=env_mode("ASTRAI_GEMM"),
m=m,
supported_m=supported_m,
auto_m=m in _AUTO_GEMV_M,
auto_m=supported_m,
shape_matches=shape_matches,
same_device=same_device,
weight_dtype=weight.dtype,
@@ -142,25 +151,25 @@ _SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()
def _linear_records() -> List[ImplRecord]:
mode = env_mode("ASTRAI_GEMV")
gemv_priority = 0 if mode == "1" else 100
mode = env_mode("ASTRAI_GEMM")
gemm_priority = 0 if mode == "1" else 100
auto_priority = 0 if mode == "auto" else 90
torch_priority = 0 if mode == "0" else 50
return [
ImplRecord(
family="linear",
name="gemv",
obj=_inference_bf16_gemv,
name="gemm",
obj=_inference_bf16_gemm,
spec=_SPEC_CAPABLE,
available=lambda: is_available("bf16_gemv"),
priority=gemv_priority,
available=lambda: is_available("bf16_gemm"),
priority=gemm_priority,
),
ImplRecord(
family="linear",
name="auto_gemv",
obj=_inference_bf16_gemv,
name="auto_gemm",
obj=_inference_bf16_gemm,
spec=_SPEC_AUTO,
available=lambda: is_available("bf16_gemv"),
available=lambda: is_available("bf16_gemm"),
priority=auto_priority,
),
ImplRecord(
@@ -187,23 +196,19 @@ register_family("linear", _axes, _linear_records, _fallback_record)
def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
"""Apply a linear projection with safe inference-only GEMV dispatch.
"""Apply a linear projection with safe inference-only GEMM dispatch.
``ASTRAI_GEMV=0`` always uses PyTorch, ``1`` forces GEMV whenever the
primitive can safely handle any M in ``{1, ..., 8}``, and ``auto`` (the
default) uses GEMV for decode batches with M in ``{2, 3, 4}``.
``ASTRAI_GEMM=0`` always uses PyTorch, ``1`` forces GEMM whenever the
primitive can safely handle the call (M in [1, 64], K % 8 == 0 and
16-byte-aligned for M > 8), and ``auto`` (the default) selects GEMM
for all capable decode batches.
"""
# Route through the shared dispatcher whenever a selection is active so
# explicit/context/env overrides stay honored; otherwise keep the hot
# path free of axes dictionaries and record sorting.
if get_override("linear") is not None or env_selection("linear") is not None:
return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
mode = env_mode("ASTRAI_GEMV")
if mode != "0" and _gemv_capable(x, weight, bias):
m = 1 if x.ndim == 1 else x.shape[0]
if mode == "1" or m in _AUTO_GEMV_M:
return _inference_bf16_gemv(x, weight, bias)
mode = env_mode("ASTRAI_GEMM")
if mode != "0" and _gemm_capable(x, weight, bias):
return _inference_bf16_gemm(x, weight, bias)
return _torch_linear(x, weight, bias)
astrai/extension/backend/swiglu.py
+9 -9
View File
@@ -52,16 +52,16 @@ def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
"""Apply the dense-MLP SwiGLU projection with a safe torch fallback.
``ASTRAI_SWIGLU=0`` and ``auto`` keep the unfused linear-backend chain;
``1`` forces the fused primitive for supported inputs. Auto will adopt
an M-banded rule mirroring the linear backend once end-to-end evidence
qualifies one.
``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain; ``1`` forces
the fused primitive for supported inputs; ``auto`` (the default) uses
the fused primitive for decode batches with M in ``{1, ..., 8}``. The
fused kernel reads x once and covers both projections plus the SiLU
gate-multiply in a single launch, measured 9-15% faster than the
unfused chain per MLP call on L20 with L2-thrashing weight rotation.
"""
if env_mode("ASTRAI_SWIGLU") != "1" or not _swiglu_capable(
x, up_weight, gate_weight
):
return _unfused_swiglu(x, up_weight, gate_weight)
return _fused_swiglu(x, up_weight, gate_weight)
if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
return _fused_swiglu(x, up_weight, gate_weight)
return _unfused_swiglu(x, up_weight, gate_weight)
__all__ = ["swiglu"]
astrai/extension/ops/__init__.py
+2 -2
View File
@@ -7,7 +7,7 @@ from astrai.extension.ops.attention import (
attn_paged_prefill,
attn_prefill,
)
from astrai.extension.ops.gemv import bf16_gemv
from astrai.extension.ops.gemm import bf16_gemm
from astrai.extension.ops.rotary import rotary_emb
from astrai.extension.ops.swiglu import bf16_swiglu
@@ -17,7 +17,7 @@ __all__ = [
"attn_paged_decode",
"attn_paged_prefill",
"attn_prefill",
"bf16_gemv",
"bf16_gemm",
"bf16_swiglu",
"rotary_emb",
]
astrai/extension/ops/gemm.py
+27
View File
@@ -0,0 +1,27 @@
"""Stateless wrapper for the directly callable BF16 GEMM primitive."""
from typing import Optional
import torch
from astrai.extension.loader import get_module
def bf16_gemm(
x: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Compute ``F.linear(x, weight, bias)`` for up to 64 BF16 rows.
``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 64]``, and
``weight`` must be a contiguous row-major ``[N, K]`` tensor. M in
``[1, 8]`` uses the register-resident skinny GEMM kernel (any K);
larger M uses the tiled kernel (K must be a multiple of 8 with
16-byte-aligned tensors). This primitive is inference-only and
intentionally performs no fallback or model-level dispatch.
"""
return get_module("bf16_gemm").bf16_gemm(x, weight, bias)
__all__ = ["bf16_gemm"]
astrai/extension/ops/gemv.py
-26
View File
@@ -1,26 +0,0 @@
"""Stateless wrapper for the directly callable BF16 GEMV primitive."""
from typing import Optional
import torch
from astrai.extension.loader import get_module
def bf16_gemv(
x: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Compute ``F.linear(x, weight, bias)`` for up to eight BF16 rows.
``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 8]``,
and ``weight`` must be a contiguous row-major ``[N, K]`` tensor. The CUDA
kernel reuses each weight row across M, accumulates in FP32, and returns
BF16. This primitive is inference-only and intentionally performs no
fallback or model-level dispatch.
"""
return get_module("bf16_gemv").bf16_gemv(x, weight, bias)
__all__ = ["bf16_gemv"]