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

fix: make K/V segments contiguous before passing to kernel (TMA needs contiguous tensors)

Browse Source
This commit is contained in:
biondizzle
2026-05-26 11:00:36 +00:00
parent 5407dc768a
commit 2252d7c865
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

@@ -167,8 +167,8 @@ def test_lse_kv_merge():
lses = []
for seg in range(s_k // seg_size):
k_seg = k[seg * seg_size:(seg + 1) * seg_size]
v_seg = v[seg * seg_size:(seg + 1) * 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)
o_seg, lse_seg = _run_fmha_with_lse(q, k_seg_3d, v_seg, m, seg_size, hd)
@@ -208,8 +208,8 @@ def test_lse_kv_merge_4tiles():
lses = []
for seg in range(s_k // seg_size):
k_seg = k[seg * seg_size:(seg + 1) * seg_size]
v_seg = v[seg * seg_size:(seg + 1) * 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)
o_seg, lse_seg = _run_fmha_with_lse(q, k_seg_3d, v_seg, m, seg_size, hd)