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
tests/extension/test_linear_dispatch.py
+61 -49
View File
@@ -12,26 +12,26 @@ from astrai.extension.dispatch import explain, op_backend, resolve
# module object explicitly for monkeypatching its private helpers.
linear_module = importlib.import_module("astrai.extension.backend.linear")
GEMV_AVAILABLE = (
GEMM_AVAILABLE = (
torch.cuda.is_available()
and is_available("bf16_gemv")
and is_available("bf16_gemm")
and torch.cuda.get_device_capability() >= (8, 0)
)
skip_no_gemv = pytest.mark.skipif(
not GEMV_AVAILABLE,
reason="BF16 GEMV requires a built kernel and compute capability 8.0+",
skip_no_gemm = pytest.mark.skipif(
not GEMM_AVAILABLE,
reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
)
def _routes_to_gemv(monkeypatch, x, weight, bias=None) -> bool:
"""Patch the GEMV entry point to a sentinel and report whether
def _routes_to_gemm(monkeypatch, x, weight, bias=None) -> bool:
"""Patch the GEMM entry point to a sentinel and report whether
``linear`` selected it (torch fallback would compute a real tensor)."""
sentinel = object()
def fake_gemv(x, weight, bias):
def fake_gemm(x, weight, bias):
return sentinel
monkeypatch.setattr(linear_module, "_inference_bf16_gemv", fake_gemv)
monkeypatch.setattr(linear_module, "_inference_bf16_gemm", fake_gemm)
return linear(x, weight, bias) is sentinel
@@ -52,7 +52,7 @@ def test_model_linear_routes_through_backend(monkeypatch):
def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
monkeypatch.setenv("ASTRAI_GEMV", "invalid-test-mode")
monkeypatch.setenv("ASTRAI_GEMM", "invalid-test-mode")
x = torch.randn(2, 8)
weight = torch.randn(4, 8)
with caplog.at_level(logging.WARNING):
@@ -62,7 +62,7 @@ def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1")
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(2, 8, requires_grad=True)
weight = torch.randn(4, 8, requires_grad=True)
actual = linear(x, weight)
@@ -73,71 +73,83 @@ def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
assert weight.grad is not None
@skip_no_gemv
@pytest.mark.parametrize("m", [2, 3, 4])
@skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 3, 4, 5, 6, 7, 8])
def test_auto_selects_small_decode_batches(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMV", "auto")
monkeypatch.setenv("ASTRAI_GEMM", "auto")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight)
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemv
@pytest.mark.parametrize("m", [1, 5, 8, 9])
@skip_no_gemm
@pytest.mark.parametrize("m", [12, 16, 24, 32])
def test_auto_selects_larger_decode_batches(monkeypatch, m):
"""Auto covers M up to 32; 48+ loses to cuBLAS on long-K shapes."""
monkeypatch.setenv("ASTRAI_GEMM", "auto")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemm
@pytest.mark.parametrize("m", [48, 64, 65])
def test_auto_falls_back_outside_band(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMV", "auto")
"""M beyond 32 falls back to cuBLAS (measured regression at M=48+)."""
monkeypatch.setenv("ASTRAI_GEMM", "auto")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert not _routes_to_gemv(monkeypatch, x, weight)
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
)
@skip_no_gemv
def test_mode_zero_disables_gemv(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "0")
@skip_no_gemm
def test_mode_zero_disables_gemm(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "0")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert not _routes_to_gemv(monkeypatch, x, weight)
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
@skip_no_gemv
@pytest.mark.parametrize("m", [1, 2, 8])
@skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 8, 16, 32])
def test_mode_one_forces_every_capable_batch(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMV", "1")
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight)
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemv
@skip_no_gemm
def test_mode_one_rejects_oversized_batch_and_grad(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1")
monkeypatch.setenv("ASTRAI_GEMM", "1")
weight = torch.randn(
256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
)
with torch.no_grad():
oversized = torch.randn(9, 1536, device="cuda", dtype=torch.bfloat16)
assert not _routes_to_gemv(monkeypatch, oversized, weight)
assert not _routes_to_gemv(
oversized = torch.randn(65, 1536, device="cuda", dtype=torch.bfloat16)
assert not _routes_to_gemm(monkeypatch, oversized, weight)
assert not _routes_to_gemm(
monkeypatch, torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16), weight
)
@skip_no_gemv
@skip_no_gemm
def test_mode_one_supports_bias_and_vector_input(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1")
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight, bias)
assert _routes_to_gemm(monkeypatch, x, weight, bias)
monkeypatch.undo()
torch.testing.assert_close(
linear(x, weight, bias),
@@ -147,9 +159,9 @@ def test_mode_one_supports_bias_and_vector_input(monkeypatch):
)
@skip_no_gemv
@skip_no_gemm
def test_dispatched_linear_cuda_graph_replay(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1")
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
@@ -176,44 +188,44 @@ def test_linear_family_is_registered_with_shared_dispatcher():
assert "linear" in explain("linear", x, weight)
@skip_no_gemv
@skip_no_gemm
def test_ops_env_override_forces_torch_for_capable_call(monkeypatch):
"""ASTR_OPS=linear=torch must keep working after the M-band rewrite
(regression: the family was silently dropped from the dispatcher, so
the override warned, fell through, and the gemv kernel still ran)."""
the override warned, fell through, and the gemm kernel still ran)."""
monkeypatch.setenv("ASTR_OPS", "linear=torch")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert not _routes_to_gemv(monkeypatch, x, weight)
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
)
@skip_no_gemv
def test_ops_env_override_forces_gemv(monkeypatch):
monkeypatch.setenv("ASTR_OPS", "linear=gemv")
@skip_no_gemm
def test_ops_env_override_forces_gemm(monkeypatch):
monkeypatch.setenv("ASTR_OPS", "linear=gemm")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
# M=1 is outside the auto band but inside the forced gemv record.
assert _routes_to_gemv(monkeypatch, x, weight)
# M=1 is outside the auto band but inside the forced gemm record.
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemv
@skip_no_gemm
def test_op_backend_context_selects_torch(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1")
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(), op_backend(linear="torch"):
assert not _routes_to_gemv(monkeypatch, x, weight)
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
)
# The override is scoped: the forced mode applies again afterwards.
with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight)
assert _routes_to_gemm(monkeypatch, x, weight)
def test_op_backend_rejects_unknown_linear_handle():