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

fix: partition_A not partition_S

Browse Source
This commit is contained in:
biondizzle
2026-05-23 03:47:53 +00:00
parent 3ee330a84c
commit 469665f69a
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/fmha.py
View File

@@ -200,9 +200,9 @@ class FmhaKernel:
tScS = qk_thr.partition_C(cS)
tTMEM_LOADcS = thr_load.partition_D(tScS)
# P → SMEM copy (using PV A-operand thread partition)
# P → SMEM: use PV A-operand partition for SMEM write
p_s = cute.slice_(p_smem_s,(None,None,None,0))
tCrP_smem = pv_thr.partition_S(sP) # softmax thread → SMEM partition for P
tCrP_smem = pv_thr.partition_A(sP) # PV thread → SMEM partition for P (A operand)
tCrP_reg = cute.make_rmem_tensor(tCrP_smem.shape, self.q_dtype)
# Online softmax state