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

REVERT to 0bdcdc0 — the version that passed n=128 cos 0.999998

Browse Source
This commit is contained in:
biondizzle
2026-05-22 20:07:18 +00:00
parent 3814838107
commit 242ffebcd9
1 changed files with 77 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
tests/fmha_v3_stage_c_example7.py
View File

@@ -203,19 +203,22 @@ class FmhaV3StageCMulti:
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
# ===== TMA LOAD warp =====
# Use the SAME pattern as the working test_fmha_v3_diag.py:
# kv_coord = Int32(0+0) + kv_coord += 1 in cutlass.range
# 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 = Int32(0 + 0)
for kt in cutlass.range(self.s_k // 128, unroll=1):
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 += 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()
@@ -289,44 +292,16 @@ class FmhaV3StageCMulti:
# 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
# O corr setup: DISABLED to debug n=128 regression
corr_tile_size = 16
cO_corr = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
tOcO_corr = pv_thr.partition_C(cO_corr)
tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size)))
tOcO_i_layout = cute.composition(tOcO_corr.layout, cute.make_layout((128, corr_tile_size)))
tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout)
tOcO_i = cute.make_tensor(tOcO_corr.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_LOAD_OtO = thr_tmem_load_o.partition_S(tOtO_i)
tTMEM_LOAD_OcO = thr_tmem_load_o.partition_D(tOcO_i)
tTMEM_STORE_OtO = thr_tmem_store_o.partition_D(tOtO_i)
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOAD_OcO.shape, 128 // corr_tile_size), self.acc_dtype)
# DISABLED: O rescale (kt > 0)
# if kt > 0:
# for ci in range(HEAD_DIM // corr_tile_size):
# ...
# DISABLED: Final O normalize (1/row_sum)
# inv_row_sum = Float32(1.0) / row_sum
# for i in range(HEAD_DIM // corr_tile_size):
# ...
row_max = -Float32.inf
row_sum = Float32(0.0)
scale_log2 = Float32(self.scale_softmax_log2)
# 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()
# Load S[kt]
tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype)
cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
@@ -354,8 +329,6 @@ class FmhaV3StageCMulti:
acc_scale = Float32(0.0)
row_sum *= acc_scale
# O rescale: DISABLED (debugging n=128 regression)
# Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum,
# store BF16 P through the FP32-backed register bridge.
rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype)
@@ -396,7 +369,65 @@ class FmhaV3StageCMulti:
# Wait for MMA's PV[N-1] to commit before reading O.
final_o_bar.arrive_and_wait()
# === Final O normalization: DISABLED (debugging) ===
# === 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(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_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 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.
# O in TMEM is now scaled by 1/row_sum.