From ba6005316e8d06290532281dd7bf575dddd87b50 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 23:26:46 +0000 Subject: [PATCH] D1.3: Re-enable coordinate-indexed SMEM-P write with identity tensor coords --- dsv4/kernels/attention/fmha.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index c21b6ab7..392ab3a3 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -351,9 +351,17 @@ class FmhaKernel: cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) cute.arch.fence_view_async_tmem_store() else: - # SMEM-P: zero-fill sP for now (testing sP→PV pipeline) - for j in cutlass.range(cute.size(sP), vectorize=True): - sP[j] = self.q_dtype(0) + # SMEM-P: write P to sP using coordinate-indexed store. + # Uses tTMEM_LOADcS identity tensor to get (m, k) coordinates. + 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 + _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(