From 342a5bf2b756d1a1fc6fdc582cfeac15182632fe Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Tue, 10 Jan 2023 17:06:42 -0800 Subject: [PATCH] Improve rpb cuda kernel (#14195) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### 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 Review of this comment https://github.com/microsoft/onnxruntime/pull/14149/#discussion_r1063152021 Co-authored-by: Ubuntu --- .../cuda/bert/relative_attn_bias.cc | 4 +- .../cuda/bert/relative_attn_bias_impl.cu | 59 ++++++++++++++----- .../cuda/bert/relative_attn_bias_impl.h | 3 +- 3 files changed, 49 insertions(+), 17 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias.cc b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias.cc index 88ebbfe831..af13efe0e2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias.cc +++ b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias.cc @@ -57,6 +57,7 @@ Status RelPosAttnBias::ComputeInternal(OpKernelContext* context) const { typedef typename ToCudaType::MappedType CudaT; + auto& device_prop = GetDeviceProp(); return LaunchRelPosAttnBiasKernel(Stream(context), reinterpret_cast(output->template MutableData()), reinterpret_cast(bias_table->template Data()), @@ -64,7 +65,8 @@ Status RelPosAttnBias::ComputeInternal(OpKernelContext* context) const { static_cast(query_len), static_cast(num_buckets), max_distance_, - is_bidirectional_); + is_bidirectional_, + device_prop.maxThreadsPerBlock); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu index a30fe5b3db..e333152cb5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu @@ -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<<>>(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<<>>(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<<>>(output, + bias_table, + num_heads, + seq_len, + num_bucket, + is_bidirectional, + max_distance); - buildRelativeAttentionBias<<>>(output, - bias_table, - num_heads, - seq_len, - num_bucket, - is_bidirectional, - max_distance); - - return CUDA_CALL(cudaGetLastError()); + return CUDA_CALL(cudaGetLastError()); + } } template Status LaunchRelPosAttnBiasKernel(cudaStream_t stream, @@ -109,7 +136,8 @@ template Status LaunchRelPosAttnBiasKernel(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(cudaStream_t stream, half* output, @@ -118,7 +146,8 @@ template Status LaunchRelPosAttnBiasKernel(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 diff --git a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.h b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.h index a1efd2755f..5a1a229ab6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.h @@ -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