From f550c7dd3302504a7501113f82b1beead31d4534 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 03:34:03 +0000 Subject: [PATCH] D1: SMEM-P coordinate-indexed store with dynamic sP shape extraction --- dsv4/kernels/attention/fmha.py | 61 ++++++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 13 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index e499fb11..0585724a 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -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):