diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 2548e3b1..c3cf87a3 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -360,20 +360,12 @@ 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 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) + # SMEM-P: TEMPORARILY zero-fill sP (debugging deadlock). + # The coordinate-indexed write causes a deadlock at hd=256. + # TODO: Fix the SMEM-P write path. 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] + _sP_nostage[(j0, j1), 0, (0, 0)] = BFloat16(0.0) cute.arch.fence_proxy("async.shared", space="cta") if kt > 0: tTMrO = cute.make_rmem_tensor(