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

D1.5: Implement correction epilog with get_tmem_load_op + get_smem_store_op paired atoms

Browse Source 
- One-way: TMEM → registers (normalize) → SMEM → GMEM
- Eliminates TMEM round-trip error for O normalization
- O rescale (kt>0) still uses old atoms (fix later)
- Based on CUTLASS FMHA reference's correction_epilog pattern
This commit is contained in:
biondizzle
2026-05-24 00:30:38 +00:00
parent 90131da010
commit ea5662ab2b
1 changed files with 68 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,7 +9,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
from cutlass.utils import LayoutEnum
from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset
from cutlass.utils.blackwell_helpers import get_smem_store_op
from cutlass.utils.blackwell_helpers import get_smem_store_op, get_tmem_load_op
import cuda.bindings.driver as cuda
import cutlass.torch as ct
import math
@@ -413,83 +413,86 @@ class FmhaKernel:
# ============================================================
# CORRECTION EPILOG: One-way TMEM → registers → normalize → SMEM → GMEM
# ============================================================
# Replace the old TMEM round-trip (3% error per trip) with the proper
# CUTLASS correction epilog pattern using paired atoms.
# This is a ONE-WAY trip: TMEM → registers (get_tmem_load_op) →
# normalize → SMEM (get_smem_store_op) → GMEM (TMA store).
# No TMEM write-back, no layout mismatch, no data corruption.
# Uses paired atoms from get_tmem_load_op + get_smem_store_op
# to preserve the C-fragment layout. No TMEM write-back.
# Based on CUTLASS FMHA reference's correction_epilog pattern.
# ============================================================
# D5a: When normalize=False, we still do the one-way trip but skip the 1/row_sum multiply.
# D5a: When normalize=False, still do one-way trip but skip 1/row_sum.
if const_expr(self.normalize):
inv_row_sum = Float32(1.0) / row_sum
# Build the TMEM→reg and reg→SMEM tiled copies using paired atoms
# Correction tile: split O into (128, corr_tile_size) sub-tiles
corr_tile_size = 32 * 8 // self.o_dtype.width # 32*8/16 = 16 for BF16
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
tCtO = utils.gemm.sm100.transform_partitioned_tensor_layout(tCtO_base)
tiled_copy_t2r, tTR_tO, tTR_rO = utils.gemm.sm100.epilogue_tmem_copy_and_partition(
self, tidx, tCtO, tCgC, epi_tile, self.use_2cta_instrs
cO = cute.make_identity_tensor(self.pv_mma_tiler[:2])
tOcO = pv_thr.partition_C(cO)
tOsO = pv_thr.partition_C(sC)
tOtO_i = cute.logical_divide(tCtO_base, cute.make_layout((128, corr_tile_size)))
tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size)))
tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size)))
# Build TMEM load copy using get_tmem_load_op (paired atom)
epi_subtile = (self.epi_tile[0], corr_tile_size)
tmem_copy_atom = utils.blackwell_helpers.get_tmem_load_op(
self.pv_mma_tiler, self.c_layout, self.o_dtype, self.acc_dtype,
epi_subtile, use_2cta_instrs=self.use_2cta_instrs,
)
tTR_rC = cute.make_rmem_tensor(tTR_rO.shape, self.c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = utils.gemm.sm100.epilogue_smem_copy_and_partition(
self, tiled_copy_t2r, tTR_rC, tidx, sC
)
tCgC_epi = cute.flat_divide(tCgC, epi_tile)
bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(
tma_c, 0, cute.make_layout(1),
cute.group_modes(sC, 0, 2),
cute.group_modes(tCgC_epi, 0, 2),
tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0])
# Build SMEM store copy using get_smem_store_op (paired with TMEM load)
smem_copy_atom = utils.blackwell_helpers.get_smem_store_op(
self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load
)
tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load)
# Partition source (TMEM) and destination (SMEM) for each thread
thr_tmem_load = tiled_tmem_load.get_slice(sfw_idx)
thr_smem_store = tiled_smem_store.get_slice(sfw_idx)
tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None])
tSMEM_STOREsO = thr_smem_store.partition_D(tOsO_i[(None, None), None])
tSMEM_STOREcO = thr_smem_store.partition_S(tOcO_i[(None, None), None])
n_corr_tiles = self.pv_n_tile // corr_tile_size
# For each correction tile: TMEM → reg (normalize) → SMEM
for i in range(n_corr_tiles):
tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i]
tSMEM_STOREsO_i = tSMEM_STOREsO[None, 0, 0, i]
tSMEM_STOREcO_i = tSMEM_STOREcO[None, 0, 0, i]
tTMrO = cute.make_rmem_tensor(tSMEM_STOREcO_i.shape, self.acc_dtype)
# Load O from TMEM using paired atom
cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO)
# Normalize: multiply by inv_row_sum
if const_expr(self.normalize):
for j in cutlass.range(cute.size(tTMrO), vectorize=True):
tTMrO[j] = tTMrO[j] * inv_row_sum
# Convert to output dtype and store to SMEM via paired atom
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, tSMrO, tSMEM_STOREsO_i)
# Fence SMEM writes and sync before TMA store
cute.arch.fence_proxy("async.shared", space="cta")
epilog_sync_bar = pipeline.NamedBarrier(
barrier_id=self.epilog_sync_bar_id,
num_threads=32 * len(self.epilogue_warp_id),
)
# Consume the accumulator pipeline
acc_cons_st = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
c_pipe = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)),
)
acc_pipe.consumer_wait(acc_cons_st)
# Slice to the current tile
tTR_tO_tile = tTR_tO[(None, None, None, None, None, acc_cons_st.index)]
bSG_gC = bSG_gC_partitioned[(None, None, None, Int32(0), Int32(0), Int32(0))]
tTR_tO_tile = cute.group_modes(tTR_tO_tile, 3, cute.rank(tTR_tO_tile))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
# Store O to global memory in subtiles, applying 1/row_sum normalize
subtile_cnt = cute.size(tTR_tO_tile.shape, mode=[3])
for subtile_idx in range(subtile_cnt):
tTR_tO_mn = tTR_tO_tile[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tO_mn, tTR_rO)
# Apply normalize: multiply by inv_row_sum, then convert to BF16
if const_expr(self.normalize):
for j in cutlass.range(cute.size(tTR_rO), vectorize=True):
tTR_rO[j] = tTR_rO[j] * inv_row_sum
acc_vec = tiled_copy_r2s.retile(tTR_rO).load()
acc_vec = acc_vec.to(self.c_dtype)
tRS_rC.store(acc_vec)
c_buffer = subtile_cnt * 0 + subtile_idx # num_prev_subtiles = 0
c_buffer = c_buffer % self.num_c_stage
cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
cute.arch.fence_proxy("async.shared", space="cta")
epilog_sync_bar.arrive_and_wait()
if warp_idx == self.epilogue_warp_id[0]:
cute.copy(tma_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, subtile_idx)])
c_pipe.producer_commit()
c_pipe.producer_acquire()
epilog_sync_bar.arrive_and_wait()
epilog_sync_bar.arrive_and_wait()
acc_pipe.consumer_release(acc_cons_st)
acc_cons_st.advance()
# TMA store: SMEM → GMEM
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)
c_pipe.producer_acquire()
cute.copy(tma_c, sC, tCgC)
c_pipe.producer_commit()
cute.arch.gpu_bar_sync()
c_pipe.producer_tail()
# D5a: Write LSE (log-softmax) when normalize=False