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

Merge branch 'master' of ssh://sweetapi.com:2222/biondizzle/nvfp4-megamoe-kernel

Browse Source
This commit is contained in:
biondizzle
2026-05-24 03:39:51 +00:00
parent 4a3033e37b 5ad8250ab0
commit 84f07aa955
1 changed files with 61 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -89,7 +89,7 @@ class FmhaKernel:
cute.size_in_bytes(self.q_dtype, v_s)) * cta
@cute.jit
def __call__(self, q, k, v, c, stream, lse=None, gP=None, tma_p=None):
def __call__(self, q, k, v, c, stream, lse=None):
self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
@@ -113,33 +113,19 @@ class FmhaKernel:
tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape)
epi_s = cute.select(self.c_smem_s,mode=[0,1])
tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
# SMEM-P: gP buffer and TMA for P (GMEM→SMEM via TMA)
if self.use_smem_p and gP is not None:
p_s = cute.slice_(self.p_smem_s,(None,None,None,0))
tma_p,gP = cute.nvgpu.make_tiled_tma_atom_A(
utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, pv_mma.thr_id),
gP, p_s, self.qk_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape
)
elif not self.use_smem_p:
tma_p = tma_q # dummy, dead code
else:
raise ValueError("use_smem_p=True but no gP provided")
# Always create a valid mLSE tensor for the kernel.
# CuTeDSL doesn't support None parameters in @cute.kernel.
# For normalize=True, mLSE is unused (dead-code-eliminated by compiler).
if const_expr(lse is None):
lse = cute.make_tensor(c.iterator, cute.make_layout((1,), stride=(0,)))
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,tma_p,gP,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream)
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream)
@cute.kernel
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, tma_p, mGP, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE):
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE):
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
tidx,_,_ = cute.arch.thread_idx()
if warp_idx == self.tma_warp_id:
cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c)
if const_expr(self.use_smem_p):
cpasync.prefetch_descriptor(tma_p)
@cute.struct
class SS:
@@ -240,12 +226,6 @@ class FmhaKernel:
sh.commit()
softmax_done_bar.arrive_and_wait()
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0)
if const_expr(self.use_smem_p):
# SMEM-P: TMA load gP → sP (MMA warp does this after barrier)
tPgP, tPsP = cpasync.tma_partition(tma_p, 0, cute.nvgpu.OperandMajorMode.M, cute.group_modes(sP,0,3), cute.group_modes(mGP,0,3))
cute.copy(tma_p, tPsP[(None,0,None,0)], tPgP[(None,0,None,0)], tma_bar_ptr=st.s_bar.data_ptr())
cpasync.commit_group()
cpasync.wait_group(0)
if not self.use_smem_p:
# TMEM-P: PV reads P from TMEM
for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True):
@@ -339,12 +319,6 @@ 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()
@@ -388,20 +362,25 @@ class FmhaKernel:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: write P to gP (global memory), then TMA loads gP→sP.
# rP_bf16 and tCgP are both partitioned by the QK MMA's C-fragment,
# so they have the same thread→value mapping. A simple element-wise
# copy from rP_bf16 to tCgP puts P values at the correct gP positions.
gP_local = cute.local_tile(mGP, (128, self.s_k), (0, 0))
tCgP = qk_thr.partition_C(gP_local)
# Flatten both tensors for element-wise copy
rP_flat = cute.make_tensor(rP_bf16.iterator, cute.coalesce(rP_bf16.layout))
gP_flat = cute.make_tensor(tCgP.iterator, cute.coalesce(tCgP.layout))
# Copy element-by-element (both should have 128 values per thread)
for idx in cutlass.range(cute.size(rP_flat), vectorize=True):
gP_flat[idx] = rP_flat[idx]
cute.arch.fence_proxy("async.global", space="cta")
# 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]
m_coord = coord[0]
k_coord = coord[1]
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(
@@ -429,20 +408,39 @@ class FmhaKernel:
final_o_bar.arrive_and_wait()
# ============================================================
# EPILOGUE: TMA store O to GMEM + compute LSE
# EPILOGUE: Normalize O + TMA store to GMEM
# ============================================================
# 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.
# 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.
# ============================================================
# 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)
# 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)
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(
@@ -455,16 +453,17 @@ class FmhaKernel:
)
c_pipe.producer_tail()
# Compute LSE: lse = ln(row_sum) + row_max * ln(2)
# Always compute LSE (needed for external normalization).
# D5a: Write LSE (log-softmax) when normalize=False
# lse = ln(row_sum) + row_max * ln(2)
# row_max is in scale_log2 domain, multiply by ln(2) to convert.
_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.to(self.q_dtype)
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
tmem.relinquish_alloc_permit()
tmem.free(tmem_ptr)