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

fix O normalization: use direct rmem tensor from partition_D shape

Browse Source
This commit is contained in:
biondizzle
2026-05-22 09:23:58 +00:00
parent 0da960d8da
commit 5e51b726ba
1 changed files with 5 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/unit/test_fmha_v3_stage_c_full.py
View File

@@ -284,9 +284,6 @@ class FmhaV3StageC:
softmax_done_bar.arrive()
# --- Normalize O in TMEM by row_sum ---
# O is at tmem_o0_offset in TMEM. Load each element, divide by row_sum, store back.
# Use the O TMEM layout (pv_thr C-fragment) for the load/store.
# We need a tiled TMEM copy for O.
cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO = pv_thr.partition_C(cO)
corr_tile_size = 16
@@ -308,13 +305,11 @@ class FmhaV3StageC:
for i in range(o_col_tiles):
tTMEM_LOAD_O_i = cute.make_tensor(tTMEM_LOAD_OtO.iterator + i * corr_tile_size, tTMEM_LOAD_OtO.layout)
tTMEM_STORE_O_i = cute.make_tensor(tTMEM_STORE_OtO.iterator + i * corr_tile_size, tTMEM_STORE_OtO.layout)
tTMrO_i_ = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype)
tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMEM_LOAD_OcO.shape[0]))
tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO_i)
for k in cutlass.range(cute.size(tTMrO_i), vectorize=True):
tTMrO_i[k] = tTMrO_i[k] * inv_row_sum
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STORE_O_i)
tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype)
cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO)
for k in cutlass.range(cute.size(tTMrO), vectorize=True):
tTMrO[k] = tTMrO[k] * inv_row_sum
cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i)
cute.arch.fence_view_async_tmem_store()
# --- Epilogue: write O from TMEM to GMEM ---