apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-12-24 13:28:13 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

skip DtoH capture (#4988)

Browse Source
This commit is contained in:
Sunny-bot1
2025-11-13 10:57:44 +08:00
committed by GitHub
parent 8329338d37
commit 5b24013d46
1 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu
View File

@@ -17,6 +17,7 @@
#ifndef PADDLE_WITH_CUSTOM_DEVICE_METAX_GPU
#include "paddle/phi/core/memory/memcpy.h"
#endif
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
#include "utils.cuh"
template <int THREADBLOCK_SIZE>
@@ -287,9 +288,12 @@ void GetBlockShapeAndSplitKVBlock(
seq_lens_encoder.data<int>(),
max_len_tensor_gpu.data<int>(),
bsz);
max_len_tensor_cpu.copy_(
max_len_tensor_gpu, max_len_tensor_cpu.place(), false);
// Note (sunxin): Skip capturing the DtoH copy (it's time-consuming); CPU data
// is only for branching in attention.
if (!phi::backends::gpu::IsCUDAGraphCapturing()) {
max_len_tensor_cpu.copy_(
max_len_tensor_gpu, max_len_tensor_cpu.place(), false);
}
auto max_len_cpu_ptr = max_len_tensor_cpu.data<int>();
int max_len_this_time = max_len_cpu_ptr[0];
@@ -398,9 +402,12 @@ void GetBlockShapeAndSplitKVBlock(
bsz,
decoder_block_shape_q,
group_size);
decoder_num_blocks_cpu.copy_(
decoder_num_blocks_device, decoder_num_blocks_cpu.place(), false);
// Note (sunxin): Skip capturing the DtoH copy (it's time-consuming); CPU
// data is only for branching in attention.
if (!phi::backends::gpu::IsCUDAGraphCapturing()) {
decoder_num_blocks_cpu.copy_(
decoder_num_blocks_device, decoder_num_blocks_cpu.place(), false);
}
PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
decoder_chunk_size_device.data<int>(), 64, sizeof(int32_t), stream));
}