From 2c1860929610341d1f8b772fcc508b1d5429056a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 17:25:01 +0000 Subject: [PATCH] P8: Fix P6 test imports after deleting multihead module --- tests/unit/test_p6_tma_epilogue.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_p6_tma_epilogue.py b/tests/unit/test_p6_tma_epilogue.py index 9c74d4db..876bbd7e 100644 --- a/tests/unit/test_p6_tma_epilogue.py +++ b/tests/unit/test_p6_tma_epilogue.py @@ -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)