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

fix: correction_epilog with paired atoms + pre-partitioned TMA store

Browse Source
This commit is contained in:
biondizzle
2026-05-23 02:34:33 +00:00
parent fb4888079a
commit 22f8870fbd
1 changed files with 58 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

116
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -170,6 +170,17 @@ class FmhaV3StageCMulti:
gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None))
gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None))
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None))
tCgC_epi = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None))
# Pre-partition TMA store tensors (needed for correction_epilog)
epi_s = cute.select(self.c_smem_s, mode=[0, 1])
tCgC_epi_flat = cute.flat_divide(tCgC_epi, epi_tile)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_c, 0, cute.make_layout(1),
cute.group_modes(sC, 0, 2),
cute.group_modes(tCgC_epi_flat, 0, 2),
)
n_kv_tiles = cute.size(gK, mode=[3])
qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0)
@@ -407,69 +418,58 @@ class FmhaV3StageCMulti:
# Wait for MMA's PV[N-1] to commit before reading O.
final_o_bar.arrive_and_wait()
# DIAG: NO-OP TMEM round-trip before normalize
tTMrO_noop = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
)
for i in range(n_corr_tiles):
tTMrO_i_ = tTMrO_noop[None, i]
tTMrO_i_layout = cute.composition(
tTMrO_i_.layout, cute.make_layout(tTMrO_noop.shape[0])
)
tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
tTMEM_LOADtO_i = cute.make_tensor(
tTMEM_LOADtO.iterator + i * corr_tile_size,
tTMEM_LOADtO.layout,
)
tTMEM_STOREtO_i = cute.make_tensor(
tTMEM_STOREtO.iterator + i * corr_tile_size,
tTMEM_STOREtO.layout,
)
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i)
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
cute.arch.fence_view_async_tmem_store()
# === Final O normalization: O *= 1/row_sum ===
# === Correction epilog: one-way TMEM → reg → SMEM with normalize ===
# Uses get_tmem_load_op + get_smem_store_op paired atoms.
inv_row_sum = Float32(1.0) / row_sum
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
epi_corr_tile_size = 32 * 8 // self.o_dtype.width # 16 for BF16
tOtO_epi = cute.logical_divide(tOtO0, cute.make_layout((128, epi_corr_tile_size)))
tmem_load_epi_atom = utils.sm100.get_tmem_load_op(
self.pv_mma_tiler, self.c_layout, self.o_dtype, self.acc_dtype,
(epi_tile[0], epi_corr_tile_size), self.use_2cta_instrs,
)
for i in range(n_corr_tiles):
tTMrO_i_ = tTMrO[None, i]
tTMrO_i_layout = cute.composition(
tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])
)
tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
tTMEM_LOADtO_i = cute.make_tensor(
tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout
)
tTMEM_STOREtO_i = cute.make_tensor(
tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout
)
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i)
for j in cutlass.range(cute.size(tTMrO_i), vectorize=True):
tTMrO_i[j] = tTMrO_i[j] * inv_row_sum
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
cute.arch.fence_view_async_tmem_store()
# Standard epilogue: TMEM → SMEM → GMEM via TMA store.
# O in TMEM is now scaled by 1/row_sum.
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
acc_cons_st = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
tiled_tmem_load_epi = tcgen05.make_tmem_copy(tmem_load_epi_atom, tOtO_epi[(None, None), 0])
smem_store_epi_atom = utils.sm100.get_smem_store_op(
self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load_epi,
)
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile,
0, const_expr(lambda x: x), (0, 0, 0),
acc_cons_st, acc_pipe, c_pipe,
tiled_smem_store_epi = cute.make_tiled_copy_D(smem_store_epi_atom, tiled_tmem_load_epi)
tOsO = pv_thr.partition_C(sC)
cO_epi = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO_epi = pv_thr.partition_C(cO_epi)
tOsO_epi = cute.logical_divide(tOsO, cute.make_layout((128, epi_corr_tile_size)))
tOcO_epi = cute.logical_divide(tOcO_epi, cute.make_layout((128, epi_corr_tile_size)))
thr_tmem_load_epi = tiled_tmem_load_epi.get_slice(sfw_idx)
tTMEM_LOADtO_epi = thr_tmem_load_epi.partition_S(tOtO_epi[(None, None), None])
tTMEM_LOADsO_epi = thr_tmem_load_epi.partition_D(tOsO_epi[(None, None), None])
tTMEM_LOADcO_epi = thr_tmem_load_epi.partition_D(tOcO_epi[(None, None), None])
n_epi_corr_tiles = self.pv_mma_tiler[1] // epi_corr_tile_size
for i in range(n_epi_corr_tiles):
tTMrO = cute.make_rmem_tensor(
tTMEM_LOADcO_epi[None, 0, 0, i].shape, self.acc_dtype
)
cute.copy(tiled_tmem_load_epi, tTMEM_LOADtO_epi[None, 0, 0, i], tTMrO)
for j in range(cute.size(tTMrO)):
tTMrO[j] = tTMrO[j] * inv_row_sum
tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype)
o_vec = tTMrO.load()
tSMrO.store(o_vec.to(self.o_dtype))
cute.copy(tiled_smem_store_epi, tSMrO, tTMEM_LOADsO_epi[None, 0, 0, i])
cute.arch.fence_proxy("async.shared", space="cta")
# TMA store SMEM → GMEM using pre-partitioned tensors
epi_bar = pipeline.NamedBarrier(
barrier_id=self.epilog_sync_bar_id,
num_threads=32 * len(self.epilogue_warp_id),
)
c_pipe.producer_tail()
epi_bar.arrive_and_wait()
cute.copy(tma_c, bSG_sC[(None, 0)], bSG_gC[(None, 0, 0, 0, 0, 0, 0)])
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)
tmem.relinquish_alloc_permit()
tmem.free(tmem_ptr)