P5: fix multi-tile test — use same Q data for kernel and reference

This commit is contained in:
2026-05-30 10:49:12 +00:00
parent e701a1411c
commit 95e0c8c464

View File

@@ -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)