biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P5: use multi-tile kernel for N>128 in integration test

Browse Source
This commit is contained in:
biondizzle
2026-05-30 10:47:00 +00:00
parent 5932e928a8
commit e701a1411c
1 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/unit/test_p3_fast_decode.py
View File

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