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

P8: Fix P6 test imports after deleting multihead module

Browse Source
This commit is contained in:
biondizzle
2026-05-30 17:25:01 +00:00
parent e1b9e94c24
commit 2c18609296
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/unit/test_p6_tma_epilogue.py
View File

@@ -48,7 +48,7 @@ def reference_attention(q_4d, k_4d, v_4d, scale):
def test_one_way_epilogue():
"""Test the one-way TMEM→regs→SMEM→GMEM epilogue path."""
from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
torch.manual_seed(42)
@@ -77,7 +77,7 @@ def test_one_way_epilogue():
o_4d, _ = fmha_multitile_decode_raw(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)
o_4d, _ = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)
o_ref = reference_attention(q_4d, k_4d, v_4d, scale)