apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2026-04-23 00:17:25 +08:00
Code Issues Actions 19 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
ming1753
2026-04-10 11:31:51 +08:00
committed by GitHub
parent 3c54a41131
commit 734fbcffde
3 changed files with 23 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu
+4 -4
View File
@@ -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];
@@ -378,7 +378,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.
@@ -409,7 +409,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);
@@ -433,7 +433,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
+2 -2
View File
@@ -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
+17
View File
@@ -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();