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

D1.5: Try O rescale with tCtO_base layout (epilogue-proven TMEM addressing)

Browse Source 
Previous attempts used tOtO0 (from pv_thr.make_fragment_C) and corrupted data.
This version uses tCtO_base (from pv_mma.make_fragment_C) which is the SAME
tensor the epilogue successfully reads O from. Both load and store atoms built
from same tCtO_i via composition — CUTLASS correction_rescale pattern.
This commit is contained in:
biondizzle
2026-05-27 02:10:39 +00:00
parent f94978ffa7
commit 79e2eb3b42
1 changed files with 65 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -410,22 +410,46 @@ class FmhaKernel:
scale_log2 = Float32(self.scale_softmax_log2)
# ============================================================
# D1.5: O RESCALE — SMEM ACCUMULATOR APPROACH
# =================================================
# TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken:
# even NO-OP round-trip corrupts data (ratio = -11 billion).
# Instead, we use one-way TMEM→REGS→SMEM after each PV,
# accumulate in SMEM with acc_scale multiplication, and
# TMA store SMEM→GMEM after all kt iterations.
#
# For n_kv_tiles=1 (s_k=128), the existing epilogue_tma_store
# path works perfectly (cos=0.999998). The SMEM accumulator
# is only needed for n_kv_tiles > 1.
# D1.5: O RESCALE ATOMS (using tCtO_base — proven-correct TMEM layout)
# ============================================================
# Use tCtO_base (from epilogue, proven to correctly read O from TMEM)
# for BOTH load and store atoms. Both copies built from same tensor
# so register layouts are compatible.
# ============================================================
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
# NOTE: The code below is the BROKEN TMEM round-trip approach.
# It's kept as reference but should NOT be used.
# The SMEM accumulator implementation is TODO.
# Build TMEM load+store atoms from tCtO_base via composition
corr_tile_size = 16
tCtO_i_layout = cute.composition(
tCtO_base.layout, cute.make_layout((128, corr_tile_size))
)
tCtO_i = cute.make_tensor(tCtO_base.iterator, tCtO_i_layout)
# Coordinate tensor for partition_D of load
cO = cute.make_identity_tensor((128, self.head_dim))
# Use pv_mma (not pv_thr) for partition_C — matches tCtO_base's layout
tOcO = pv_mma.get_slice(0).partition_C(cO)
tOcO_i_layout = cute.composition(
tOcO.layout, cute.make_layout((128, corr_tile_size))
)
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.qk_acc_dtype,
)
tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tCtO_i)
thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
tTMEM_LOADtO = thr_tmem_load_o.partition_S(tCtO_i)
tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i)
tmem_store_o_atom = cute.make_copy_atom(
tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
self.qk_acc_dtype,
)
tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tCtO_i)
thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
tTMEM_STOREtO = thr_tmem_store_o.partition_D(tCtO_i)
# prev_acc_scale: unused, kept for clarity. acc_scale at kt is used
# to rescale O from kt=0..kt-1 before PV[kt].
@@ -529,12 +553,33 @@ class FmhaKernel:
k2 = k_coord // 64
_sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0]
cute.arch.fence_proxy("async.shared", space="cta")
# D1.5: O rescale for kt > 0 — NOT YET IMPLEMENTED.
# TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken:
# even NO-OP round-trip corrupts O accumulator data.
# Production path for multi-KV-tile: Python KV merge (cos 0.999998).
# Future: SMEM accumulator approach (one-way TMEM→REGS→SMEM per kt).
# n_kv_tiles=1 is the only supported path for in-kernel processing.
# D1.5: O rescale for kt > 0 — using tCtO_base (proven-correct TMEM layout).
# Both load and store atoms built from same tCtO_i tensor via composition.
if const_expr(self.n_kv_tiles > 1):
if kt > 0:
pv_done_bar.arrive_and_wait() # Wait for PV[kt-1]
n_slices = self.head_dim // corr_tile_size
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, n_slices), self.qk_acc_dtype
)
for i in range(n_slices):
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)
cute.arch.fence_view_async_tmem_load()
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()
si_handle.release()
softmax_done_bar.arrive()