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

fix: use logical_divide (not composition) for O rescale/normalize atoms to match get_tmem_load_op layout

Browse Source
This commit is contained in:
biondizzle
2026-05-23 02:53:59 +00:00
parent 71934f78b2
commit de7cc61480
1 changed files with 9 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -235,12 +235,11 @@ class FmhaV3StageCMulti:
scale_log2 = Float32(self.scale_softmax_log2)
# O rescale atoms (hand-constructed, for per-tile O *= acc_scale)
# Use logical_divide (not composition) to match get_tmem_load_op layout
corr_tile_size = 16
tOcO = pv_thr.partition_C(cS)
tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size)))
tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size)))
tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout)
tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout)
tOtO_epi_r = cute.logical_divide(tOtO0, cute.make_layout((128, corr_tile_size)))
tOcO_epi_r = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size)))
tmem_load_o_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
self.acc_dtype,
@@ -249,13 +248,13 @@ class FmhaV3StageCMulti:
tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
self.acc_dtype,
)
tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i)
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i)
tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_epi_r[(None, None), 0])
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_epi_r[(None, None), 0])
thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i)
tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i)
tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i)
tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_epi_r[(None, None), None])
tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_epi_r[(None, None), None])
tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_epi_r[(None, None), None])
n_corr_tiles = HEAD_DIM // corr_tile_size
for kt in range(self.n_kv_tiles):
@@ -299,7 +298,7 @@ class FmhaV3StageCMulti:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
# Per-tile O rescale (hand-constructed atoms, only for kt > 0)
# Per-tile O rescale (hand-constructed atoms with logical_divide layout)
if kt > 0:
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
@@ -331,10 +330,6 @@ class FmhaV3StageCMulti:
final_o_bar.arrive_and_wait()
# === NO-OP TMEM round-trip: re-map O from MMA layout to epilog layout ===
# The MMA writes O in the C-fragment TMEM layout, but epilogue_tma_store
# reads using get_tmem_load_op which expects a different layout. A NO-OP
# load-then-store through the hand-constructed atoms re-maps the data.
# TODO: eliminate this by using get_tmem_load_op for the normalize.
tTMrO_noop = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
)
@@ -357,11 +352,6 @@ class FmhaV3StageCMulti:
cute.arch.fence_view_async_tmem_store()
# === Final O normalization: O *= 1/row_sum ===
# TMEM round-trip using hand-constructed atoms.
# Known issue: hand-constructed Ld32x32bOp/St32x32bOp atoms introduce
# ~3% error due to TMEM column mapping mismatch with get_tmem_load_op.
# TODO: replace with correction_epilog (paired atoms) once TMA store
# region isolation is resolved.
inv_row_sum = Float32(1.0) / row_sum
tTMrO = cute.make_rmem_tensor(