From e701a1411ce8c4000641da54ee0790559262066d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 10:47:00 +0000 Subject: [PATCH] P5: use multi-tile kernel for N>128 in integration test --- tests/unit/test_p3_fast_decode.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py index 1a1cda5b..0df5ca3a 100644 --- a/tests/unit/test_p3_fast_decode.py +++ b/tests/unit/test_p3_fast_decode.py @@ -85,8 +85,16 @@ def test_kernel_correctness(): k_4d = torch.randn(1, n_kv, N, hd, dtype=torch.bfloat16, device='cuda').contiguous() v_4d = torch.randn(1, n_kv, hd, N, dtype=torch.bfloat16, device='cuda').contiguous() - sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda') - o_4d, _ = fmha_multihead_decode_raw(q_4d, k_4d, v_4d, scale, 0, 0, False, sb) + # Use the correct kernel for the KV size + if N > 128 or hd == 512: + from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw as kernel_fn + # Multi-tile kernel needs T in the Q shape + q_4d_mt = torch.randn(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous() + o_4d, _ = kernel_fn(q_4d_mt, k_4d, v_4d, scale, 0, 0, False, + torch.zeros(1, n_q, dtype=torch.float32, device='cuda')) + else: + sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda') + o_4d, _ = fmha_multihead_decode_raw(q_4d, k_4d, v_4d, scale, 0, 0, False, sb) o_ref = reference_attention(q_4d, k_4d, v_4d, scale)