From 95e0c8c464728fee05b10c10c7af0f25fe66b9de Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 10:49:12 +0000 Subject: [PATCH] =?UTF-8?q?P5:=20fix=20multi-tile=20test=20=E2=80=94=20use?= =?UTF-8?q?=20same=20Q=20data=20for=20kernel=20and=20reference?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_p3_fast_decode.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py index 0df5ca3a..e1115ac7 100644 --- a/tests/unit/test_p3_fast_decode.py +++ b/tests/unit/test_p3_fast_decode.py @@ -88,10 +88,7 @@ def test_kernel_correctness(): # 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')) + o_4d, _ = kernel_fn(q_4d, k_4d, v_4d, scale) 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)