Add missing TMEM fence after P store in TMEM-P path

2026-05-23 05:17:45 +00:00
parent 0cd0e8b35f
commit a284580422
1 changed files with 1 additions and 0 deletions
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -430,6 +430,7 @@ class FmhaKernel:
                        rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype))

                    cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
+                    cute.arch.fence_view_async_tmem_store()
                else:
                    # SMEM-P: compute softmax, write P to SMEM (TODO)
                    for j in range(frg_cnt):