Synchronizes async preload before Q scaling

Waits for outstanding async loads and syncs threads so Q scaling never races ahead of shared-memory tiles.

Author: LoserCheems

csrc/flash_dmattn/src/flash_fwd_kernel.h

@@ -422,6 +422,9 @@ inline __device__ void compute_attn_1rowblock(const Params &params, const int bi
         binfo.actual_seqlen_k, binfo.actual_seqlen_q
     );
 
+    FLASH_NAMESPACE::cp_async_wait<0>();
+    __syncthreads();
+
     // Scale Q once before streaming loop KV
     if constexpr (Kernel_traits::Is_Q_in_regs) {
         #pragma unroll
@@ -1140,6 +1143,9 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params &params, cons
         binfo.actual_seqlen_k, binfo.actual_seqlen_q
     );
 
+    FLASH_NAMESPACE::cp_async_wait<0>();
+    __syncthreads();
+
     // Scale Q once before streaming loop KV
     #pragma unroll
     for (int i = 0; i < size(tSrQ); ++i) {