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

REVERT to working example7 (n=128 cos 0.999998). Example8 TMA fix didn't work.

Browse Source
This commit is contained in:
biondizzle
2026-05-22 20:28:15 +00:00
parent fa1f5b8ef6
commit 2aa6e4d234
1 changed files with 120 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

226
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -10,32 +10,30 @@ Two structural rules we had to learn the hard way:
(B) Hand-constructed TMEM load/store atoms (Ld32x32bOp + St32x32bOp built
independently) DO NOT preserve register tile shape across a round-trip.
Use paired atoms (or, as we discovered: independently constructed atoms
DO work if they're built from the SAME `Repetition(N)` count — the
Ld32x32bOp(Rep(16)) + St32x32bOp(Rep(16)) pair preserves the register
tile shape exactly because the atom width matches). This is what the
CUTLASS Blackwell FMHA reference does in `correction_rescale`.
(C) Multi-tile GMEM indexing: `kt` from cutlass.range constant-folds at trace
time, so all TMA loads address tile 0. Workaround: track an Int32
coordinate manually, BUT seed it from an SSA expression
(`n_kv_tiles - n_kv_tiles`) rather than a literal `Int32(0)`, so the JIT
sees it as a runtime register and propagates the `+= 1` as a tracked
loop-carried iter_args update.
A no-op TMEM-load-then-TMEM-store visibly corrupts data. Use the paired
atoms from `utils.sm100.get_tmem_load_op` + `get_smem_store_op` — they
are configured together for the same (mma_tiler, layout, dtype) combo
and the register tile shape lines up. This is what the CUTLASS Blackwell
FMHA reference does in `correction_epilog`.
Kernel structure:
1. Combined K+V pipeline (tx_count = K_bytes + V_bytes; one acquire per kt;
K and V share the same barrier slot). SMEM slot via kvh.index, GMEM via
manually-tracked kv_coord (SSA-seeded).
the cutlass.range loop variable.
2. Reference-style scaled epilogue: TMEM correction_rescale (O *= 1/row_sum
via paired Ld32x32b + St32x32b atoms), then standard epilogue_tma_store
to send O from TMEM through SMEM to GMEM.
2. Reference-style epilogue (TMEM → reg → scale by 1/row_sum → FP32→BF16 in
reg → SMEM via paired atoms → TMA SMEM→GMEM). One pass, no TMEM
round-trip, no `epilogue_tma_store` helper. Inline TMA store + named
barrier sync to substitute for what the helper would have done.
3. Per-tile O rescale (multiplying existing O by exp2(old_max - new_max)
before PV[kt]) lives in the softmax warp BEFORE softmax_done_bar.arrive().
Reuses the same paired-atom pattern as the final normalize.
3. Online softmax row_max / row_sum tracking is correct, but the per-tile
in-place TMEM O rescale (multiplying existing O by exp2(old_max - new_max)
before PV[kt]) is currently DISABLED. Fixing that requires applying the
same paired-atom pattern to a separate scratch SMEM buffer and bouncing
PV's accumulator through it, which is substantial work. For now, the
kernel is correct when row_max growth across tiles is mild. Long n with
pronounced max growth will drift; the fix path is well-defined.
4. final_o_bar (32 MMA + 128 softmax threads). MMA arrives between
acc_pipe.producer_commit and producer_tail; softmax arrives_and_waits
@@ -58,7 +56,6 @@ class FmhaV3StageCMulti:
def __init__(self, s_k=128, scale_softmax=None):
# s_k MUST equal actual sequence length n.
self.s_k = s_k
self.n_kv_tiles = s_k // 128
self.acc_dtype = Float32; self.qk_acc_dtype = Float32
self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16
self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1
@@ -151,9 +148,12 @@ class FmhaV3StageCMulti:
smem = utils.SmemAllocator(); st = smem.allocate(SS)
qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants()
# Combined K+V pipeline: each stage carries BOTH K and V loaded together.
kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants()
s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants()
softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id))
# Final-O sync: MMA arrives between producer_commit and producer_tail;
# softmax arrives_and_waits before reading O for the final normalize.
final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id))
acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True)
tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id)))
@@ -203,29 +203,22 @@ class FmhaV3StageCMulti:
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
# ===== TMA LOAD warp =====
# Multi-tile GMEM indexing trick:
# - kt from cutlass.range constant-folds at trace time → all TMA
# loads address tile 0 in compiled code. This is the actual
# observed behavior in CuTeDSL 4.5.1, not a hypothesis.
# - Manual kv_coord works IF its initial value is an SSA Int32
# (a runtime register) rather than a literal Int32(0).
# `n_kv_tiles - n_kv_tiles` is an opaque SSA zero — n_kv_tiles is
# itself an SSA value from cute.size(gK, mode=[3]). With the seed
# in SSA, the JIT treats kv_coord as a tracked loop-carried iter
# variable and propagates `kv_coord = kv_coord + 1` properly.
# - Read kv_coord BEFORE the increment; assignment via `=` (not
# augmented `+=`) avoids any in-place mutation ambiguity.
# GMEM tile coordinate: use the cutlass.range induction variable kt
# directly. CuTeDSL's `cutlass.range` doesn't auto-detect a Python `+=`
# rebinding as a loop-carried iter_args update — the JIT traces the
# body once and captures whatever value `kv_coord` had at trace time,
# so an outer `kv_coord = Int32(0)` plus a `kv_coord += 1` inside the
# loop bakes 0 into every iteration's TMA descriptor at runtime.
# The induction variable IS the loop-carried state, properly tracked.
if warp_idx == self.tma_warp_id:
qp.reset(); qh = qp.acquire_and_advance()
cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
qp.tail()
kvp.reset(); pk = kvp.try_acquire()
kv_coord = n_kv_tiles - n_kv_tiles # SSA runtime zero
for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1):
kvh = kvp.acquire_and_advance(pk)
cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
kv_coord = kv_coord + 1
cute.copy(tma_k, tBgK[(None, kt)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[(None, kt)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
pk = cutlass.Boolean(1)
kvp.tail()
@@ -255,6 +248,12 @@ class FmhaV3StageCMulti:
cute.arch.fence_view_async_tmem_store()
kvh.release()
acc_pipe.producer_commit(acc_st); acc_st.advance()
# Signal softmax FIRST so it can run normalize + epilogue. Then
# wait for the epilogue's consumer-release in producer_tail.
# Reverse order deadlocks: producer_tail blocks waiting for
# consumer release; softmax blocks at final_o_bar waiting for
# MMA arrive; the epilogue (which does the release) is gated
# behind softmax's final_o_bar wait. Cycle.
final_o_bar.arrive()
acc_pipe.producer_tail(acc_st)
@@ -286,36 +285,20 @@ class FmhaV3StageCMulti:
tScP = cute.make_tensor(tScS.iterator, tScP_layout)
tTMEM_STOREcP = thr_store.partition_S(tScP)
# === O rescale path setup (used per-tile AND for final normalize) ===
corr_tile_size = 16
cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO = pv_thr.partition_C(cO)
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 = HEAD_DIM // corr_tile_size
row_max = -Float32.inf
row_sum = Float32(0.0)
scale_log2 = Float32(self.scale_softmax_log2)
# Per-tile softmax loop with online rescale.
# Per-tile softmax loop.
# Online softmax row_max/row_sum tracking is maintained, but the
# in-place TMEM O rescale (which would multiply existing O by
# exp2(old_max - new_max) before PV[kt]) is DISABLED — this is the
# correctness compromise for hand-paired TMEM atoms not working.
# The fix path is to integrate the rescale into the same paired
# tmem_load/smem_store epilogue pattern we use below for normalize.
# For now: kernel is correct when row_max growth across tiles is
# mild (typical for short n with random data); for very long n
# the missing rescale shows as accuracy drift.
for kt in range(n_kv_tiles):
si_handle = s_cons.wait_and_advance()
@@ -324,7 +307,7 @@ class FmhaV3StageCMulti:
cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
cute.arch.fence_view_async_tmem_load()
# Pass 1: update row_max in log2-domain.
# Pass 1: update row_max (in log2-domain, fused with scale).
old_row_max = row_max
frg_cnt = 4
frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
@@ -337,8 +320,8 @@ class FmhaV3StageCMulti:
if row_max == -cutlass.Float32.inf:
row_max_safe = Float32(0.0)
# acc_scale = exp2(old_max - new_max). On first tile this is 0
# (old_max = -inf), so row_sum stays 0 and rescale is skipped.
# row_sum rescale (correct even without O rescale — row_sum
# is a register variable, not in TMEM).
# row_max is already in scaled domain, so no extra scale_log2.
acc_scale_ = old_row_max - row_max_safe
acc_scale = cute.math.exp2(acc_scale_, fastmath=True)
@@ -347,7 +330,7 @@ class FmhaV3StageCMulti:
row_sum *= acc_scale
# Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum,
# cast to BF16 via FP32-backed register bridge.
# store BF16 P through the FP32-backed register bridge.
rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype)
rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout)
minus_row_max = Float32(0.0) - row_max_safe
@@ -364,56 +347,86 @@ class FmhaV3StageCMulti:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
# === Per-tile O rescale: O *= acc_scale for kt > 0 ===
# Uses the SAME paired-atom pattern as the final normalize.
# Must run BEFORE softmax_done_bar.arrive() so MMA's PV[kt]
# reads the rescaled O.
# Visibility of MMA's PV[kt-1] writes: provided by
# s_cons.wait_and_advance at the top of this iteration, which
# acquires on MMA's S[kt] commit. S[kt] is sequenced after
# PV[kt-1] in MMA's iteration, so PV[kt-1]'s tmem_store_fence
# has been observed by the time we read O here.
if kt > 0:
for i in range(n_corr_tiles):
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.shape, self.acc_dtype)
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO)
cute.arch.fence_view_async_tmem_load()
for k in cutlass.range(cute.size(tTMrO), vectorize=True):
tTMrO[k] = tTMrO[k] * acc_scale
cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i)
cute.arch.fence_view_async_tmem_store()
si_handle.release()
softmax_done_bar.arrive()
# Wait for MMA's PV[N-1] to commit before reading O for normalize.
# === Reference-style scaled epilogue (no TMEM round-trip) ===
#
# Pattern (mirrors CUTLASS Blackwell FMHA reference's
# correction_epilog): for each column sub-tile,
# 1. TMEM -> registers via PAIRED tmem_load atom
# 2. scale in registers (1/row_sum)
# 3. FP32 -> BF16 conversion in registers
# 4. registers -> SMEM via PAIRED smem_store atom
# Then TMA SMEM -> GMEM as a separate step.
#
# Critical: the load and store atoms MUST be a matched pair.
# Independently constructed Ld32x32bOp + St32x32bOp atoms (the
# previous code) don't preserve the register tile shape, so even a
# no-op load+store corrupts data. Using utils.blackwell_helpers
# (sm100_utils) gives a paired set keyed to the same epi_subtile.
# Wait for MMA's PV[N-1] to commit before reading O.
final_o_bar.arrive_and_wait()
# === Final O normalization: O *= 1/row_sum ===
# === O normalization via TMEM load → scale → TMEM store ===
# Matches CUTLASS reference's correction_rescale pattern exactly.
corr_tile_size = 16
cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO = pv_thr.partition_C(cO)
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_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_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
)
inv_row_sum = Float32(1.0) / row_sum
for i in range(n_corr_tiles):
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_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,
tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout
)
tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype)
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO)
cute.arch.fence_view_async_tmem_load()
for k in cutlass.range(cute.size(tTMrO), vectorize=True):
tTMrO[k] = tTMrO[k] * inv_row_sum
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()
# Standard epilogue: TMEM → SMEM → GMEM via TMA store.
@@ -459,6 +472,7 @@ def test():
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
# Each n requires its own compiled kernel (s_k is compile-time).
kernel = FmhaV3StageCMulti(s_k=n)
print(f'n={n}: Compiling...', flush=True)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)