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

Fix O rescale: use Stage C proven correction_rescale pattern

Browse Source
This commit is contained in:
biondizzle
2026-05-23 05:10:46 +00:00
parent 300482e40a
commit d15bb7b84a
1 changed files with 39 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
dsv4/kernels/attention/fmha.py
View File

@@ -362,13 +362,23 @@ class FmhaKernel:
# For now, zero sP as a stub — PV will read garbage/zero
pass
# ── O rescale / normalization setup (correction_rescale atoms) ──
# ── O rescale / normalization setup (correction_rescale pattern from Stage C) ──
corr_tile_size = 16
o_rescale_atom_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
o_rescale_atom_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
o_rescale_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], corr_tile_size)))
tiled_o_ld = tcgen05.make_tmem_copy(o_rescale_atom_ld, cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
tiled_o_st = tcgen05.make_tmem_copy(o_rescale_atom_st, cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
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)
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_i)
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_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)
n_corr_tiles = self.head_dim // corr_tile_size
# ── Online softmax state ──
row_max = -Float32.inf
@@ -441,30 +451,35 @@ class FmhaKernel:
# ── Per-tile O rescale (multiply O by acc_scale when kt > 0) ──
if kt > 0:
thr_ld = tiled_o_ld.get_slice(sfw_idx)
thr_st = tiled_o_st.get_slice(sfw_idx)
tOrO_src = thr_ld.partition_S(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
tOrO_dst = thr_st.partition_D(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
rO = cute.make_rmem_tensor(tOrO_src.shape, self.qk_acc_dtype)
cute.copy(tiled_o_ld, tOrO_src, rO)
for i in cutlass.range(cute.size(rO), vectorize=True):
rO[i] = rO[i] * acc_scale
cute.copy(tiled_o_st, rO, tOrO_dst)
tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype)
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 k in cutlass.range(cute.size(tTMrO_i), vectorize=True):
tTMrO_i[k] = tTMrO_i[k] * acc_scale
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
cute.arch.fence_view_async_tmem_store()
# ── Wait for MMA's final PV GEMM ──
final_o_bar.arrive_and_wait()
# ── O normalization: multiply O by 1/row_sum (TMEM round-trip) ──
inv_row_sum = Float32(1.0) / row_sum
thr_ld = tiled_o_ld.get_slice(sfw_idx)
thr_st = tiled_o_st.get_slice(sfw_idx)
tOrO_src = thr_ld.partition_S(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
tOrO_dst = thr_st.partition_D(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
rO = cute.make_rmem_tensor(tOrO_src.shape, self.qk_acc_dtype)
cute.copy(tiled_o_ld, tOrO_src, rO)
for i in cutlass.range(cute.size(rO), vectorize=True):
rO[i] = rO[i] * inv_row_sum
cute.copy(tiled_o_st, rO, tOrO_dst)
tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype)
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 k in cutlass.range(cute.size(tTMrO_i), vectorize=True):
tTMrO_i[k] = tTMrO_i[k] * inv_row_sum
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
cute.arch.fence_view_async_tmem_store()
# ── Epilogue: TMA store O → global ──