mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-26 19:52:38 +00:00
Remove some D2D memcpy in T5 BeamSearch (#15706)
### Description <!-- Describe your changes. --> replace D2D memcpy with cuda kernels in creating decoder initial feeds before:  after:  ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> --------- Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang
GitHub
parent
0ffae8073b
commit
7f293d065a
3 changed files with 247 additions and 30 deletions
|
|
@ -6,6 +6,7 @@
|
|||
#include "cub/util_type.cuh"
|
||||
#include <cub/cub.cuh>
|
||||
#include <cub/device/device_segmented_radix_sort.cuh>
|
||||
#include "contrib_ops/cuda/bert/utils.cuh"
|
||||
#include "contrib_ops/cuda/transformers/generation_cuda_impl.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
|
@ -805,6 +806,206 @@ void UpdateDecoderMaskedMultiHeadAttentionCacheIndirection(int32_t* tgt_indir_ca
|
|||
max_seq_length,
|
||||
current_length);
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
namespace {
|
||||
template <typename T, size_t size>
|
||||
struct TypeMapper : public V_vec_m_<T, size> {};
|
||||
|
||||
template <>
|
||||
struct TypeMapper<int32_t, 2> {
|
||||
using Type = uint2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TypeMapper<int32_t, 4> {
|
||||
using Type = uint4;
|
||||
};
|
||||
} // namespace
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
__global__ void KeyCacheExpansionKernel(const T* input,
|
||||
T* output,
|
||||
int beam_width,
|
||||
int max_seq_length,
|
||||
int head_size) {
|
||||
const int num_heads = gridDim.y;
|
||||
const int sequence_length = gridDim.z;
|
||||
|
||||
const int bbid = blockIdx.x;
|
||||
const int batch_id = bbid / beam_width;
|
||||
const int head_id = blockIdx.y;
|
||||
const int s = blockIdx.z;
|
||||
const int tidx = threadIdx.x;
|
||||
|
||||
const int input_offset = ((batch_id * num_heads + head_id) * sequence_length + s) * head_size + tidx;
|
||||
const int output_offset = ((bbid * num_heads + head_id) * max_seq_length + s) * head_size + tidx;
|
||||
|
||||
if (tidx < head_size) {
|
||||
output[output_offset] = input[input_offset];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void KeyCacheExpansionKernelLauncher(const T* key_cache,
|
||||
T* key_cache_expanded,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int num_heads,
|
||||
int sequence_length,
|
||||
int max_seq_length,
|
||||
int head_size,
|
||||
cudaStream_t stream) {
|
||||
const dim3 grid(batch_size * beam_width, num_heads, sequence_length);
|
||||
|
||||
int equiv_head_size = (head_size & 1) == 0 ? (head_size >> 1) : head_size;
|
||||
equiv_head_size = (equiv_head_size & 1) == 0 ? (equiv_head_size >> 1) : equiv_head_size;
|
||||
|
||||
// Here we know head_size is smaller than max_thread_num_per_block
|
||||
int tpb = std::max(32, equiv_head_size);
|
||||
|
||||
// round up tpb to power of 2
|
||||
--tpb;
|
||||
tpb |= (tpb >> 1);
|
||||
tpb |= (tpb >> 2);
|
||||
tpb |= (tpb >> 4);
|
||||
tpb |= (tpb >> 8);
|
||||
tpb |= (tpb >> 16);
|
||||
tpb++;
|
||||
|
||||
#ifndef USE_ROCM
|
||||
if ((head_size % 4) == 0) {
|
||||
using vec_type = typename TypeMapper<T, 4>::Type;
|
||||
const dim3 block(tpb);
|
||||
KeyCacheExpansionKernel<<<grid, block, 0, stream>>>(reinterpret_cast<const vec_type*>(key_cache),
|
||||
reinterpret_cast<vec_type*>(key_cache_expanded),
|
||||
beam_width,
|
||||
max_seq_length,
|
||||
equiv_head_size);
|
||||
} else if ((head_size & 1) == 0) {
|
||||
using vec_type = typename TypeMapper<T, 2>::Type;
|
||||
const dim3 block(tpb);
|
||||
KeyCacheExpansionKernel<<<grid, block, 0, stream>>>(reinterpret_cast<const vec_type*>(key_cache),
|
||||
reinterpret_cast<vec_type*>(key_cache_expanded),
|
||||
beam_width,
|
||||
max_seq_length,
|
||||
equiv_head_size);
|
||||
} else {
|
||||
#endif
|
||||
const dim3 block(tpb);
|
||||
KeyCacheExpansionKernel<<<grid, block, 0, stream>>>(key_cache,
|
||||
key_cache_expanded,
|
||||
beam_width,
|
||||
max_seq_length,
|
||||
head_size);
|
||||
#ifndef USE_ROCM
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template void KeyCacheExpansionKernelLauncher(const float* key_cache,
|
||||
float* key_cache_expanded,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int num_heads,
|
||||
int sequence_length,
|
||||
int max_seq_length,
|
||||
int head_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
template void KeyCacheExpansionKernelLauncher(const half* key_cache,
|
||||
half* key_cache_expanded,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int num_heads,
|
||||
int sequence_length,
|
||||
int max_seq_length,
|
||||
int head_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
template void KeyCacheExpansionKernelLauncher(const int32_t* key_cache,
|
||||
int32_t* key_cache_expanded,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int num_heads,
|
||||
int sequence_length,
|
||||
int max_seq_length,
|
||||
int head_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
template <typename T>
|
||||
__global__ void BufferExpansionKernel(const T* input,
|
||||
T* output,
|
||||
int chunk_size) {
|
||||
const int batch_id = blockIdx.x;
|
||||
const int beam_id = blockIdx.y;
|
||||
const int tidx = threadIdx.x;
|
||||
const int beam_size = gridDim.y;
|
||||
const int idx = blockIdx.z * blockDim.x + tidx;
|
||||
|
||||
const int input_offset = batch_id * chunk_size + idx;
|
||||
const int output_offset = batch_id * beam_size * chunk_size + beam_id * chunk_size + idx;
|
||||
|
||||
if (idx < chunk_size) {
|
||||
output[output_offset] = input[input_offset];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void BufferExpansionKernelLauncher(const T* input,
|
||||
T* output,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int chunk_size,
|
||||
cudaStream_t stream) {
|
||||
const dim3 block(128);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
if ((chunk_size % 4) == 0) {
|
||||
using vec_type = typename TypeMapper<T, 4>::Type;
|
||||
const dim3 grid(batch_size, beam_width, (chunk_size / 4 + block.x - 1) / block.x);
|
||||
BufferExpansionKernel<<<grid, block, 0, stream>>>(reinterpret_cast<const vec_type*>(input),
|
||||
reinterpret_cast<vec_type*>(output),
|
||||
chunk_size / 4);
|
||||
} else if ((chunk_size & 1) == 0) {
|
||||
using vec_type = typename TypeMapper<T, 2>::Type;
|
||||
const dim3 grid(batch_size, beam_width, (chunk_size / 2 + block.x - 1) / block.x);
|
||||
BufferExpansionKernel<<<grid, block, 0, stream>>>(reinterpret_cast<const vec_type*>(input),
|
||||
reinterpret_cast<vec_type*>(output),
|
||||
chunk_size / 2);
|
||||
} else {
|
||||
#endif
|
||||
const dim3 grid(batch_size, beam_width, (chunk_size + block.x - 1) / block.x);
|
||||
BufferExpansionKernel<<<grid, block, 0, stream>>>(input,
|
||||
output,
|
||||
chunk_size);
|
||||
#ifndef USE_ROCM
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template void BufferExpansionKernelLauncher(const float* input,
|
||||
float* output,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int chunk_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
template void BufferExpansionKernelLauncher(const half* input,
|
||||
half* output,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int chunk_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
template void BufferExpansionKernelLauncher(const int32_t* input,
|
||||
int32_t* output,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int chunk_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -119,6 +119,25 @@ void UpdateDecoderMaskedMultiHeadAttentionCacheIndirection(int32_t* tgt_indir_ca
|
|||
int current_length,
|
||||
cudaStream_t stream);
|
||||
|
||||
template <typename T>
|
||||
void KeyCacheExpansionKernelLauncher(const T* key_cache,
|
||||
T* key_cache_expanded,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int num_heads,
|
||||
int sequence_length,
|
||||
int max_seq_length,
|
||||
int head_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
template <typename T>
|
||||
void BufferExpansionKernelLauncher(const T* input,
|
||||
T* output,
|
||||
int batch_size,
|
||||
int beam_width,
|
||||
int chunk_size,
|
||||
cudaStream_t stream);
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -1152,6 +1152,16 @@ Status UpdateDecoderFeeds(
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
namespace {
|
||||
template <typename T>
|
||||
struct ToCudaTypeWrapper : public ToCudaType<T> {};
|
||||
|
||||
template <>
|
||||
struct ToCudaTypeWrapper<int32_t> {
|
||||
using MappedType = int32_t;
|
||||
};
|
||||
} // namespace
|
||||
|
||||
template <typename T>
|
||||
Status ExpandBuffer(Stream* ort_stream,
|
||||
const OrtValue& input,
|
||||
|
|
@ -1188,23 +1198,18 @@ Status ExpandBuffer(Stream* ort_stream,
|
|||
|
||||
const T* input_data = input.Get<Tensor>().Data<T>();
|
||||
T* expanded_data = expanded.GetMutable<Tensor>()->MutableData<T>();
|
||||
T* target = expanded_data;
|
||||
|
||||
using CudaT = typename ToCudaTypeWrapper<T>::MappedType;
|
||||
|
||||
if (max_sequence_length == 0) {
|
||||
const int64_t& chunk_size = static_cast<int64_t>(input_shape.Size() / batch_size);
|
||||
|
||||
for (int i = 0; i < batch_size; i++) {
|
||||
for (int j = 0; j < num_beams; j++) {
|
||||
CUDA_RETURN_IF_ERROR(
|
||||
cudaMemcpyAsync(
|
||||
target,
|
||||
input_data + i * chunk_size,
|
||||
sizeof(T) * chunk_size,
|
||||
cudaMemcpyDeviceToDevice,
|
||||
cuda_stream));
|
||||
target += chunk_size;
|
||||
}
|
||||
}
|
||||
cuda::BufferExpansionKernelLauncher<CudaT>(reinterpret_cast<const CudaT*>(input_data),
|
||||
reinterpret_cast<CudaT*>(expanded_data),
|
||||
static_cast<int>(batch_size),
|
||||
num_beams,
|
||||
static_cast<int>(chunk_size),
|
||||
cuda_stream);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
|
|
@ -1213,24 +1218,16 @@ Status ExpandBuffer(Stream* ort_stream,
|
|||
// Expand from [B, N, S, H] to [B*beam, N, S_max, H]
|
||||
const int64_t& num_heads = input_shape[1];
|
||||
const int64_t& head_size = input_shape[3];
|
||||
const int64_t& input_offset = sequence_length * head_size;
|
||||
const int64_t& output_offset = max_sequence_length * head_size;
|
||||
const int64_t& NSH = input_offset * num_heads;
|
||||
|
||||
for (int i = 0; i < batch_size; i++) {
|
||||
for (int j = 0; j < num_beams; j++) {
|
||||
for (int k = 0; k < num_heads; k++) {
|
||||
CUDA_RETURN_IF_ERROR(
|
||||
cudaMemcpyAsync(
|
||||
target,
|
||||
input_data + i * NSH + k * input_offset,
|
||||
sizeof(T) * SafeInt<size_t>(input_offset),
|
||||
cudaMemcpyDeviceToDevice,
|
||||
cuda_stream));
|
||||
target += output_offset;
|
||||
}
|
||||
}
|
||||
}
|
||||
cuda::KeyCacheExpansionKernelLauncher<CudaT>(reinterpret_cast<const CudaT*>(input_data),
|
||||
reinterpret_cast<CudaT*>(expanded_data),
|
||||
static_cast<int>(batch_size),
|
||||
num_beams,
|
||||
static_cast<int>(num_heads),
|
||||
static_cast<int>(sequence_length),
|
||||
max_sequence_length,
|
||||
static_cast<int>(head_size),
|
||||
cuda_stream);
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue