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

revert: fmha.py to D1 (0f52c34, before D1.5 changes)

Browse Source
This commit is contained in:
biondizzle
2026-05-24 03:40:12 +00:00
parent 84f07aa955
commit 18ab507896
1 changed files with 26 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -319,6 +319,12 @@ class FmhaKernel:
tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i)
n_corr_tiles = self.pv_n_tile // corr_tile_size
# tTMrO register tensor (defined unconditionally for CuTeDSL scoping).
# Used for O rescale (kt > 0) and O normalization (after loop).
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
)
for kt in range(self.n_kv_tiles):
si_handle = s_cons.wait_and_advance()
@@ -363,9 +369,6 @@ class FmhaKernel:
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: write P to sP using coordinate-indexed store.
# Uses tTMEM_LOADcS identity tensor to get (m, k) coordinates.
# DEBUG: Write a known pattern to sP to verify the coordinate mapping.
# Pattern: sP[m, k] = (m + k) % 256 as BF16 (unique per position)
for j0 in range(32):
for j1 in range(4):
coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0]
@@ -374,13 +377,9 @@ class FmhaKernel:
k0 = k_coord % 16
k1 = (k_coord // 16) % 4
k2 = k_coord // 64
# Debug: write (m + k) mod 256 instead of actual P value
_sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0]
cute.arch.fence_proxy("async.shared", space="cta")
if kt > 0:
tTMrO = cute.make_rmem_tensor(
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
)
for i in range(n_corr_tiles):
tTMrO_i_ = tTMrO[None, i]
tTMrO_i_layout = cute.composition(
@@ -408,39 +407,20 @@ class FmhaKernel:
final_o_bar.arrive_and_wait()
# ============================================================
# EPILOGUE: Normalize O + TMA store to GMEM
# EPILOGUE: TMA store O to GMEM + compute LSE
# ============================================================
# 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.
# The raw un-normalized O in TMEM is perfect (cos 0.999998).
# TMEM round-trip normalization with hand-constructed atoms causes
# severe data corruption (53% error) due to layout mismatch with
# epilogue_tma_store's paired-atom addressing.
# Solution: always write raw O via epilogue_tma_store, compute LSE,
# and let the caller normalize externally using LSE.
# This is the D5a path — production-quality with zero precision loss.
# The TMEM round-trip normalization (normalize=True) is tracked as D1.5.
# ============================================================
# 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()
# TMA store via CUTLASS epilogue_tma_store
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout)
# TMA store via CUTLASS epilogue_tma_store (reads raw O from TMEM)
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.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(
@@ -453,17 +433,16 @@ class FmhaKernel:
)
c_pipe.producer_tail()
# D5a: Write LSE (log-softmax) when normalize=False
# lse = ln(row_sum) + row_max * ln(2)
# Compute LSE: lse = ln(row_sum) + row_max * ln(2)
# Always compute LSE (needed for external normalization).
# row_max is in scale_log2 domain, multiply by ln(2) to convert.
if const_expr(not self.normalize):
_row_max_safe = row_max
if row_max == -cutlass.Float32.inf:
_row_max_safe = Float32(0.0)
if sfw_idx == 0:
_ln2 = Float32(0.6931471805599453) # ln(2)
lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2
mLSE[0] = lse_val
_row_max_safe = row_max
if row_max == -cutlass.Float32.inf:
_row_max_safe = Float32(0.0)
if sfw_idx == 0:
_ln2 = Float32(0.6931471805599453) # ln(2)
lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2
mLSE[0] = lse_val
tmem.relinquish_alloc_permit()
tmem.free(tmem_ptr)