From 686fd3c22a2b2961cbbf1730d7017eda5dd61858 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 24 Apr 2023 10:02:35 -0700 Subject: [PATCH] Fix cuda 12.1 windows Build (#15614) ### Description Fix CUDA 12.1 Windows build error of cuda namespace ambiguous. Use a new namespace for attention softmax. Tested with VS 2019 and VS 2022 with the following settings: - OS: Microsoft Windows 11 Enterprise (Version 10.0.22621 Build 22621) - CUDA: cuda_12.1.0_531.14_windows - TensorRT: TensorRT-8.6.0.12.Windows10.x86_64.cuda-12.0 - CUDNN: 8.8.1.3 for cuda 12 - Visual Studio Enterprise 2019, version 16.11.26 (MSVC v142) or Visual Studio Enterprise 2022 (64-bit), version 17.5.4 - Python: 3.10 - CMake: 3.25.2 VS 2019: ``` build.bat --cmake_generator "Visual Studio 16 2019" --config Release --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=52;60;61;70;75;80;86" --skip_submodule_sync --parallel --build_shared_lib --update --build --build_dir .\build\trt --use_cuda --cuda_version "12.1" --cuda_home "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1" --cudnn_home "C:\CuDNN\8.8.1.3_cuda12" --use_tensorrt --tensorrt_home "C:\TensorRT-8.6.0.12.Windows10.x86_64.cuda-12.0\TensorRT-8.6.0.12" ``` VS 2022: ``` build.bat --cmake_generator "Visual Studio 17 2022" --config Release --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=52;60;61;70;75;80;86" --skip_submodule_sync --parallel --build_shared_lib --update --build --build_dir .\build\trt_2022 --use_cuda --cuda_version "12.1" --cuda_home "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1" --cudnn_home "C:\CuDNN\8.8.1.3_cuda12" --use_tensorrt --tensorrt_home "C:\TensorRT-8.6.0.12.Windows10.x86_64.cuda-12.0\TensorRT-8.6.0.12" ``` ### Motivation and Context https://github.com/microsoft/onnxruntime/issues/15242 --- cmake/onnxruntime_rocm_hipify.cmake | 1 + .../contrib_ops/cuda/bert/attention_impl.cu | 2 +- .../cuda/bert/attention_softmax.cu | 998 ++++++++++++++++++ .../contrib_ops/cuda/bert/attention_softmax.h | 859 +-------------- .../cuda/bert/packed_attention_impl.cu | 2 +- 5 files changed, 1009 insertions(+), 853 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu diff --git a/cmake/onnxruntime_rocm_hipify.cmake b/cmake/onnxruntime_rocm_hipify.cmake index 496dac5f85..caaa89c84f 100644 --- a/cmake/onnxruntime_rocm_hipify.cmake +++ b/cmake/onnxruntime_rocm_hipify.cmake @@ -9,6 +9,7 @@ set(contrib_ops_excluded_files "bert/attention.h" "bert/attention_impl.cu" "bert/attention_softmax.h" + "bert/attention_softmax.cu" "bert/decoder_masked_multihead_attention.h" "bert/decoder_masked_multihead_attention.cc" "bert/decoder_masked_self_attention.h" diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 7a7818f9af..d6741d68f3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -44,7 +44,7 @@ limitations under the License. #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" using namespace onnxruntime::cuda; -using namespace cub; +using namespace onnxruntime::contrib::attention_softmax_cuda; namespace onnxruntime { namespace contrib { diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu new file mode 100644 index 0000000000..d50cba24b0 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu @@ -0,0 +1,998 @@ +/* + The implementation of this file is based on qkvToContext plugin in TensorRT demo: + https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/ + +Copyright 2019 NVIDIA Corporation + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +#include +#include +#include +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/math/softmax.h" +#include "contrib_ops/cuda/bert/attention_softmax.h" +using namespace onnxruntime::cuda; + +namespace onnxruntime { +namespace contrib { +namespace attention_softmax_cuda { + +template +__device__ inline void Softmax(const int all_sequence_length, + const int valid_end, + const int valid_start, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + float thread_data_max(-CUDART_INF_F); + + const bool no_rpb = (rel_pos_bias == nullptr); + + // e^x is represented as infinity if x is large enough, like 100.f. + // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. + // a math transform as below is leveraged to get a stable softmax: + // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; + const int size_per_batch = gridDim.x * all_sequence_length; + for (int i = threadIdx.x; i < valid_end; i += TPB) { + if (i >= valid_start) { + const int index = offset + i; + float input_at_idx = no_rpb + ? float(input[index]) + : float(input[index] + (broadcast_rel_pos_bias + ? rel_pos_bias[index % size_per_batch] + : rel_pos_bias[index])); + if (thread_data_max < input_at_idx) { + thread_data_max = input_at_idx; + } + } + } + + const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max()); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float thread_data_sum(0.f); + for (int i = threadIdx.x; i < valid_end; i += TPB) { + if (i >= valid_start) { + const int index = offset + i; + float val = no_rpb ? input[index] : input[index] + rel_pos_bias[index % size_per_batch]; + thread_data_sum += expf(val - max_block); + } + } + + const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_sum, cub::Sum()); + if (threadIdx.x == 0) { + sum_reverse_block = 1.f / sum; + } + __syncthreads(); + + for (int i = threadIdx.x; i < all_sequence_length; i += TPB) { + const int index = offset + i; + float input_at_idx = no_rpb ? float(input[index]) : float(input[index] + rel_pos_bias[index % size_per_batch]); + const float val = (i >= valid_start && i < valid_end) ? expf(input_at_idx - max_block) * sum_reverse_block : 0.f; + output[index] = T(val); + } +} + +template +__device__ inline void SoftmaxSmall(const int all_sequence_length, + const int sequence_length, + const int valid_end, + const int valid_start, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + bool is_unidirectional) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; + const int index = offset + threadIdx.x; + + bool is_valid = false; // whether it has attention mask == 1. + + // Update end position for unidirectional. + int end = valid_end; + if (is_unidirectional) { + int end_unid = all_sequence_length - sequence_length + (blockIdx.x % sequence_length) + 1; + if (end_unid <= valid_start) { + // In this situation, mask of [0, end_unid) and [valid_start, valid_end) has -10000, + // and [end_unid, valid_start) and [valid_end, all_seq_len) has -20000. + // So [0, end_unid) will also have value after softmax. + is_valid = threadIdx.x < end_unid; + } else { + end = min(valid_end, end_unid); + } + } + + is_valid = is_valid || (threadIdx.x >= valid_start && threadIdx.x < end); + + // e^x is represented as infinity if x is large enough, like 100.f. + // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. + // a math transform as below is leveraged to get a stable softmax: + // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) + const bool no_rpb = (rel_pos_bias == nullptr); + const int size_per_batch = gridDim.x * all_sequence_length; + float input_data = no_rpb + ? float(input[index]) + : float(input[index] + (broadcast_rel_pos_bias + ? rel_pos_bias[index % size_per_batch] + : rel_pos_bias[index])); + float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F); + const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float thread_data_exp(0.f); + if (is_valid) { + thread_data_exp = expf(input_data - max_block); + } + + const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); + + // Store value of 1.0/sum. + if (threadIdx.x == 0) { + sum_reverse_block = (1.f) / sum; + } + __syncthreads(); + + // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. + if (threadIdx.x < all_sequence_length) { + output[index] = T(thread_data_exp * sum_reverse_block); + } +} + +template +__global__ void SoftmaxLargeKernel(const int all_sequence_length, + const int sequence_length, + const int valid_end, + const int valid_start, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + bool is_unidirectional) { + extern __shared__ float cached_data[]; // float[all_sequence_length] + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + // Update end position for unidirectional. + int end = valid_end; + int end_unid = -1; + if (is_unidirectional) { + end_unid = all_sequence_length - sequence_length + (blockIdx.x % sequence_length) + 1; + if (end_unid <= valid_start) { + ; + // In this situation, mask of [0, end_unid) and [valid_start, valid_end) has -10000, + // and [end_unid, valid_start) and [valid_end, all_seq_len) has -20000. + // So [0, end_unid) will also have value after softmax. + // KEEP SMALL KERNEL CODE LOGIC HERE as COMMENT + // is_valid = threadIdx.x < end_unid; // is_valid initialized with false + } else { + end = min(valid_end, end_unid); + } + } + + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; + const int size_per_batch = gridDim.x * all_sequence_length; + + float thread_data_max = -CUDART_INF_F; + for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { + const int index = offset + seq_idx; + bool is_valid = (seq_idx < end_unid) || (seq_idx >= valid_start && seq_idx < end); + + // e^x is represented as infinity if x is large enough, like 100.f. + // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. + // a math transform as below is leveraged to get a stable softmax: + // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) + float input_data = is_valid + ? (rel_pos_bias + ? float(input[index] + (broadcast_rel_pos_bias + ? rel_pos_bias[index % size_per_batch] + : rel_pos_bias[index])) + : float(input[index])) + : float(-CUDART_INF_F); + cached_data[seq_idx] = input_data; + thread_data_max = max(thread_data_max, input_data); + } + const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float thread_data_exp(0.f); + for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { + bool is_valid = (seq_idx < end_unid) || (seq_idx >= valid_start && seq_idx < end); + cached_data[seq_idx] = is_valid ? expf(cached_data[seq_idx] - max_block) : 0.0f; + thread_data_exp += cached_data[seq_idx]; + } + const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); + + // Store value of 1.0/sum. + if (threadIdx.x == 0) { + sum_reverse_block = (1.f) / sum; + } + __syncthreads(); + + // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. + for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { + output[offset + seq_idx] = T(cached_data[seq_idx] * sum_reverse_block); + } +} + +template +__global__ void SoftmaxWithRawMaskLargeKernel(const int all_sequence_length, + const int sequence_length, + const int* attention_mask, // 2D, 3D or 4D attention mask + const bool* key_padding_mask, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + const bool is_unidirectional, + const float rsqrt_head_size, + const int mask_dimension, + const int max_sequence_length, + const bool skip_softmax, + const float mask_filter_value) { + extern __shared__ float cached_data[]; // float[all_sequence_length] + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + float max_thread_data = -CUDART_INF_F; + const int size_per_batch = gridDim.x * all_sequence_length; + + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + int base_index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; + for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { + float thread_data = -CUDART_INF_F; + int index = base_index + seq_idx; + if (rel_pos_bias == nullptr) { + thread_data = float(input[index]) * rsqrt_head_size; + } else { + T rel_pos_bias_value = broadcast_rel_pos_bias ? rel_pos_bias[index % size_per_batch] : rel_pos_bias[index]; + thread_data = float(input[index] + rel_pos_bias_value) * rsqrt_head_size; + } + + const int sequence_index = blockIdx.x % sequence_length; + if (is_unidirectional) { + int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length. + if (seq_idx > from_index) { + thread_data = mask_filter_value; + } + } + + int mask_offset = 0; + const int batch_index = blockIdx.y; + if (mask_dimension == 2) { + mask_offset = batch_index * all_sequence_length + seq_idx; + } else if (mask_dimension == 3) { + mask_offset = (batch_index * sequence_length + sequence_index) * all_sequence_length + seq_idx; + } else if (mask_dimension == 4) { + int from_index = all_sequence_length - sequence_length + sequence_index; + mask_offset = (batch_index * max_sequence_length + from_index) * max_sequence_length + seq_idx; + } + + if (nullptr == key_padding_mask) { + const int& mask = attention_mask[mask_offset]; + if (mask == 0) + thread_data += mask_filter_value; + } else { + const bool mask = key_padding_mask[mask_offset]; + if (mask) { + thread_data = -CUDART_INF_F; + } + } + + if (skip_softmax) { + output[index] = T(thread_data); + } + cached_data[seq_idx] = thread_data; + max_thread_data = max(max_thread_data, thread_data); + } + + if (skip_softmax) { + return; + } + + const float max = BlockReduce(tmp_storage).Reduce(max_thread_data, cub::Max(), all_sequence_length); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float sum_thread_data_exp = 0.0f; + for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { + auto ev = expf(cached_data[seq_idx] - max_block); + cached_data[seq_idx] = ev; + sum_thread_data_exp += ev; + } + const auto sum = BlockReduce(tmp_storage).Reduce(sum_thread_data_exp, cub::Sum(), TPB); + + // Store value of 1.0/sum + if (threadIdx.x == 0) { + sum_reverse_block = (1.f) / sum; + } + __syncthreads(); + + for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { + output[base_index + seq_idx] = T(cached_data[seq_idx] * sum_reverse_block); + } +} + +template +__device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length, + const int sequence_length, + const int* attention_mask, // 2D, 3D or 4D attention mask + const bool* key_padding_mask, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + const bool is_unidirectional, + const float rsqrt_head_size, + const int mask_dimension, + const int max_sequence_length, + const bool skip_softmax, + const float mask_filter_value) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + int index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length + threadIdx.x; + const int size_per_batch = gridDim.x * all_sequence_length; + + float thread_data = -CUDART_INF_F; + if (threadIdx.x < all_sequence_length) { + thread_data = float(input[index]) * rsqrt_head_size; + + const int sequence_index = blockIdx.x % sequence_length; + if (is_unidirectional) { + int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length. + if (threadIdx.x > from_index) { + thread_data = mask_filter_value; + } + } + + int mask_offset = 0; + const int batch_index = blockIdx.y; + if (mask_dimension == 2) { + mask_offset = batch_index * all_sequence_length + threadIdx.x; + } else if (mask_dimension == 3) { + mask_offset = (batch_index * sequence_length + sequence_index) * all_sequence_length + threadIdx.x; + } else if (mask_dimension == 4) { + int from_index = all_sequence_length - sequence_length + sequence_index; + mask_offset = (batch_index * max_sequence_length + from_index) * max_sequence_length + threadIdx.x; + } + + if (nullptr == key_padding_mask) { + const int& mask = attention_mask[mask_offset]; + if (mask == 0) + thread_data += mask_filter_value; + } else { + const bool mask = key_padding_mask[mask_offset]; + if (mask) { + thread_data = -CUDART_INF_F; + } + } + + if (rel_pos_bias != nullptr) { + float bias = broadcast_rel_pos_bias ? float(rel_pos_bias[index % size_per_batch]) : float(rel_pos_bias[index]); + thread_data += bias; + } + } + + if (skip_softmax) { + if (threadIdx.x < all_sequence_length) { + output[index] = T(thread_data); + } + return; + } + + const float max = BlockReduce(tmp_storage).Reduce(thread_data, cub::Max(), all_sequence_length); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float thread_data_exp = threadIdx.x < all_sequence_length ? expf(thread_data - max_block) : 0.0f; + const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), all_sequence_length); + + // Store value of 1.0/sum + if (threadIdx.x == 0) { + sum_reverse_block = (1.f) / sum; + } + __syncthreads(); + + if (threadIdx.x < all_sequence_length) { + output[index] = T(thread_data_exp * sum_reverse_block); + } +} + +template +__global__ void SoftmaxKernelSmall(const int all_sequence_length, + const int sequence_length, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + bool is_unidirectional) { + SoftmaxSmall(all_sequence_length, sequence_length, all_sequence_length, 0, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); +} + +template +__global__ void SoftmaxKernel(const int all_sequence_length, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output) { + Softmax(all_sequence_length, all_sequence_length, 0, + rel_pos_bias, broadcast_rel_pos_bias, input, output); +} + +template +Status ComputeSoftmax(cudaStream_t stream, const int all_sequence_length, const int sequence_length, + const int batch_size, const int num_heads, const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, T* input, T* output, bool is_unidirectional) { + const dim3 grid(sequence_length * num_heads, batch_size, 1); + if (all_sequence_length <= 32) { + const int blockSize = 32; + SoftmaxKernelSmall<<>>( + all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 64) { + const int blockSize = 64; + SoftmaxKernelSmall<<>>( + all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 128) { + const int blockSize = 128; + SoftmaxKernelSmall<<>>( + all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 256) { + const int blockSize = 256; + SoftmaxKernelSmall<<>>( + all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 512) { + const int blockSize = 512; + SoftmaxKernelSmall<<>>( + all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 1024) { + const int blockSize = 1024; + SoftmaxKernelSmall<<>>( + all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (!is_unidirectional) { + const int blockSize = 1024; + SoftmaxKernel<<>>( + all_sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output); + } else { + const int blockSize = 256; + const int sh_bytes = sizeof(float) * all_sequence_length; + SoftmaxLargeKernel<<>>( + all_sequence_length, sequence_length, all_sequence_length, 0, rel_pos_bias, broadcast_rel_pos_bias, + input, output, true); + } + + return CUDA_CALL(cudaGetLastError()); +} + +template +__global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length, + const int sequence_length, + const int* mask_end, + const int* mask_start, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + bool is_unidirectional) { + __shared__ int start_position; + __shared__ int end_position; + + if (threadIdx.x == 0) { + const int batch = blockIdx.y; + start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0; + end_position = min(all_sequence_length, mask_end[batch]); + + // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result. + if (start_position >= end_position) { + start_position = 0; + end_position = all_sequence_length; + } + } + __syncthreads(); + + SoftmaxSmall(all_sequence_length, sequence_length, end_position, start_position, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); +} + +template +__device__ inline void SoftmaxSmallPacked(const int sequence_length, + const int end, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length; + const int index = offset + threadIdx.x; + + bool is_valid = threadIdx.x < end; + + // e^x is represented as infinity if x is large enough, like 100.f. + // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. + // a math transform as below is leveraged to get a stable softmax: + // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) + const bool no_rpb = (rel_pos_bias == nullptr); + const int size_per_batch = gridDim.x * sequence_length; + float input_data = no_rpb + ? float(input[index]) + : float(input[index] + (broadcast_rel_pos_bias + ? rel_pos_bias[index % size_per_batch] + : rel_pos_bias[index])); + + float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F); + const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float thread_data_exp(0.f); + if (is_valid) { + thread_data_exp = expf(input_data - max_block); + } + + const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); + + // Store value of 1.0/sum. + if (threadIdx.x == 0) { + sum_reverse_block = (1.f) / sum; + } + __syncthreads(); + + // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. + if (threadIdx.x < sequence_length) { + output[index] = T(thread_data_exp * sum_reverse_block); + } +} + +template +__global__ void SoftmaxKernelSmallWithCumSeqLen(const T* input, + const T* rel_pos_bias, const bool broadcast_rel_pos_bias, + const int* cum_seq_length, const int sequence_length, + T* output) { + __shared__ int end_position; + + if (threadIdx.x == 0) { + const int batch = blockIdx.y; + end_position = cum_seq_length[batch + 1] - cum_seq_length[batch]; + } + __syncthreads(); + + SoftmaxSmallPacked(sequence_length, end_position, + rel_pos_bias, broadcast_rel_pos_bias, + input, output); +} + +template +__global__ void SoftmaxKernelWithCumSeqLen(const T* input, + const T* rel_pos_bias, const bool broadcast_rel_pos_bias, + const int* cum_seq_length, const int sequence_length, + T* output) { + __shared__ int end_position; + + if (threadIdx.x == 0) { + const int batch = blockIdx.y; + end_position = cum_seq_length[batch + 1] - cum_seq_length[batch]; + } + __syncthreads(); + + Softmax(sequence_length, end_position, 0 /*start_position*/, + rel_pos_bias, broadcast_rel_pos_bias, input, output); +} + +template +__global__ void MaskedSoftmaxKernel(const int all_sequence_length, + const int* mask_end, + const int* mask_start, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, T* output) { + __shared__ int start_position; + __shared__ int end_position; + + if (threadIdx.x == 0) { + const int batch = blockIdx.y; + start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0; + end_position = min(all_sequence_length, mask_end[batch]); + + // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result. + if (start_position >= end_position) { + start_position = 0; + end_position = all_sequence_length; + } + } + __syncthreads(); + + Softmax(all_sequence_length, end_position, start_position, + rel_pos_bias, broadcast_rel_pos_bias, input, output); +} + +template +__global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length, + const int sequence_length, + const int* attention_mask, + const bool* key_padding_mask, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + const bool is_unidirectional, + const float rsqrt_head_size, + const int mask_dimension, + const int max_sequence_length, + const bool skip_softmax, + const float mask_filter_value) { + SoftmaxWithRawMaskSmall( + all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, output, + is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + skip_softmax, mask_filter_value); +} + +template +Status ComputeSoftmaxWithCumSeqLength( + const T* input, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const int32_t* cum_seq_length, + const int batch_size, + const int sequence_length, + const int num_heads, + T* output, cudaStream_t stream) { + const dim3 grid(sequence_length * num_heads, batch_size, 1); + + if (sequence_length <= 32) { + const int blockSize = 32; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + + } else if (sequence_length <= 64) { + const int blockSize = 64; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 128) { + const int blockSize = 128; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 256) { + const int blockSize = 256; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 512) { + const int blockSize = 512; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 1024) { + const int blockSize = 1024; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else { + SoftmaxKernelWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } + + return CUDA_CALL(cudaGetLastError()); +} + +template +Status ComputeSoftmaxWithMask1D(cudaStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int num_heads, + const int* mask_index, + const int* mask_start, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + const bool is_unidirectional) { + const dim3 grid(sequence_length * num_heads, batch_size, 1); + + if (all_sequence_length <= 32) { + const int blockSize = 32; + MaskedSoftmaxKernelSmall + <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 64) { + const int blockSize = 64; + MaskedSoftmaxKernelSmall + <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 128) { + const int blockSize = 128; + MaskedSoftmaxKernelSmall + <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 256) { + const int blockSize = 256; + MaskedSoftmaxKernelSmall + <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 512) { + const int blockSize = 512; + MaskedSoftmaxKernelSmall + <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (all_sequence_length <= 1024) { + const int blockSize = 1024; + MaskedSoftmaxKernelSmall + <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); + } else if (!is_unidirectional) { + const int blockSize = 1024; + MaskedSoftmaxKernel + <<>>(all_sequence_length, mask_index, mask_start, + rel_pos_bias, broadcast_rel_pos_bias, input, output); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024."); + } + + return CUDA_CALL(cudaGetLastError()); +} + +template +Status ComputeSoftmaxWithRawMask(cudaStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int num_heads, + const int* attention_mask, + const bool* key_padding_mask, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output, + const bool is_unidirectional, + const float rsqrt_head_size, + const int mask_dimension, + const int max_sequence_length, + const bool use_persistent_softmax, + T* persistent_softmax_workspace, + const float mask_filter_value) { + const dim3 grid(sequence_length * num_heads, batch_size, 1); + + T* out = use_persistent_softmax ? persistent_softmax_workspace : output; + if (all_sequence_length <= 32) { + const int blockSize = 32; + SoftmaxWithRawMaskSmallKernel + <<>>(all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } else if (all_sequence_length <= 64) { + const int blockSize = 64; + SoftmaxWithRawMaskSmallKernel + <<>>(all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } else if (all_sequence_length <= 128) { + const int blockSize = 128; + SoftmaxWithRawMaskSmallKernel + <<>>(all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } else if (all_sequence_length <= 256) { + const int blockSize = 256; + SoftmaxWithRawMaskSmallKernel + <<>>(all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } else if (all_sequence_length <= 512) { + const int blockSize = 512; + SoftmaxWithRawMaskSmallKernel + <<>>(all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } else if (all_sequence_length <= 1024) { + const int blockSize = 1024; + SoftmaxWithRawMaskSmallKernel + <<>>(all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } else { + const int blockSize = 256; + const int sh_bytes = sizeof(float) * all_sequence_length; + SoftmaxWithRawMaskLargeKernel + <<>>( + all_sequence_length, sequence_length, + attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, + out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, + use_persistent_softmax, mask_filter_value); + } + + if (use_persistent_softmax) { + return onnxruntime::cuda::dispatch_warpwise_softmax_forward( + stream, + output, + persistent_softmax_workspace, + all_sequence_length, + all_sequence_length, + batch_size * num_heads * sequence_length); + } + + return CUDA_CALL(cudaGetLastError()); +} + +// Template Instantiation +template Status ComputeSoftmax( + cudaStream_t stream, const int all_sequence_length, const int sequence_length, + const int batch_size, const int num_heads, const float* rel_pos_bias, + const bool broadcast_rel_pos_bias, float* input, float* output, bool is_unidirectional); + +template Status ComputeSoftmax( + cudaStream_t stream, const int all_sequence_length, const int sequence_length, + const int batch_size, const int num_heads, const half* rel_pos_bias, + const bool broadcast_rel_pos_bias, half* input, half* output, bool is_unidirectional); + +template Status ComputeSoftmaxWithCumSeqLength( + const float* input, + const float* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const int32_t* cum_seq_length, + const int batch_size, + const int sequence_length, + const int num_heads, + float* output, cudaStream_t stream); + +template Status ComputeSoftmaxWithCumSeqLength( + const half* input, + const half* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const int32_t* cum_seq_length, + const int batch_size, + const int sequence_length, + const int num_heads, + half* output, cudaStream_t stream); + +template Status ComputeSoftmaxWithMask1D(cudaStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int num_heads, + const int* mask_index, + const int* mask_start, + const float* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const float* input, + float* output, + const bool is_unidirectional); + +template Status ComputeSoftmaxWithMask1D(cudaStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int num_heads, + const int* mask_index, + const int* mask_start, + const half* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const half* input, + half* output, + const bool is_unidirectional); + +template Status ComputeSoftmaxWithRawMask(cudaStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int num_heads, + const int* attention_mask, + const bool* key_padding_mask, + const float* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const float* input, + float* output, + const bool is_unidirectional, + const float rsqrt_head_size, + const int mask_dimension, + const int max_sequence_length, + const bool use_persistent_softmax, + float* persistent_softmax_workspace, + const float mask_filter_value); + +template Status ComputeSoftmaxWithRawMask(cudaStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int num_heads, + const int* attention_mask, + const bool* key_padding_mask, + const half* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const half* input, + half* output, + const bool is_unidirectional, + const float rsqrt_head_size, + const int mask_dimension, + const int max_sequence_length, + const bool use_persistent_softmax, + half* persistent_softmax_workspace, + const float mask_filter_value); + +} // namespace attention_softmax_cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h index 06fd8aa518..4af9e4d695 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h @@ -1,709 +1,16 @@ -/* - The implementation of this file is based on qkvToContext plugin in TensorRT demo: - https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/ - -Copyright 2019 NVIDIA Corporation - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. #pragma once -#include -#include -#include -#include -#include "core/providers/cuda/cu_inc/common.cuh" -#include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/math/softmax.h" - -using namespace onnxruntime::cuda; -using namespace cub; - namespace onnxruntime { namespace contrib { -namespace cuda { - -template -__device__ inline void Softmax(const int all_sequence_length, - const int valid_end, - const int valid_start, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tmp_storage; - - __shared__ float sum_reverse_block; - __shared__ float max_block; - - float thread_data_max(-CUDART_INF_F); - - const bool no_rpb = (rel_pos_bias == nullptr); - - // e^x is represented as infinity if x is large enough, like 100.f. - // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. - // a math transform as below is leveraged to get a stable softmax: - // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) - const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; - const int size_per_batch = gridDim.x * all_sequence_length; - for (int i = threadIdx.x; i < valid_end; i += TPB) { - if (i >= valid_start) { - const int index = offset + i; - float input_at_idx = no_rpb - ? float(input[index]) - : float(input[index] + (broadcast_rel_pos_bias - ? rel_pos_bias[index % size_per_batch] - : rel_pos_bias[index])); - if (thread_data_max < input_at_idx) { - thread_data_max = input_at_idx; - } - } - } - - const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max()); - - // Store max value - if (threadIdx.x == 0) { - max_block = max; - } - __syncthreads(); - - float thread_data_sum(0.f); - for (int i = threadIdx.x; i < valid_end; i += TPB) { - if (i >= valid_start) { - const int index = offset + i; - float val = no_rpb ? input[index] : input[index] + rel_pos_bias[index % size_per_batch]; - thread_data_sum += expf(val - max_block); - } - } - - const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_sum, cub::Sum()); - if (threadIdx.x == 0) { - sum_reverse_block = 1.f / sum; - } - __syncthreads(); - - for (int i = threadIdx.x; i < all_sequence_length; i += TPB) { - const int index = offset + i; - float input_at_idx = no_rpb ? float(input[index]) : float(input[index] + rel_pos_bias[index % size_per_batch]); - const float val = (i >= valid_start && i < valid_end) ? expf(input_at_idx - max_block) * sum_reverse_block : 0.f; - output[index] = T(val); - } -} - -template -__device__ inline void SoftmaxSmall(const int all_sequence_length, - const int sequence_length, - const int valid_end, - const int valid_start, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - bool is_unidirectional) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tmp_storage; - - __shared__ float sum_reverse_block; - __shared__ float max_block; - - // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; - const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; - const int index = offset + threadIdx.x; - - bool is_valid = false; // whether it has attention mask == 1. - - // Update end position for unidirectional. - int end = valid_end; - if (is_unidirectional) { - int end_unid = all_sequence_length - sequence_length + (blockIdx.x % sequence_length) + 1; - if (end_unid <= valid_start) { - // In this situation, mask of [0, end_unid) and [valid_start, valid_end) has -10000, - // and [end_unid, valid_start) and [valid_end, all_seq_len) has -20000. - // So [0, end_unid) will also have value after softmax. - is_valid = threadIdx.x < end_unid; - } else { - end = min(valid_end, end_unid); - } - } - - is_valid = is_valid || (threadIdx.x >= valid_start && threadIdx.x < end); - - // e^x is represented as infinity if x is large enough, like 100.f. - // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. - // a math transform as below is leveraged to get a stable softmax: - // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) - const bool no_rpb = (rel_pos_bias == nullptr); - const int size_per_batch = gridDim.x * all_sequence_length; - float input_data = no_rpb - ? float(input[index]) - : float(input[index] + (broadcast_rel_pos_bias - ? rel_pos_bias[index % size_per_batch] - : rel_pos_bias[index])); - float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F); - const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); - - // Store max value - if (threadIdx.x == 0) { - max_block = max; - } - __syncthreads(); - - float thread_data_exp(0.f); - if (is_valid) { - thread_data_exp = expf(input_data - max_block); - } - - const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); - - // Store value of 1.0/sum. - if (threadIdx.x == 0) { - sum_reverse_block = (1.f) / sum; - } - __syncthreads(); - - // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. - if (threadIdx.x < all_sequence_length) { - output[index] = T(thread_data_exp * sum_reverse_block); - } -} - -template -__global__ void SoftmaxLargeKernel(const int all_sequence_length, - const int sequence_length, - const int valid_end, - const int valid_start, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - bool is_unidirectional) { - extern __shared__ float cached_data[]; // float[all_sequence_length] - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tmp_storage; - - __shared__ float sum_reverse_block; - __shared__ float max_block; - - // Update end position for unidirectional. - int end = valid_end; - int end_unid = -1; - if (is_unidirectional) { - end_unid = all_sequence_length - sequence_length + (blockIdx.x % sequence_length) + 1; - if (end_unid <= valid_start) { - ; - // In this situation, mask of [0, end_unid) and [valid_start, valid_end) has -10000, - // and [end_unid, valid_start) and [valid_end, all_seq_len) has -20000. - // So [0, end_unid) will also have value after softmax. - // KEEP SMALL KERNEL CODE LOGIC HERE as COMMENT - // is_valid = threadIdx.x < end_unid; // is_valid initialized with false - } else { - end = min(valid_end, end_unid); - } - } - - // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; - const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; - const int size_per_batch = gridDim.x * all_sequence_length; - - float thread_data_max = -CUDART_INF_F; - for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { - const int index = offset + seq_idx; - bool is_valid = (seq_idx < end_unid) || (seq_idx >= valid_start && seq_idx < end); - - // e^x is represented as infinity if x is large enough, like 100.f. - // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. - // a math transform as below is leveraged to get a stable softmax: - // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) - float input_data = is_valid - ? (rel_pos_bias - ? float(input[index] + (broadcast_rel_pos_bias - ? rel_pos_bias[index % size_per_batch] - : rel_pos_bias[index])) - : float(input[index])) - : float(-CUDART_INF_F); - cached_data[seq_idx] = input_data; - thread_data_max = max(thread_data_max, input_data); - } - const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); - - // Store max value - if (threadIdx.x == 0) { - max_block = max; - } - __syncthreads(); - - float thread_data_exp(0.f); - for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { - bool is_valid = (seq_idx < end_unid) || (seq_idx >= valid_start && seq_idx < end); - cached_data[seq_idx] = is_valid ? expf(cached_data[seq_idx] - max_block) : 0.0f; - thread_data_exp += cached_data[seq_idx]; - } - const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); - - // Store value of 1.0/sum. - if (threadIdx.x == 0) { - sum_reverse_block = (1.f) / sum; - } - __syncthreads(); - - // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. - for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { - output[offset + seq_idx] = T(cached_data[seq_idx] * sum_reverse_block); - } -} - -template -__global__ void SoftmaxWithRawMaskLargeKernel(const int all_sequence_length, - const int sequence_length, - const int* attention_mask, // 2D, 3D or 4D attention mask - const bool* key_padding_mask, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - const bool is_unidirectional, - const float rsqrt_head_size, - const int mask_dimension, - const int max_sequence_length, - const bool skip_softmax, - const float mask_filter_value) { - extern __shared__ float cached_data[]; // float[all_sequence_length] - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tmp_storage; - - __shared__ float sum_reverse_block; - __shared__ float max_block; - - float max_thread_data = -CUDART_INF_F; - const int size_per_batch = gridDim.x * all_sequence_length; - - // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; - int base_index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; - for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { - float thread_data = -CUDART_INF_F; - int index = base_index + seq_idx; - if (rel_pos_bias == nullptr) { - thread_data = float(input[index]) * rsqrt_head_size; - } else { - T rel_pos_bias_value = broadcast_rel_pos_bias ? rel_pos_bias[index % size_per_batch] : rel_pos_bias[index]; - thread_data = float(input[index] + rel_pos_bias_value) * rsqrt_head_size; - } - - const int sequence_index = blockIdx.x % sequence_length; - if (is_unidirectional) { - int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length. - if (seq_idx > from_index) { - thread_data = mask_filter_value; - } - } - - int mask_offset = 0; - const int batch_index = blockIdx.y; - if (mask_dimension == 2) { - mask_offset = batch_index * all_sequence_length + seq_idx; - } else if (mask_dimension == 3) { - mask_offset = (batch_index * sequence_length + sequence_index) * all_sequence_length + seq_idx; - } else if (mask_dimension == 4) { - int from_index = all_sequence_length - sequence_length + sequence_index; - mask_offset = (batch_index * max_sequence_length + from_index) * max_sequence_length + seq_idx; - } - - if (nullptr == key_padding_mask) { - const int& mask = attention_mask[mask_offset]; - if (mask == 0) - thread_data += mask_filter_value; - } else { - const bool mask = key_padding_mask[mask_offset]; - if (mask) { - thread_data = -CUDART_INF_F; - } - } - - if (skip_softmax) { - output[index] = T(thread_data); - } - cached_data[seq_idx] = thread_data; - max_thread_data = max(max_thread_data, thread_data); - } - - if (skip_softmax) { - return; - } - - const float max = BlockReduce(tmp_storage).Reduce(max_thread_data, cub::Max(), all_sequence_length); - - // Store max value - if (threadIdx.x == 0) { - max_block = max; - } - __syncthreads(); - - float sum_thread_data_exp = 0.0f; - for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { - auto ev = expf(cached_data[seq_idx] - max_block); - cached_data[seq_idx] = ev; - sum_thread_data_exp += ev; - } - const auto sum = BlockReduce(tmp_storage).Reduce(sum_thread_data_exp, cub::Sum(), TPB); - - // Store value of 1.0/sum - if (threadIdx.x == 0) { - sum_reverse_block = (1.f) / sum; - } - __syncthreads(); - - for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) { - output[base_index + seq_idx] = T(cached_data[seq_idx] * sum_reverse_block); - } -} - -template -__device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length, - const int sequence_length, - const int* attention_mask, // 2D, 3D or 4D attention mask - const bool* key_padding_mask, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - const bool is_unidirectional, - const float rsqrt_head_size, - const int mask_dimension, - const int max_sequence_length, - const bool skip_softmax, - const float mask_filter_value) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tmp_storage; - - __shared__ float sum_reverse_block; - __shared__ float max_block; - - // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; - int index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length + threadIdx.x; - const int size_per_batch = gridDim.x * all_sequence_length; - - float thread_data = -CUDART_INF_F; - if (threadIdx.x < all_sequence_length) { - thread_data = float(input[index]) * rsqrt_head_size; - - const int sequence_index = blockIdx.x % sequence_length; - if (is_unidirectional) { - int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length. - if (threadIdx.x > from_index) { - thread_data = mask_filter_value; - } - } - - int mask_offset = 0; - const int batch_index = blockIdx.y; - if (mask_dimension == 2) { - mask_offset = batch_index * all_sequence_length + threadIdx.x; - } else if (mask_dimension == 3) { - mask_offset = (batch_index * sequence_length + sequence_index) * all_sequence_length + threadIdx.x; - } else if (mask_dimension == 4) { - int from_index = all_sequence_length - sequence_length + sequence_index; - mask_offset = (batch_index * max_sequence_length + from_index) * max_sequence_length + threadIdx.x; - } - - if (nullptr == key_padding_mask) { - const int& mask = attention_mask[mask_offset]; - if (mask == 0) - thread_data += mask_filter_value; - } else { - const bool mask = key_padding_mask[mask_offset]; - if (mask) { - thread_data = -CUDART_INF_F; - } - } - - if (rel_pos_bias != nullptr) { - float rel_pos_bias_value = broadcast_rel_pos_bias ? float(rel_pos_bias[index % size_per_batch]) : float(rel_pos_bias[index]); - thread_data += rel_pos_bias_value; - } - } - - if (skip_softmax) { - if (threadIdx.x < all_sequence_length) { - output[index] = T(thread_data); - } - return; - } - - const float max = BlockReduce(tmp_storage).Reduce(thread_data, cub::Max(), all_sequence_length); - - // Store max value - if (threadIdx.x == 0) { - max_block = max; - } - __syncthreads(); - - float thread_data_exp = threadIdx.x < all_sequence_length ? expf(thread_data - max_block) : 0.0f; - const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), all_sequence_length); - - // Store value of 1.0/sum - if (threadIdx.x == 0) { - sum_reverse_block = (1.f) / sum; - } - __syncthreads(); - - if (threadIdx.x < all_sequence_length) { - output[index] = T(thread_data_exp * sum_reverse_block); - } -} - -template -__global__ void SoftmaxKernelSmall(const int all_sequence_length, - const int sequence_length, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - bool is_unidirectional) { - SoftmaxSmall(all_sequence_length, sequence_length, all_sequence_length, 0, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); -} - -template -__global__ void SoftmaxKernel(const int all_sequence_length, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output) { - Softmax(all_sequence_length, all_sequence_length, 0, - rel_pos_bias, broadcast_rel_pos_bias, input, output); -} +namespace attention_softmax_cuda { template Status ComputeSoftmax(cudaStream_t stream, const int all_sequence_length, const int sequence_length, const int batch_size, const int num_heads, const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, T* input, T* output, bool is_unidirectional) { - const dim3 grid(sequence_length * num_heads, batch_size, 1); - if (all_sequence_length <= 32) { - const int blockSize = 32; - SoftmaxKernelSmall<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 64) { - const int blockSize = 64; - SoftmaxKernelSmall<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 128) { - const int blockSize = 128; - SoftmaxKernelSmall<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 256) { - const int blockSize = 256; - SoftmaxKernelSmall<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 512) { - const int blockSize = 512; - SoftmaxKernelSmall<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 1024) { - const int blockSize = 1024; - SoftmaxKernelSmall<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (!is_unidirectional) { - const int blockSize = 1024; - SoftmaxKernel<<>>( - all_sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output); - } else { - const int blockSize = 256; - const int sh_bytes = sizeof(float) * all_sequence_length; - SoftmaxLargeKernel<<>>( - all_sequence_length, sequence_length, all_sequence_length, 0, rel_pos_bias, broadcast_rel_pos_bias, - input, output, true); - } - - return CUDA_CALL(cudaGetLastError()); -} - -template -__global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length, - const int sequence_length, - const int* mask_end, - const int* mask_start, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - bool is_unidirectional) { - __shared__ int start_position; - __shared__ int end_position; - - if (threadIdx.x == 0) { - const int batch = blockIdx.y; - start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0; - end_position = min(all_sequence_length, mask_end[batch]); - - // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result. - if (start_position >= end_position) { - start_position = 0; - end_position = all_sequence_length; - } - } - __syncthreads(); - - SoftmaxSmall(all_sequence_length, sequence_length, end_position, start_position, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); -} - -template -__device__ inline void SoftmaxSmallPacked(const int sequence_length, - const int end, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tmp_storage; - - __shared__ float sum_reverse_block; - __shared__ float max_block; - - // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; - const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length; - const int index = offset + threadIdx.x; - - bool is_valid = threadIdx.x < end; - - // e^x is represented as infinity if x is large enough, like 100.f. - // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. - // a math transform as below is leveraged to get a stable softmax: - // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) - const bool no_rpb = (rel_pos_bias == nullptr); - const int size_per_batch = gridDim.x * sequence_length; - float input_data = no_rpb - ? float(input[index]) - : float(input[index] + (broadcast_rel_pos_bias - ? rel_pos_bias[index % size_per_batch] - : rel_pos_bias[index])); - - float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F); - const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); - - // Store max value - if (threadIdx.x == 0) { - max_block = max; - } - __syncthreads(); - - float thread_data_exp(0.f); - if (is_valid) { - thread_data_exp = expf(input_data - max_block); - } - - const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); - - // Store value of 1.0/sum. - if (threadIdx.x == 0) { - sum_reverse_block = (1.f) / sum; - } - __syncthreads(); - - // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. - if (threadIdx.x < sequence_length) { - output[index] = T(thread_data_exp * sum_reverse_block); - } -} - -template -__global__ void SoftmaxKernelSmallWithCumSeqLen(const T* input, - const T* rel_pos_bias, const bool broadcast_rel_pos_bias, - const int* cum_seq_length, const int sequence_length, - T* output) { - __shared__ int end_position; - - if (threadIdx.x == 0) { - const int batch = blockIdx.y; - end_position = cum_seq_length[batch + 1] - cum_seq_length[batch]; - } - __syncthreads(); - - SoftmaxSmallPacked(sequence_length, end_position, - rel_pos_bias, broadcast_rel_pos_bias, - input, output); -} - -template -__global__ void SoftmaxKernelWithCumSeqLen(const T* input, - const T* rel_pos_bias, const bool broadcast_rel_pos_bias, - const int* cum_seq_length, const int sequence_length, - T* output) { - __shared__ int end_position; - - if (threadIdx.x == 0) { - const int batch = blockIdx.y; - end_position = cum_seq_length[batch + 1] - cum_seq_length[batch]; - } - __syncthreads(); - - Softmax(sequence_length, end_position, 0 /*start_position*/, - rel_pos_bias, broadcast_rel_pos_bias, input, output); -} - -template -__global__ void MaskedSoftmaxKernel(const int all_sequence_length, - const int* mask_end, - const int* mask_start, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, T* output) { - __shared__ int start_position; - __shared__ int end_position; - - if (threadIdx.x == 0) { - const int batch = blockIdx.y; - start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0; - end_position = min(all_sequence_length, mask_end[batch]); - - // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result. - if (start_position >= end_position) { - start_position = 0; - end_position = all_sequence_length; - } - } - __syncthreads(); - - Softmax(all_sequence_length, end_position, start_position, - rel_pos_bias, broadcast_rel_pos_bias, input, output); -} - -template -__global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length, - const int sequence_length, - const int* attention_mask, - const bool* key_padding_mask, - const T* rel_pos_bias, - const bool broadcast_rel_pos_bias, - const T* input, - T* output, - const bool is_unidirectional, - const float rsqrt_head_size, - const int mask_dimension, - const int max_sequence_length, - const bool skip_softmax, - const float mask_filter_value) { - SoftmaxWithRawMaskSmall( - all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, output, - is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - skip_softmax, mask_filter_value); -} + const bool broadcast_rel_pos_bias, T* input, T* output, bool is_unidirectional); template Status ComputeSoftmaxWithCumSeqLength( @@ -714,48 +21,7 @@ Status ComputeSoftmaxWithCumSeqLength( const int batch_size, const int sequence_length, const int num_heads, - T* output, cudaStream_t stream) { - const dim3 grid(sequence_length * num_heads, batch_size, 1); - - if (sequence_length <= 32) { - const int blockSize = 32; - SoftmaxKernelSmallWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - - } else if (sequence_length <= 64) { - const int blockSize = 64; - SoftmaxKernelSmallWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - } else if (sequence_length <= 128) { - const int blockSize = 128; - SoftmaxKernelSmallWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - } else if (sequence_length <= 256) { - const int blockSize = 256; - SoftmaxKernelSmallWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - } else if (sequence_length <= 512) { - const int blockSize = 512; - SoftmaxKernelSmallWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - } else if (sequence_length <= 1024) { - const int blockSize = 1024; - SoftmaxKernelSmallWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - } else { - SoftmaxKernelWithCumSeqLen - <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, - cum_seq_length, sequence_length, output); - } - - return CUDA_CALL(cudaGetLastError()); -} + T* output, cudaStream_t stream); template Status ComputeSoftmaxWithMask1D(cudaStream_t stream, @@ -769,50 +35,7 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream, const bool broadcast_rel_pos_bias, const T* input, T* output, - const bool is_unidirectional) { - const dim3 grid(sequence_length * num_heads, batch_size, 1); - - if (all_sequence_length <= 32) { - const int blockSize = 32; - MaskedSoftmaxKernelSmall - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 64) { - const int blockSize = 64; - MaskedSoftmaxKernelSmall - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 128) { - const int blockSize = 128; - MaskedSoftmaxKernelSmall - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 256) { - const int blockSize = 256; - MaskedSoftmaxKernelSmall - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 512) { - const int blockSize = 512; - MaskedSoftmaxKernelSmall - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (all_sequence_length <= 1024) { - const int blockSize = 1024; - MaskedSoftmaxKernelSmall - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); - } else if (!is_unidirectional) { - const int blockSize = 1024; - MaskedSoftmaxKernel - <<>>(all_sequence_length, mask_index, mask_start, - rel_pos_bias, broadcast_rel_pos_bias, input, output); - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024."); - } - - return CUDA_CALL(cudaGetLastError()); -} + const bool is_unidirectional); template Status ComputeSoftmaxWithRawMask(cudaStream_t stream, @@ -832,74 +55,8 @@ Status ComputeSoftmaxWithRawMask(cudaStream_t stream, const int max_sequence_length, const bool use_persistent_softmax, T* persistent_softmax_workspace, - const float mask_filter_value) { - const dim3 grid(sequence_length * num_heads, batch_size, 1); + const float mask_filter_value); - T* out = use_persistent_softmax ? persistent_softmax_workspace : output; - if (all_sequence_length <= 32) { - const int blockSize = 32; - SoftmaxWithRawMaskSmallKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } else if (all_sequence_length <= 64) { - const int blockSize = 64; - SoftmaxWithRawMaskSmallKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } else if (all_sequence_length <= 128) { - const int blockSize = 128; - SoftmaxWithRawMaskSmallKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } else if (all_sequence_length <= 256) { - const int blockSize = 256; - SoftmaxWithRawMaskSmallKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } else if (all_sequence_length <= 512) { - const int blockSize = 512; - SoftmaxWithRawMaskSmallKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } else if (all_sequence_length <= 1024) { - const int blockSize = 1024; - SoftmaxWithRawMaskSmallKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } else { - const int blockSize = 256; - const int sh_bytes = sizeof(float) * all_sequence_length; - SoftmaxWithRawMaskLargeKernel - <<>>(all_sequence_length, sequence_length, - attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, - out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length, - use_persistent_softmax, mask_filter_value); - } - - if (use_persistent_softmax) { - return dispatch_warpwise_softmax_forward(stream, - output, - persistent_softmax_workspace, - all_sequence_length, - all_sequence_length, - batch_size * num_heads * sequence_length); - } - - return CUDA_CALL(cudaGetLastError()); -} - -} // namespace cuda +} // namespace attention_softmax_cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu index 1729975453..ff0900a6b9 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu @@ -18,7 +18,7 @@ #include "contrib_ops/cuda/bert/rotary_embedding_util.h" using namespace onnxruntime::cuda; -using namespace cub; +using namespace onnxruntime::contrib::attention_softmax_cuda; #define CHECK_CUDA(expr) CUDA_RETURN_IF_ERROR(expr)