[BugFix] Fix Async D2H copy bug & flash mash atten cache V out of bound bug (#7221) (#7294)

Co-authored-by: ming1753 <61511741+ming1753@users.noreply.github.com>

Author: EmmonsCurse

custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu

@@ -296,7 +296,7 @@ void GetBlockShapeAndSplitKVBlock(
   if (!phi::backends::gpu::IsCUDAGraphCapturing())
 #endif
     max_len_tensor_cpu.copy_(
-        max_len_tensor_gpu, max_len_tensor_cpu.place(), false);
+        max_len_tensor_gpu, max_len_tensor_cpu.place(), true);
 
   auto max_len_cpu_ptr = max_len_tensor_cpu.data<int>();
   int max_len_this_time = max_len_cpu_ptr[0];
@@ -379,7 +379,7 @@ void GetBlockShapeAndSplitKVBlock(
       if (!phi::backends::gpu::IsCUDAGraphCapturing())
 #endif
         decoder_num_blocks_cpu.copy_(
-            decoder_num_blocks_device, decoder_num_blocks_cpu.place(), false);
+            decoder_num_blocks_device, decoder_num_blocks_cpu.place(), true);
     }
   }
   // mla_backend not need run the following code.
@@ -410,7 +410,7 @@ void GetBlockShapeAndSplitKVBlock(
         block_size);
 
     kv_num_blocks_x_cpu.copy_(
-        kv_num_blocks_x, kv_num_blocks_x_cpu.place(), false);
+        kv_num_blocks_x, kv_num_blocks_x_cpu.place(), true);
     // Clear buffer
     const uint32_t encoder_max_tile_size_per_bs_q =
         div_up((max_enc_dec_len_this_time * group_size), encoder_block_shape_q);
@@ -434,7 +434,7 @@ void GetBlockShapeAndSplitKVBlock(
                                         encoder_block_shape_q,
                                         group_size);
     encoder_num_blocks_x_cpu.copy_(
-        encoder_num_blocks_x, encoder_num_blocks_x_cpu.place(), false);
+        encoder_num_blocks_x, encoder_num_blocks_x_cpu.place(), true);
   }
 }
 

custom_ops/gpu_ops/append_attn/pre_cache_len_concat.cu

@@ -87,9 +87,9 @@ std::vector<paddle::Tensor> PreCacheLenConcat(
       bsz,
       block_size);
   paddle::Tensor pre_cache_num_blocks_cpu =
-      pre_cache_num_blocks.copy_to(paddle::CPUPlace(), false);
+      pre_cache_num_blocks.copy_to(paddle::CPUPlace(), true);
   paddle::Tensor kv_token_num_cpu =
-      kv_token_num.copy_to(paddle::CPUPlace(), false);
+      kv_token_num.copy_to(paddle::CPUPlace(), true);
 
   return {
       cu_seqlens_k,

custom_ops/gpu_ops/flash_mask_attn/mainloop_attn.hpp

@@ -490,6 +490,23 @@ struct CollectiveMainloopAttn {
 
       softmax.rescale_o(tOrO, scores_scale);
       consumer_wait(pipeline_v, smem_pipe_read_v);
+      if (seq_len_k - n_block * kBlockN < kBlockN) {
+        int valid_k = seq_len_k - n_block * kBlockN;
+        auto sVt_this = sVt(_, _, smem_pipe_read_v.index());
+        constexpr int kHdLo = decltype(get<0, 0>(shape(sVt_this)))::value;
+        constexpr int kHdHi = decltype(get<0, 1>(shape(sVt_this)))::value;
+        if (thread_idx >= valid_k && thread_idx < kBlockN) {
+#pragma unroll
+          for (int hd_hi = 0; hd_hi < kHdHi; ++hd_hi) {
+#pragma unroll
+            for (int hd_lo = 0; hd_lo < kHdLo; ++hd_lo) {
+              sVt_this(make_coord(make_coord(hd_lo, hd_hi), thread_idx)) =
+                  Element(0);
+            }
+          }
+        }
+        cutlass::arch::fence_view_async_shared();
+      }
       gemm</*zero_init=*/false, /*wg_wait=*/-1>(
           tiled_mma1, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), tOrO);
       warp_scheduler_barrier_arrive();