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

fix: k_seg is already 3D from slicing, don't add extra unsqueeze(-1)

Browse Source
This commit is contained in:
biondizzle
2026-05-26 11:02:44 +00:00
parent 375a682206
commit 5b55cf0767
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/unit/test_d5b_perrow_lse.py
View File

@@ -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())