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:
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user