From 5b55cf07676958d8064700f4a402db07c55c8da6 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 26 May 2026 11:02:44 +0000 Subject: [PATCH] fix: k_seg is already 3D from slicing, don't add extra unsqueeze(-1) --- tests/unit/test_d5b_perrow_lse.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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())