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
csrc/bench/benchmark_gemv.py → csrc/bench/benchmark_gemm.py
+2 -2
View File
@@ -1,7 +1,7 @@
"""Benchmark decode-time linear shapes before enabling custom GEMV dispatch.
"""Benchmark decode-time linear shapes before enabling custom GEMM dispatch.
The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
establishes the per-architecture cuBLAS baseline that later GEMV primitives and
establishes the per-architecture cuBLAS baseline that later GEMM primitives and
dispatch decisions must beat.
"""
csrc/bench/benchmark_gemv_common.py → csrc/bench/benchmark_gemm_common.py
+6 -6
View File
@@ -1,4 +1,4 @@
"""Benchmark the BF16 GEMV primitive and guarded linear dispatcher.
"""Benchmark the BF16 GEMM primitive and guarded linear dispatcher.
The kernel suite covers AstrAI's native projections plus common LLaMA and
GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain;
@@ -19,7 +19,7 @@ from pathlib import Path
import torch
import torch.nn.functional as F
from astrai.extension import bf16_gemv, is_available, linear
from astrai.extension import bf16_gemm, is_available, linear
@dataclass(frozen=True)
@@ -198,7 +198,7 @@ def _kernel_functions(
return F.linear(x, weight)
def candidate() -> torch.Tensor:
return bf16_gemv(x, weight.detach())
return bf16_gemm(x, weight.detach())
return baseline, candidate
@@ -256,7 +256,7 @@ def benchmark_kernels(
def _set_mode(mode: str) -> None:
os.environ["ASTRAI_GEMV"] = mode
os.environ["ASTRAI_GEMM"] = mode
def _chain_weights(
@@ -398,8 +398,8 @@ def parse_args() -> argparse.Namespace:
def main() -> None:
args = parse_args()
if not torch.cuda.is_available() or not is_available("bf16_gemv"):
raise RuntimeError("benchmark requires CUDA and the built bf16_gemv extension")
if not torch.cuda.is_available() or not is_available("bf16_gemm"):
raise RuntimeError("benchmark requires CUDA and the built bf16_gemm extension")
if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1:
raise ValueError("warmup must be non-negative and sample/inner counts positive")
csrc/bench/benchmark_swiglu.py
+10 -10
View File
@@ -1,4 +1,4 @@
"""Benchmark fused BF16 SwiGLU against torch and unfused GEMV chains."""
"""Benchmark fused BF16 SwiGLU against torch and unfused GEMM chains."""
from __future__ import annotations
@@ -14,7 +14,7 @@ import click
import torch
import torch.nn.functional as F
from astrai.extension import bf16_gemv, bf16_swiglu, is_available
from astrai.extension import bf16_gemm, bf16_swiglu, is_available
@dataclass(frozen=True)
@@ -124,8 +124,8 @@ def capture(operation: Callable[[], torch.Tensor]):
def make_operations(x, up_weight, gate_weight, mode: str):
operations: dict[str, Callable[[], torch.Tensor]] = {
"torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)),
"gemv_chain": lambda: (
bf16_gemv(x, up_weight) * F.silu(bf16_gemv(x, gate_weight))
"gemm_chain": lambda: (
bf16_gemm(x, up_weight) * F.silu(bf16_gemm(x, gate_weight))
),
"fused": lambda: bf16_swiglu(x, up_weight, gate_weight),
}
@@ -226,19 +226,19 @@ def render_markdown(payload: dict[str, object]) -> str:
f"- Compute capability: {metadata['compute_capability']}",
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
"",
"| Shape | M | Mode | torch ms | GEMV chain ms | fused ms | "
"| Shape | M | Mode | torch ms | GEMM chain ms | fused ms | "
"vs best unfused | fused kernels | max abs | cosine |",
"|---|---:|---|---:|---:|---:|---:|---:|---:|---:|",
]
for shape, m, mode in cases:
torch_item = by_case[(shape, m, mode, "torch")]
gemv_item = by_case[(shape, m, mode, "gemv_chain")]
gemm_item = by_case[(shape, m, mode, "gemm_chain")]
fused_item = by_case[(shape, m, mode, "fused")]
best = min(torch_item["median_ms"], gemv_item["median_ms"])
best = min(torch_item["median_ms"], gemm_item["median_ms"])
improvement = (best / fused_item["median_ms"] - 1) * 100
lines.append(
f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | "
f"{gemv_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
f"{gemm_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
f"{improvement:+.2f}% | "
f"{fused_item['cuda_kernel_launches_per_call']:.1f} | "
f"{fused_item['max_abs_error']:.5f} | "
@@ -273,8 +273,8 @@ def benchmark_command(
) -> None:
if not torch.cuda.is_available():
raise click.ClickException("CUDA is required")
if not is_available("bf16_gemv") or not is_available("bf16_swiglu"):
raise click.ClickException("built bf16_gemv and bf16_swiglu are required")
if not is_available("bf16_gemm") or not is_available("bf16_swiglu"):
raise click.ClickException("built bf16_gemm and bf16_swiglu are required")
shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES
m_values_parsed = parse_positive_ints(m_values)
if any(m > 8 for m in m_values_parsed):