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

D1.5: Use proven Stage C approach - normalize via TMEM round-trip + epilogue_tma_store

Browse Source
This commit is contained in:
biondizzle
2026-05-24 02:30:20 +00:00
parent 8010e3dda2
commit ece137b2c2
1 changed files with 36 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -15,16 +15,6 @@ import cutlass.torch as ct
import math
def _transform_partitioned_tensor_layout(tensor):
"""Transform MMA layout: ((ATOM_M, ATOM_N), MMA_M, MMA_N, ...rest)
-> ((ATOM_M, MMA_M), (ATOM_N, MMA_N), ...rest).
Same as CUTLASS utils.gemm.sm100.transform_partitioned_tensor_layout."""
layout = tensor.layout; shape = layout.shape; stride = layout.stride
new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:])
new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:])
return cute.make_tensor(tensor.iterator, cute.make_layout(shape=new_shape, stride=new_stride))
class FmhaKernel:
def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True):
self.head_dim = head_dim
@@ -416,104 +406,50 @@ class FmhaKernel:
final_o_bar.arrive_and_wait()
# ============================================================
# CORRECTION EPILOG: One-way TMEM → registers → normalize → SMEM → GMEM
# EPILOGUE: Normalize O + TMA store to GMEM
# ============================================================
# Follows CUTLASS epilogue_tma_store pattern exactly:
# transform_partitioned_tensor_layout → flat_divide →
# get_tmem_load_op → make_tmem_copy → partition_S →
# get_smem_store_op → make_tiled_copy_D → partition_D →
# cpasync.tma_partition → copy loop
# Eliminates the 3% per-tile TMEM round-trip error by using
# paired atoms that preserve the C-fragment layout.
# Step 1: Normalize O in TMEM via round-trip (3% error from hand-constructed
# atoms — D1.5 tracks the paired-atom fix).
# Step 2: Use CUTLASS epilogue_tma_store for TMEM→SMEM→GMEM write.
# ============================================================
# D5a: When normalize=False, still do one-way trip but skip 1/row_sum.
# D5a: When normalize=False, skip 1/row_sum (emit un-normalized O + LSE).
if const_expr(self.normalize):
inv_row_sum = Float32(1.0) / row_sum
# Normalize O: TMEM round-trip O *= inv_row_sum
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()
# Step 1: Transform partitioned tensor layouts (CUTLASS pattern)
# ((ATOM_M, ATOM_N), MMA_M, MMA_N, ...) -> ((ATOM_M, MMA_M), (ATOM_N, MMA_N), ...)
tOtO_xfm = _transform_partitioned_tensor_layout(
cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout))
tCgC_xfm = _transform_partitioned_tensor_layout(tCgC)
# Step 2: TMEM load copy (epilogue_tmem_copy_and_partition pattern)
from cutlass.utils.blackwell_helpers import get_tmem_load_op as _get_tmem_load_op
tmem_copy_atom = _get_tmem_load_op(
self.cta_tile_shape_mnk, self.c_layout, self.o_dtype, self.acc_dtype,
epi_tile, self.use_2cta_instrs,
# TMA store via CUTLASS epilogue_tma_store
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout)
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 = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
# flat_divide by epi_tile to create sub-tiled views
tOtO_epi = cute.flat_divide(tOtO_xfm, epi_tile)
tCgC_epi = cute.flat_divide(tCgC_xfm, epi_tile)
# make_tmem_copy with the first sub-tile shape
tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi[(None, None, 0, 0)])
thr_t2r = tiled_copy_t2r.get_slice(sfw_idx)
# Partition source (TMEM) and destination (GMEM-derived register shape)
tTR_tAcc = thr_t2r.partition_S(tOtO_epi)
tTR_gC = thr_t2r.partition_D(tCgC_epi)
# Register tensor shape: keep (T2R, T2R_M, T2R_N) dims, zero the rest
# After partition_D, tTR_gC has (T2R, T2R_M, T2R_N, EPI_M, EPI_N, REST...)
# We slice to get just the per-subtile register count.
_tTR_gC_shape = tTR_gC.shape
_n_rest = len(_tTR_gC_shape) - 3 # 3 leading dims: T2R, T2R_M, T2R_N
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[(None, None, None) + (0,) * _n_rest].shape, self.acc_dtype)
# Step 3: SMEM store copy (epilogue_smem_copy_and_partition pattern)
smem_copy_atom = get_smem_store_op(
self.c_layout, self.o_dtype, self.acc_dtype, tiled_copy_t2r)
tiled_copy_r2s = cute.make_tiled_copy_D(smem_copy_atom, tiled_copy_t2r)
thr_r2s = tiled_copy_r2s.get_slice(sfw_idx)
tRS_sC = thr_r2s.partition_D(sC)
tTR_rC = cute.make_rmem_tensor(tRS_sC[(None, None, None, 0)].shape, self.o_dtype)
# Step 4: TMA store partition (cpasync.tma_partition for S2G)
# flat_divide tCgC for TMA partition (need the un-xfm version for tma_partition)
tCgC_epi_tma = cute.flat_divide(tCgC, epi_tile)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_c, 0, cute.make_layout(1),
cute.group_modes(sC, 0, 2),
cute.group_modes(tCgC_epi_tma, 0, 2),
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
self, sfw_idx, 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,
)
# Step 5: Correction loop — for each sub-tile: TMEM → reg → normalize → SMEM
tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
subtile_cnt = cute.size(tTR_tAcc_g.shape, mode=[3])
for subtile_idx in range(subtile_cnt):
# Load O from TMEM (preserves C-fragment layout via paired atom)
tTR_tAcc_mn = tTR_tAcc_g[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
# Normalize: multiply by inv_row_sum (exact in FP32)
if const_expr(self.normalize):
for j in cutlass.range(cute.size(tTR_rAcc), vectorize=True):
tTR_rAcc[j] = tTR_rAcc[j] * inv_row_sum
# Convert to output dtype and store to SMEM
acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
acc_vec = acc_vec.to(self.o_dtype)
tTR_rC.store(acc_vec)
# Store to SMEM
c_buffer = subtile_idx % self.num_c_stage
cute.copy(tiled_copy_r2s, tTR_rC, tRS_sC[(None, None, None, c_buffer)])
# Fence and barrier
cute.arch.fence_proxy("async.shared", space="cta")
corr_epi_bar = pipeline.NamedBarrier(
barrier_id=5, num_threads=32 * len(self.epilogue_warp_id))
corr_epi_bar.arrive_and_wait()
# TMA store SMEM → GMEM
if warp_idx == self.epilogue_warp_id[0]:
# Group trailing modes and slice (CUTLASS pattern)
bSG_gC_flat = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
cute.copy(tma_c, bSG_sC[(None, c_buffer)], bSG_gC_flat[(None, Int32(0))])
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)
corr_epi_bar.arrive_and_wait()
c_pipe.producer_tail()
# D5a: Write LSE (log-softmax) when normalize=False
# lse = ln(row_sum) + row_max * ln(2)