perf: accelerate decode linear with bf16 gemv

- add decode-shape benchmark harness - add bf16 GEMV CUDA primitive with head-dim generic kernel - dispatch decode-time linear layers to gemv for M=1 - extend gemv coverage to small decode batches
2026-09-02 13:11:25 +08:00
parent 9c3ef0c2a1
commit a144d7f306
14 changed files with 1242 additions and 3 deletions
@@ -0,0 +1,185 @@
+import logging
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from astrai.extension import explain, is_available, linear, op_backend
+from astrai.extension.backend import linear as public_linear
+from astrai.model.components.linear import Linear
+
+GEMV_AVAILABLE = (
+    torch.cuda.is_available()
+    and is_available("bf16_gemv")
+    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+",
+)
+
+
+def test_linear_backend_is_public():
+    assert linear is public_linear
+
+
+def test_model_linear_routes_through_backend(monkeypatch):
+    sentinel = torch.randn(2, 4)
+
+    def fake_linear(x, weight, bias):
+        assert x.shape == (2, 3)
+        assert weight.shape == (4, 3)
+        assert bias is None
+        return sentinel
+
+    monkeypatch.setattr("astrai.model.components.linear.linear", fake_linear)
+    layer = Linear(3, 4)
+    assert layer(torch.randn(2, 3)) is sentinel
+
+
+def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "1")
+    x = torch.randn(2, 8, requires_grad=True)
+    weight = torch.randn(4, 8, requires_grad=True)
+    actual = linear(x, weight)
+    expected = F.linear(x, weight)
+    torch.testing.assert_close(actual, expected)
+    actual.sum().backward()
+    assert x.grad is not None
+    assert weight.grad is not None
+
+
+def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
+    monkeypatch.setenv("ASTRAI_GEMV", "invalid-test-mode")
+    with caplog.at_level(logging.WARNING):
+        trace = explain("linear", torch.randn(1, 8), torch.randn(4, 8))
+    assert "using auto" in caplog.text
+    assert "=> torch" in trace
+
+
+@skip_no_gemv
+def test_mode_zero_disables_gemv(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "0")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert "=> torch" in explain("linear", x, weight)
+        torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
+
+
+@skip_no_gemv
+def test_mode_one_forces_capable_unmeasured_shape(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "1")
+    x = torch.randn(1, 64, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert "=> gemv" in explain("linear", x, weight)
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
+        )
+
+
+@skip_no_gemv
+@pytest.mark.parametrize("m", [2, 4, 8])
+def test_mode_one_dispatches_supported_small_batches(monkeypatch, m):
+    monkeypatch.setenv("ASTRAI_GEMV", "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 "=> gemv" in explain("linear", x, weight)
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
+        )
+
+
+@skip_no_gemv
+def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "auto")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    winning = torch.randn(
+        1536,
+        1536,
+        device="cuda",
+        dtype=torch.bfloat16,
+        requires_grad=True,
+    )
+    with torch.no_grad():
+        assert "=> torch" in explain("linear", x, winning)
+        torch.testing.assert_close(linear(x, winning), F.linear(x, winning))
+
+
+@skip_no_gemv
+def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "auto")
+    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        trace = explain("linear", x, weight)
+        if torch.cuda.get_device_capability() == (8, 9):
+            assert "=> auto_gemv" in trace
+        else:
+            assert "=> torch" in trace
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.5
+        )
+
+
+@skip_no_gemv
+@pytest.mark.parametrize(
+    "m,n,k",
+    [
+        (2, 6912, 1536),  # up/gate loses at every measured M
+        (2, 1536, 6912),  # long-K accumulation changed checkpoint greedy output
+        (4, 100000, 1536),  # LM head misses the 5% M=4 gate
+        (4, 1536, 6912),  # long-K accumulation changed checkpoint greedy output
+        (8, 256, 1536),  # remaining M=8 winners miss the 3% end-to-end gate
+    ],
+)
+def test_auto_rejects_measured_small_batch_losers(monkeypatch, m, n, k):
+    monkeypatch.setenv("ASTRAI_GEMV", "auto")
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert "=> torch" in explain("linear", x, weight)
+
+
+@skip_no_gemv
+def test_grad_enabled_and_unsupported_multirow_always_fall_back(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "1")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(
+        256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
+    )
+    assert "=> torch" in explain("linear", x, weight)
+    with torch.no_grad():
+        multirow = x.expand(3, -1).contiguous()
+        assert "=> torch" in explain("linear", multirow, weight)
+
+
+@skip_no_gemv
+def test_explicit_gemv_selection_respects_capability(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "0")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad(), op_backend(linear="gemv"):
+        assert "=> gemv" in explain("linear", x, weight)
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
+        )
+
+
+@skip_no_gemv
+def test_dispatched_linear_cuda_graph_replay(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "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():
+        for _ in range(3):
+            linear(x, weight)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = linear(x, weight)
+        x.copy_(torch.randn_like(x))
+        graph.replay()
+        expected = F.linear(x, weight)
+    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)