diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py
index e1115ac7..e198fcfb 100644
--- a/tests/unit/test_p3_fast_decode.py
+++ b/tests/unit/test_p3_fast_decode.py
@@ -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)