From 9346a81d4305df60a2258a2c56932094b115d58a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 09:19:35 +0000 Subject: [PATCH] D1.3: Implement SMEM-P path (write P to SMEM via tiled_smem_copy instead of zeroing sP) --- dsv4/kernels/attention/fmha.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 14bad065..b857a374 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -15,7 +15,7 @@ import math class FmhaKernel: - def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None): + def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None): self.head_dim = head_dim self.s_k = s_k self.n_kv_tiles = s_k // 128 @@ -344,12 +344,14 @@ class FmhaKernel: cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) cute.arch.fence_view_async_tmem_store() else: - # SMEM-P: TODO — write P to SMEM via make_tiled_copy_C(store_atom, qk_mma) - # For now, zero sP as stub. PV will produce garbage with SMEM-P path. - for j in cutlass.range(cute.size(sP), vectorize=True): - sP[j] = BFloat16(0.0) + # SMEM-P: write P to SMEM via tiled_smem_copy + # rP_bf16 contains P values in QK C-fragment layout (BF16) + # Flatten to 2D for copy operation + rP_bf16_2d = cute.group_modes(rP_bf16, 0, 2) + tSMEM_CPYrP = thr_smem_copy.partition_S(rP_bf16_2d) + cute.copy(tiled_smem_copy, tSMEM_CPYrP, tSMEM_CPYsP) cute.arch.fence_proxy("async.shared", space="cta") - + softmax_done_bar.arrive() # Per-tile O rescale (hand-constructed atoms with logical_divide layout) if kt > 0: tTMrO = cute.make_rmem_tensor(