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

P8: Fix test imports after deleting multihead module

Browse Source
This commit is contained in:
biondizzle
2026-05-30 17:23:13 +00:00
parent 95725f1df0
commit e1b9e94c24
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_p3_fast_decode.py
View File

@@ -13,7 +13,7 @@ import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
def cosine_sim(a, b):
@@ -91,7 +91,7 @@ def test_kernel_correctness():
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)
o_4d, _ = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)
o_ref = reference_attention(q_4d, k_4d, v_4d, scale)