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

O normalize: exact CUTLASS correction_rescale pattern with 2D reg tensor

Browse Source
This commit is contained in:
biondizzle
2026-05-22 19:50:37 +00:00
parent 009cf9f80d
commit 0bdcdc0efd
1 changed files with 36 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
tests/fmha_v3_stage_c_example7.py
View File

@@ -370,56 +370,62 @@ class FmhaV3StageCMulti:
final_o_bar.arrive_and_wait()
# === 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).
# Matches CUTLASS reference's correction_rescale pattern exactly.
corr_tile_size = 16
# 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)))
# TMEM load + store atoms (paired — same Repetition)
tmem_load_o_atom = cute.make_copy_atom(
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)
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
self.acc_dtype,
)
tmem_store_o_atom = cute.make_copy_atom(
tmem_store_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_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])
tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i)
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i)
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)
# 2D register tensor: (frg_shape, n_corr_tiles)
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
)
# Scale = 1/row_sum
inv_row_sum = Float32(1.0) / row_sum
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]
for i in range(HEAD_DIM // corr_tile_size):
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
)
tTMrO = cute.make_rmem_tensor(
tTMEM_LOADcO[None, 0, 0, i].shape, self.acc_dtype
)
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO)
# Scale in FP32 registers
for j in range(cute.size(tTMrO), vectorize=True):
tTMrO[j] = tTMrO[j] * inv_row_sum
# Write back to TMEM
cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i)
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()