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

O normalize: TMEM round-trip with paired Ld/St atoms + standard epilogue_tma_store

Browse Source
This commit is contained in:
biondizzle
2026-05-22 19:49:51 +00:00
parent f821dd00fe
commit 009cf9f80d
1 changed files with 36 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
tests/fmha_v3_stage_c_example7.py
View File

@@ -369,44 +369,36 @@ class FmhaV3StageCMulti:
# Wait for MMA's PV[N-1] to commit before reading O.
final_o_bar.arrive_and_wait()
# === Reference-style scaled epilogue (mirrors CUTLASS FMHA correction_epilog) ===
# Pattern: TMEM → reg (paired load atom) → scale in reg → FP32→BF16 in reg
# → SMEM (paired store atom) → TMA SMEM→GMEM. No TMEM round-trip.
# === O normalization via TMEM load → scale → TMEM store ===
# Matches CUTLASS reference's correction_rescale pattern.
# Uses Ld32x32bOp / St32x32bOp with the SAME Repetition so the
# register tile shapes match (paired atoms).
corr_tile_size = 16 # matches the reference
corr_tile_size = 16
# Sub-tile the O C-fragment for column-wise iteration
# Sub-tile the O C-fragment
tOtO_i = cute.logical_divide(tOtO0, cute.make_layout((128, corr_tile_size)))
cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO = pv_thr.partition_C(cO)
tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size)))
tOsO = pv_thr.partition_C(sC[None, None, 0])
tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size)))
# Paired atoms via sm100_utils (same as CUTLASS reference)
epi_subtile = (self.epi_tile[0], corr_tile_size)
tmem_load_op = utils.sm100.get_tmem_load_op(
self.pv_mma_tiler,
self.c_layout,
self.o_dtype,
# TMEM load + store atoms (paired — same Repetition)
tmem_load_o_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
self.acc_dtype,
epi_subtile,
use_2cta_instrs=False,
)
tiled_tmem_load_o = tcgen05.make_tmem_copy(
tmem_load_op, tOtO_i[(None, None), 0]
)
thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
smem_store_op = utils.sm100.get_smem_store_op(
self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load_o
)
tiled_smem_store_o = cute.make_tiled_copy_D(
smem_store_op, tiled_tmem_load_o
tmem_store_o_atom = cute.make_copy_atom(
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[(None, None), 0])
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i[(None, None), 0])
thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i[(None, None), None])
tTMEM_LOADsO = thr_tmem_load_o.partition_D(tOsO_i[(None, None), None])
tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i[(None, None), None])
tTMEM_LOADtO = thr_load_o.partition_S(tOtO_i[(None, None), None])
tTMEM_LOADcO = thr_load_o.partition_D(tOcO_i[(None, None), None])
tTMEM_STOREtO = thr_store_o.partition_D(tOtO_i[(None, None), None])
# Scale = 1/row_sum
inv_row_sum = Float32(1.0) / row_sum
@@ -414,7 +406,9 @@ class FmhaV3StageCMulti:
n_corr = self.pv_mma_tiler[1] // corr_tile_size
for i in range(n_corr):
tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i]
tTMEM_LOADsO_i = tTMEM_LOADsO[None, 0, 0, i]
tTMEM_STOREtO_i = cute.make_tensor(
tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout
)
tTMrO = cute.make_rmem_tensor(
tTMEM_LOADcO[None, 0, 0, i].shape, self.acc_dtype
)
@@ -424,47 +418,25 @@ class FmhaV3StageCMulti:
for j in range(cute.size(tTMrO), vectorize=True):
tTMrO[j] = tTMrO[j] * inv_row_sum
# FP32 → BF16 in registers
tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype)
o_vec = tTMrO.load()
tSMrO.store(o_vec.to(self.o_dtype))
# Write back to TMEM
cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i)
# Registers → SMEM via paired atom
cute.copy(tiled_smem_store_o, tSMrO, tTMEM_LOADsO_i)
cute.arch.fence_view_async_tmem_store()
cute.arch.fence_view_async_tmem_load()
# Async-proxy fence so the TMA store sees the SMEM writes.
cute.arch.fence_proxy("async.shared", space="cta")
# Use NamedBarrier to sync softmax warps with TMA store warp
epi_sync_bar = pipeline.NamedBarrier(
barrier_id=self.epilog_sync_bar_id,
num_threads=32 * len(self.epilogue_warp_id),
)
epi_sync_bar.arrive_and_wait()
# TMA SMEM -> GMEM. One warp issues the copy; the rest waited at
# the named barrier above. (Match epilogue_tma_store's behavior
# with all-thread arrive.)
if warp_idx == self.epilogue_warp_id[0]:
# Partition sC and gC for TMA.
tOsO_tma, tOgO_tma = cpasync.tma_partition(
tma_c,
0,
cute.make_layout(1),
cute.group_modes(sC, 0, 2),
cute.group_modes(tCgC, 0, 2),
)
# tOgO_tma still has the trailing tile/batch coords; we want
# the first (and only) tile here.
cute.copy(tma_c, tOsO_tma[None, 0], tOgO_tma[None, 0, 0, 0])
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)
# Release the acc pipe so MMA's producer_tail can complete.
# 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
)
acc_pipe.consumer_release(acc_cons_st)
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,
)
c_pipe.producer_tail()
tmem.relinquish_alloc_permit()
tmem.free(tmem_ptr)