diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index f319a43f..565f66a8 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -394,7 +394,8 @@ class FmhaKernel: gP_flat = cute.make_tensor(tCgP.iterator, cute.coalesce(tCgP.layout)) for idx in cutlass.range(cute.size(rP_flat), vectorize=True): gP_flat[idx] = rP_flat[idx] - cute.arch.fence_proxy("async.global", space="cta") + # Ensure gP writes are visible before TMA reads + cute.arch.fence_proxy("async", space="cta") cute.arch.fence_proxy("async.shared", space="cta") if kt > 0: for i in range(n_corr_tiles):