refactor: accept arbitrary K in bf16 gemv with aligned head-tail sweeps

- Drop the K % 2 entry rejection and the per-K if/else load-width branch: the weight stream now anchors uint4 loads at each row's first 16-byte-aligned address, with scalar head/tail sweeps covering at most 14 remainder elements, so any positive K and any storage offset is correct
- Keep one pure-uint4 loop (no branching inside the loop) for the production case where every x row base is 16-byte aligned (K % 8 == 0 with allocator-aligned tensors) and a scalar-x pairing loop only for unaligned K, where per-row uint4 loads are not addressable; measured cost of scalar x everywhere was up to 2.5x on multi-row shapes (down M=4 28.4us vs 11.3us)
- Remove the now-obsolete k_aligned axis and K divisibility gate from the linear dispatch spec since the primitive no longer rejects any K
- Add test coverage for unaligned K (7, 12, 100, 1534) at M=1 and M=3

Benchmark: 8x L20 (sm_89, CUDA 12.8), L2-resident microbench, 300 iters; hot path unchanged within noise vs the pure-uint4 kernel (q M=2 5.8us, down M=4 11.3us, lm M=1 391us); full gate green
This commit is contained in:
2026-09-02 14:48:01 +08:00
committed by 0z5a
parent 1c3515714f
commit 800981d85a
4 changed files with 90 additions and 48 deletions
+15 -7
View File
@@ -96,6 +96,21 @@ def test_bf16_gemv_cuda_graph_replay():
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv
@pytest.mark.parametrize("n,k", [(64, 7), (64, 12), (33, 100), (256, 1534)])
def test_bf16_gemv_handles_unaligned_k(n, k):
torch.manual_seed(29)
x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemv(x, weight)
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
x3 = torch.randn(3, k, device="cuda", dtype=torch.bfloat16)
actual3 = bf16_gemv(x3, weight)
torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5)
@skip_no_gemv
def test_bf16_gemv_small_batch_cuda_graph_replay():
torch.manual_seed(31)
@@ -125,13 +140,6 @@ def test_bf16_gemv_small_batch_cuda_graph_replay():
),
"M must",
),
(
lambda: (
torch.randn(15, device="cuda", dtype=torch.bfloat16),
torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
),
"even",
),
(
lambda: (
torch.randn(16, device="cuda", dtype=torch.float16),