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

Guard SMEM accumulation with n_kv_tiles>1 to avoid TMEM destructive read

Browse Source
This commit is contained in:
biondizzle
2026-05-27 05:02:51 +00:00
parent 02a34512cb
commit 101840c78c
1 changed files with 20 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
dsv4/kernels/attention/fmha_smem_acc.py
View File

@@ -458,25 +458,28 @@ class FmhaKernel:
# ========================================================
# SMEM ACCUMULATOR: load O_kt from TMEM, accumulate in SMEM
# ========================================================
# O_kt is in TMEM (PV with ACCUMULATE=False → fresh output)
# Load via one-way Ld32x32bOp (read-only, NO write-back to TMEM)
# Then: sO_acc = acc_scale * sO_acc + O_kt
# Using coordinate-indexed writes to sO_acc
# Only needed for n_kv_tiles > 1.
# For n_kv_tiles=1, epilogue_tma_store reads directly from TMEM.
#
# WARNING: Ld32x32bOp may be a destructive TMEM read!
# If so, the O load here would consume TMEM data,
# making epilogue_tma_store read garbage.
# Guard with const_expr to avoid this for n_kv_tiles=1.
# ========================================================
rO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype)
cute.copy(tiled_o_load, tTMEM_LOADtO, rO)
cute.arch.fence_view_async_tmem_load()
if const_expr(self.n_kv_tiles > 1):
rO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype)
cute.copy(tiled_o_load, tTMEM_LOADtO, rO)
cute.arch.fence_view_async_tmem_load()
# Rescale existing sO_acc and add O_kt
# Use coordinate tensor to map each register to (row, col) in sO_acc
for j0 in range(32):
for j1 in range(4):
coord = tTMEM_LOADcO[(j0, 0), j1, 0, 0]
row = coord[0]
col = coord[1]
old_val = sO_acc[row, col]
new_val = acc_scale * old_val + rO[(j0, 0), j1, 0, 0]
sO_acc[row, col] = new_val
# Rescale existing sO_acc and add O_kt
for j0 in range(32):
for j1 in range(4):
coord = tTMEM_LOADcO[(j0, 0), j1, 0, 0]
row = coord[0]
col = coord[1]
old_val = sO_acc[row, col]
new_val = acc_scale * old_val + rO[(j0, 0), j1, 0, 0]
sO_acc[row, col] = new_val
# Wait for MMA's final signal
final_o_bar.arrive_and_wait()