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

D1: SMEM-P coordinate-indexed store with dynamic sP shape extraction

Browse Source
This commit is contained in:
biondizzle
2026-05-24 03:34:03 +00:00
parent 00c9991396
commit f550c7dd33
1 changed files with 48 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -366,19 +366,54 @@ class FmhaKernel:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: write P to sP using make_tiled_copy_C + retile.
# The retile() call transforms rP_bf16 from QK C-fragment layout
# to the SMEM copy's source layout, matching partition_D(sP).
_smem_p_store_atom = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
self.q_dtype,
num_bits_per_copy=16,
)
_tiled_smem_p = cute.make_tiled_copy_C(_smem_p_store_atom, qk_mma)
_thr_smem_p = _tiled_smem_p.get_slice(sfw_idx)
_tRS_sP = _thr_smem_p.partition_D(_sP_nostage)
_tRS_rP = _tiled_smem_p.retile(rP_bf16)
cute.copy(_tiled_smem_p, _tRS_rP, _tRS_sP)
# SMEM-P: write P to sP using coordinate-indexed store.
# Uses tTMEM_LOADcS identity tensor to get (m, k) coordinates.
# The sP layout is PV A-operand SMEM: ((128,16),1,(4,2),1) for hd=64,
# but changes at larger hd. We write using sP's actual subtile structure.
# sP[(m_sub, k_sub), 0, (k_group1, k_group2)]
# where m is decomposed as m_sub in [0, 128) with subtile 16,
# and k is decomposed as k_sub in [0, 16), k_group1 in [0, 4), k_group2.
#
# The QK C-fragment has 128 columns (s_k=128). At hd=64, P uses 64 of them.
# At hd=128, P uses all 128.
# The sP subtile pattern depends on pv_mma's K-dim (pv_n_tile / 128 * qk_mma_tiler[2]).
#
# We iterate over the identity tensor to get (m, k) for each P value,
# then compute sP indices from the actual sP shape.
#
# For the sP layout ((M_atom, K_atom), 1, (K_group1, K_group2), stage):
# m_idx = m_coord (0..127)
# k_idx = k_coord (0..min(s_k, pv_n_tile)-1)
# sub_m = m_idx % M_atom, group_m = m_idx // M_atom
# sub_k = k_idx % K_atom, group_k1 = (k_idx // K_atom) % K_group1_size, group_k2 = k_idx // (K_atom * K_group1_size)
#
# We read from sP's shape to determine the tiling.
_sP_shape = cute.shape(_sP_nostage)
_M_atom = _sP_shape[0][0] # e.g. 128
_K_atom = _sP_shape[0][1] # e.g. 16
_K_g1 = _sP_shape[2][0] # e.g. 4
_K_g2 = _sP_shape[2][1] # e.g. 2
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]
# Skip if k_coord is beyond P's columns
# P has p_cols_fp32 * (qk_acc_dtype.width / q_dtype.width) BF16 elements per row
# p_cols_fp32 = pv_n_tile * q_dtype.width / qk_acc_dtype.width
# At hd=64: p_cols_bf16 = 64, so k ranges 0..63
# At hd=128: p_cols_bf16 = 128, so k ranges 0..127
# The identity tensor maps all 128 QK columns. We write all of them
# (PV only reads the first pv_n_tile columns from SMEM).
sub_m = m_coord % _K_atom # within-atom M index
grp_m = m_coord // _K_atom # M group (should be 0 for 128-row tile)
# Actually, sP mode 0 is (M_atom, K_atom), so M is the first sub-mode
# Let me re-examine: mode 0 is (128, 16) which is (M, K_subtile)
# So m_idx maps directly to mode 0 first dim, k_idx maps to mode 0 second dim + modes 2
k_sub = k_coord % _K_atom
k_g1 = (k_coord // _K_atom) % _K_g1
k_g2 = k_coord // (_K_atom * _K_g1)
_sP_nostage[(m_coord, k_sub), 0, (k_g1, k_g2)] = rP_bf16[(j0, 0), j1, 0, 0]
cute.arch.fence_proxy("async.shared", space="cta")
if kt > 0:
for i in range(n_corr_tiles):