saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-08 17:17:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve rpb cuda kernel (#14195)

Browse source 
### Description
Average latency (ms) of float16 relative position bias cuda kernel on
V100:

Kernel\Seq_Len  | 16 | 32 | 64 | 128 | 256 | 384 | 512 | 768 | 1024 |
2048 | 4096
-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --
Before| 0.0494 | 0.0654 | 0.1519 | 0.4322 | 1.1865 | 2.4091 | 4.3676 |
14.912 | 36.517 | 142.09 | 561.80
After | 0.0483 | 0.0651 | 0.1294 | 0.3858 | 1.1128 | 2.2988 | 3.8391 |
14.290 | 34.542 | 136.13 | 529.54

### 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. -->
Review of this comment
https://github.com/microsoft/onnxruntime/pull/14149/#discussion_r1063152021

Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2023-01-10 17:06:42 -08:00 committed by GitHub
parent d152452d4b
commit 342a5bf2b7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 49 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/contrib_ops/cuda/bert/relative_attn_bias.cc
View file

@ -57,6 +57,7 @@ Status RelPosAttnBias<T>::ComputeInternal(OpKernelContext* context) const {
typedef typename ToCudaType<T>::MappedType CudaT;
auto& device_prop = GetDeviceProp();
return LaunchRelPosAttnBiasKernel<CudaT>(Stream(context),
reinterpret_cast<CudaT*>(output->template MutableData<T>()),
reinterpret_cast<const CudaT*>(bias_table->template Data<T>()),
@ -64,7 +65,8 @@ Status RelPosAttnBias<T>::ComputeInternal(OpKernelContext* context) const {
static_cast<int>(query_len),
static_cast<int>(num_buckets),
max_distance_,
is_bidirectional_);
is_bidirectional_,
device_prop.maxThreadsPerBlock);
}
} // namespace cuda

59
onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu
View file

@ -36,7 +36,7 @@ __global__ void buildRelativeAttentionBias(T* relative_attention_bias,
const bool is_bidirectional,
const int max_distance) {
const int head_id = blockIdx.x;
for (int seq_id = threadIdx.x; seq_id < seq_len * seq_len; seq_id += blockDim.x) {
for (int seq_id = threadIdx.x; seq_id < seq_len * seq_len; seq_id += blockDim.x * gridDim.y) {
int row_id = seq_id / seq_len;
int col_id = seq_id % seq_len;
@ -86,20 +86,47 @@ Status LaunchRelPosAttnBiasKernel(
const int seq_len,
const int num_bucket,
const int max_distance,
const bool is_bidirectional)
const bool is_bidirectional,
const int max_threads_per_block)
{
dim3 grid(num_heads);
dim3 block(256);
const int squared_sq_len = seq_len * seq_len;
if (squared_sq_len <= max_threads_per_block) {
dim3 grid(num_heads);
dim3 block(squared_sq_len);
buildRelativeAttentionBias<<<grid, block, 0, stream>>>(output,
bias_table,
num_heads,
seq_len,
num_bucket,
is_bidirectional,
max_distance);
return CUDA_CALL(cudaGetLastError());
} else if (seq_len >= 128 && seq_len <= 384) {
dim3 grid(num_heads, seq_len);
dim3 block(seq_len);
buildRelativeAttentionBias<<<grid, block, 0, stream>>>(output,
bias_table,
num_heads,
seq_len,
num_bucket,
is_bidirectional,
max_distance);
return CUDA_CALL(cudaGetLastError());
} else {
int blockSize = max_threads_per_block;
const int grid_y_Size = (squared_sq_len + blockSize - 1) / blockSize;
dim3 grid(num_heads, grid_y_Size);
dim3 block(blockSize);
buildRelativeAttentionBias<<<grid, block, 0, stream>>>(output,
bias_table,
num_heads,
seq_len,
num_bucket,
is_bidirectional,
max_distance);
buildRelativeAttentionBias<<<grid, block, 0, stream>>>(output,
bias_table,
num_heads,
seq_len,
num_bucket,
is_bidirectional,
max_distance);
return CUDA_CALL(cudaGetLastError());
return CUDA_CALL(cudaGetLastError());
}
}
template Status LaunchRelPosAttnBiasKernel<float>(cudaStream_t stream,
@ -109,7 +136,8 @@ template Status LaunchRelPosAttnBiasKernel<float>(cudaStream_t stream,
const int seq_len,
const int num_bucket,
const int max_distance,
const bool is_bidirectional);
const bool is_bidirectional,
const int max_threads_per_block);
template Status LaunchRelPosAttnBiasKernel<half>(cudaStream_t stream,
half* output,
@ -118,7 +146,8 @@ template Status LaunchRelPosAttnBiasKernel<half>(cudaStream_t stream,
const int seq_len,
const int num_bucket,
const int max_distance,
const bool is_bidirectional);
const bool is_bidirectional,
const int max_threads_per_block);
} // namespace cuda
} // namespace contrib

3
onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.h
View file

@ -18,7 +18,8 @@ Status LaunchRelPosAttnBiasKernel(
const int seq_len,
const int num_bucket,
const int max_distance,
const bool is_bidirectional
const bool is_bidirectional,
const int max_threads_per_block
);
} // namespace cuda