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

D1.3: Use const_expr for lse None check

Browse Source
This commit is contained in:
biondizzle
2026-05-23 22:30:55 +00:00
parent 54f1d0d669
commit bfacef28d3
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -113,8 +113,10 @@ class FmhaKernel:
tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape)
epi_s = cute.select(self.c_smem_s,mode=[0,1])
tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
# Always create a valid mLSE tensor (even for normalize=True, it's just unused)
if lse is None:
# Always create a valid mLSE tensor for the kernel.
# CuTeDSL doesn't support None parameters in @cute.kernel.
# For normalize=True, mLSE is unused (dead-code-eliminated by compiler).
if const_expr(lse is None):
lse = cute.make_tensor(c.iterator, cute.make_layout((1, 1, 1), stride=(1, 1, 1)))
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).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream)