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

Match working test: don't pass row_sums to kernel

Browse Source
This commit is contained in:
biondizzle
2026-05-27 06:41:44 +00:00
parent 312ac52d15
commit 6111db571c
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
dsv4/kernels/attention/production.py
View File

@@ -63,7 +63,7 @@ def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))
mRS = ct.from_dlpack(row_sums).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums))
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE)
_kernel_cache[key] = (compiled, kernel)
return (compiled, kernel)
@@ -204,7 +204,7 @@ def _attention_single_head(
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
compiled(mQ, mK, mV, mC, stream, lse=mLSE)
torch.cuda.synchronize()
seg_o[:, v_start:v_end] = c_tile[:, :, 0].float()