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

O normalize using tmem_ptr base (same as epilogue) + CUTLASS sub-tile pattern

Browse Source
This commit is contained in:
biondizzle
2026-05-22 19:18:16 +00:00
parent 365e8f53af
commit 183292d919
1 changed files with 39 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
tests/fmha_v3_real_softmax.py
View File

@@ -226,19 +226,27 @@ class FmhaV3RealSoftmax:
tScP = cute.make_tensor(tScS.iterator, tScP_layout)
tTMEM_STOREcP = thr_store.partition_S(tScP)
# O normalize setup: use FULL O layout (no sub-tiling)
# Match the P store pattern for TMEM access
tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(64)), self.acc_dtype)
tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(64)), self.acc_dtype)
tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO0)
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO0)
thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO0)
# O normalize setup: use the SAME base pointer as the epilogue
# The epilogue reads O from tmem_ptr + tmem_o0_offset.
# We must use the same base to access the correct TMEM columns.
tCtO_norm = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout)
cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO = pv_thr.partition_C(cO)
tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO)
tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO0)
# Sub-tile the O layout for the normalize copy
corr_tile_size = 16
tOtO_i_layout = cute.composition(tCtO_norm.layout, cute.make_layout((128, corr_tile_size)))
tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size)))
tOtO_norm_i = cute.make_tensor(tCtO_norm.iterator, tOtO_i_layout)
tOcO_norm_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout)
tmem_load_o_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(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_norm_i)
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_norm_i)
thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_norm_i)
tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_norm_i)
tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_norm_i)
row_max = -Float32.inf
row_sum = Float32(0.0)
@@ -298,11 +306,26 @@ class FmhaV3RealSoftmax:
# Final O normalization: O = O / row_sum
if row_sum != Float32(0.0):
inv_row_sum = Float32(1.0) / row_sum
tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype)
cute.copy(tiled_tmem_load_o, tTMEM_LOAD_OtO, tTMrO)
for j in cutlass.range(cute.size(tTMrO), vectorize=True):
tTMrO[j] = tTMrO[j] * inv_row_sum
cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_OtO)
n_corr = HEAD_DIM // corr_tile_size
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOAD_OcO.shape, n_corr), self.acc_dtype
)
for ci in range(n_corr):
tTMrO_ci_ = tTMrO[None, ci]
tTMrO_ci_layout = cute.composition(
tTMrO_ci_.layout, cute.make_layout(tTMrO.shape[0])
)
tTMrO_ci = cute.make_tensor(tTMrO_ci_.iterator, tTMrO_ci_layout)
tTMEM_LOAD_OtO_ci = cute.make_tensor(
tTMEM_LOAD_OtO.iterator + ci * corr_tile_size, tTMEM_LOAD_OtO.layout
)
tTMEM_STORE_OtO_ci = cute.make_tensor(
tTMEM_STORE_OtO.iterator + ci * corr_tile_size, tTMEM_STORE_OtO.layout
)
cute.copy(tiled_tmem_load_o, tTMEM_LOAD_OtO_ci, tTMrO_ci)
for j in cutlass.range(cute.size(tTMrO_ci), vectorize=True):
tTMrO_ci[j] = tTMrO_ci[j] * inv_row_sum
cute.copy(tiled_tmem_store_o, tTMrO_ci, tTMEM_STORE_OtO_ci)
# Epilogue: TMEM -> SMEM -> GMEM via TMA store
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)