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

Fix rP scope issue: use rP_bf16.iterator instead of rP.iterator

Browse Source
This commit is contained in:
biondizzle
2026-05-23 09:36:22 +00:00
parent 0b09e7e4a2
commit 0bee39d2d5
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

@@ -346,7 +346,7 @@ class FmhaKernel:
# rP_bf16 uses tTMEM_LOADrS.layout (TMEM layout) causing rank mismatch
# Create view with QK C-fragment layout (tStS0.layout)
rP_qk_layout = tStS0.layout # QK C-fragment layout for this thread
rP_qk = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=self.q_dtype), rP_qk_layout)
rP_qk = cute.make_tensor(cute.recast_ptr(rP_bf16.iterator, dtype=self.q_dtype), rP_qk_layout)
# Partition source with QK layout
tSMEM_CPYrP_qk = thr_smem_copy.partition_S(rP_qk)