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

fix: mRowSums dummy tensor must match mLSE layout (3D, not 1D)

Browse Source
This commit is contained in:
biondizzle
2026-05-26 15:14:35 +00:00
parent ffc4b542bc
commit 8f7df4d8b5
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
dsv4/kernels/attention/fmha.py
View File

@@ -153,7 +153,7 @@ class FmhaKernel:
# Grid: (M_tiles, 1, batch) where M = n_h * T packed into M dimension
# For single-head (n_h=1): grid=(1,1,1) — backward compatible
if const_expr(row_sums is None):
row_sums = cute.make_tensor(lse.iterator, cute.make_layout((1,), stride=(0,)))
row_sums = cute.make_tensor(lse.iterator, lse.layout)
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_len,sink_bias,row_sums).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream)