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

D1.3: Revert to zero-fill for sP - need to verify sP→PV pipeline first

Browse Source
This commit is contained in:
biondizzle
2026-05-23 23:26:07 +00:00
parent 5be5d42e94
commit a17dca508d
1 changed files with 3 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -351,33 +351,9 @@ class FmhaKernel:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: write P to sP using direct (m,k) computation.
# Instead of looking up coordinates from identity tensor,
# compute (m, k) from the TMEM load's known thread mapping.
# The Ld32x32bOp(Repetition(32)) distributes elements as:
# Each warp (32 threads) handles 32 rows of the S matrix.
# 4 softmax warps (warps 0-3) handle all 128 rows.
# Within a warp, each thread handles 1 row and 4 column groups.
# Thread sfw_idx: warp_id = sfw_idx // 32, lane = sfw_idx % 32
# Row: lane + 32 * warp_id
# Columns: 4 fragments of 32 elements each.
# Fragment j1: columns [j1*32, (j1+1)*32)
# Within fragment: column = j0 + j1 * 32
_warp_id = sfw_idx // 32
_lane = sfw_idx % 32
_row = _lane + 32 * _warp_id
for j0 in range(32):
for j1 in range(4):
k_coord = j0 + 32 * j1 # column index
m_coord = _row # row index
# Compute sP sub-coordinates
k0 = k_coord % 16
k1 = (k_coord // 16) % 4
k2 = k_coord // 64
# rP_bf16 uses the same indexing as tTMEM_LOADrS
# which has the same layout as tTMEM_LOADcS.
# rP_bf16[(j0, 0), j1, 0, 0] should give P at (m, k).
_sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0]
# SMEM-P: zero-fill sP for now (testing sP→PV pipeline)
for j in cutlass.range(cute.size(sP), vectorize=True):
sP[j] = self.q_dtype(0)
cute.arch.fence_proxy("async.shared", space="cta")
if kt > 0:
tTMrO = cute.make_rmem_tensor(