diff --git a/tests/unit/test_d5b_perrow_lse.py b/tests/unit/test_d5b_perrow_lse.py
index 886943dc..2735328e 100644
--- a/tests/unit/test_d5b_perrow_lse.py
+++ b/tests/unit/test_d5b_perrow_lse.py
@@ -169,9 +169,9 @@ def test_lse_kv_merge():
     for seg in range(s_k // seg_size):
         k_seg = k[seg * seg_size:(seg + 1) * seg_size].contiguous()
         v_seg = v[seg * seg_size:(seg + 1) * seg_size].contiguous()
-        k_seg_3d = k_seg.unsqueeze(-1)
+        # k_seg is already 3D from slicing (s_k, hd, 1) - no unsqueeze needed
 
-        o_seg, lse_seg = _run_fmha_with_lse(q, k_seg_3d, v_seg, m, seg_size, hd)
+        o_seg, lse_seg = _run_fmha_with_lse(q, k_seg, v_seg, m, seg_size, hd)
         o_norms.append(o_seg.float())
         lses.append(lse_seg.float())
 
@@ -210,9 +210,9 @@ def test_lse_kv_merge_4tiles():
     for seg in range(s_k // seg_size):
         k_seg = k[seg * seg_size:(seg + 1) * seg_size].contiguous()
         v_seg = v[seg * seg_size:(seg + 1) * seg_size].contiguous()
-        k_seg_3d = k_seg.unsqueeze(-1)
+        # k_seg is already 3D from slicing (s_k, hd, 1) - no unsqueeze needed
 
-        o_seg, lse_seg = _run_fmha_with_lse(q, k_seg_3d, v_seg, m, seg_size, hd)
+        o_seg, lse_seg = _run_fmha_with_lse(q, k_seg, v_seg, m, seg_size, hd)
         o_norms.append(o_seg.float())
         lses.append(lse_seg.float())