diff --git a/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.cc b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.cc new file mode 100644 index 0000000000..385653991f --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.cc @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "longformer_attention_base.h" + +namespace onnxruntime { +namespace contrib { + +LongformerAttentionBase::LongformerAttentionBase(const OpKernelInfo& info) { + int64_t num_heads = 0; + ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); + num_heads_ = static_cast(num_heads); + + int64_t window = 0; + ORT_ENFORCE(info.GetAttr("window", &window).IsOK() && window > 0); + window_ = static_cast(window); +} + +Status LongformerAttentionBase::CheckInputs(const TensorShape& input_shape, + const TensorShape& weights_shape, + const TensorShape& bias_shape, + const TensorShape& mask_shape, + const TensorShape& global_weights_shape, + const TensorShape& global_bias_shape, + const TensorShape& global_shape) const { + // Input shapes: + // input : (batch_size, sequence_length, hidden_size) + // weights : (hidden_size, 3 * hidden_size) + // bias : (3 * hidden_size) + // mask : (batch_size, sequence_length) + // global_weights : (hidden_size, 3 * hidden_size) + // global_bias : (3 * hidden_size) + // global : (batch_size, sequence_length) + + const auto& dims = input_shape.GetDims(); + if (dims.size() != 3) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'input' is expected to have 3 dimensions, got ", + dims.size()); + } + + int batch_size = static_cast(dims[0]); + int sequence_length = static_cast(dims[1]); + int hidden_size = static_cast(dims[2]); + if (sequence_length % (2 * window_) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'input' dimension 1 should be divisiable by 2W, where W is value of the window attribute."); + } + if (hidden_size % num_heads_ != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'input' dimension 2 should be divisiable by value of the num_heads attribute."); + } + + const auto& weights_dims = weights_shape.GetDims(); + if (weights_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' is expected to have 2 dimensions, got ", + weights_dims.size()); + } + if (weights_dims[0] != dims[2]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'weights' dimension 0 should have same length as dimension 2 of input 0"); + } + if (weights_dims[1] != 3 * weights_dims[0]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' dimension 1 should be 3 times of dimension 0"); + } + + const auto& bias_dims = bias_shape.GetDims(); + if (bias_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' is expected to have 1 dimension, got ", + bias_dims.size()); + } + if (bias_dims[0] != weights_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'bias' dimension 0 should have same length as dimension 1 of input 'weights'"); + } + + const auto& mask_dims = mask_shape.GetDims(); + if (mask_dims.size() == 2) { + if (static_cast(mask_dims[0]) != batch_size || static_cast(mask_dims[1]) != sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask' shall have shape batch_size x sequence_length"); + } + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask' is expected to have 2 dimensions, got ", + mask_dims.size()); + } + + const auto& global_weights_dims = global_weights_shape.GetDims(); + if (global_weights_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'global_weights' is expected to have 2 dimensions, got ", + weights_dims.size()); + } + if (global_weights_dims[0] != dims[2]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'global_weights' dimension 0 should have same length as dimension 2 of input 0"); + } + if (global_weights_dims[1] != 3 * global_weights_dims[0]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'global_weights' dimension 1 should be 3 times of dimension 0"); + } + + const auto& global_bias_dims = global_bias_shape.GetDims(); + if (global_bias_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'global_bias' is expected to have 1 dimension, got ", + global_bias_dims.size()); + } + if (global_bias_dims[0] != global_weights_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'global_bias' dimension 0 should have same length as dimension 1 of input 'global_weights'"); + } + + const auto& global_dims = global_shape.GetDims(); + if (global_dims.size() == 2) { + if (static_cast(global_dims[0]) != batch_size || static_cast(global_dims[1]) != sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'global' shall have shape batch_size x sequence_length"); + } + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'global' is expected to have 2 dimensions, got ", + global_dims.size()); + } + + return Status::OK(); +} + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h new file mode 100644 index 0000000000..182727ab58 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { + +class LongformerAttentionBase { + protected: + LongformerAttentionBase(const OpKernelInfo& info); + + Status CheckInputs(const TensorShape& input_shape, + const TensorShape& weights_shape, + const TensorShape& bias_shape, + const TensorShape& mask_shape, + const TensorShape& global_weights_shape, + const TensorShape& global_bias_shape, + const TensorShape& global_shape) const; + + int num_heads_; // Number of attention heads + int window_; // Attention windows length (W). It is half (one-sided) of total window size. +}; + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc new file mode 100644 index 0000000000..ef2eecb1ec --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc @@ -0,0 +1,138 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "longformer_attention.h" +#include "core/framework/tensorprotoutils.h" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" +#include "longformer_attention_impl.h" + +using namespace onnxruntime::cuda; +using namespace ::onnxruntime::common; +using namespace ONNX_NAMESPACE; + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + LongformerAttention, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + LongformerAttention); + +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +LongformerAttention::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {} + +template +Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const Tensor* weights = context->Input(1); + const Tensor* bias = context->Input(2); + const Tensor* mask = context->Input(3); + const Tensor* global_weights = context->Input(4); + const Tensor* global_bias = context->Input(5); + const Tensor* global_attention = context->Input(6); + ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), weights->Shape(), bias->Shape(), mask->Shape(), + global_weights->Shape(), global_bias->Shape(), global_attention->Shape())); + + // Input and output shapes: + // Input 0 - input : (batch_size, sequence_length, hidden_size) + // Output 0 - output : (batch_size, sequence_length, hidden_size) + const auto& shape = input->Shape(); + int batch_size = static_cast(shape[0]); + int sequence_length = static_cast(shape[1]); + int hidden_size = static_cast(shape[2]); + int head_size = hidden_size / num_heads_; + + Tensor* output = context->Output(0, shape); + + cublasHandle_t cublas = CublasHandle(); + constexpr size_t element_size = sizeof(T); + + // Use GEMM for fully connection. + int m = batch_size * sequence_length; + int n = 3 * hidden_size; + int k = hidden_size; + + size_t qkv_size = batch_size * sequence_length * 3 * hidden_size * element_size; + auto gemm_buffer = GetScratchBuffer(qkv_size); + + typedef typename ToCudaType::MappedType CudaT; + CudaT one = ToCudaType::FromFloat(1.0f); + CudaT zero = ToCudaType::FromFloat(0.0f); + + // Bias shape is (N), broadcast using B(N, M) = 1 * bias(N, 1) x ones(1, M) + 0 * B. + auto& device_prop = GetDeviceProp(); + CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( + cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, 1, &one, + reinterpret_cast(bias->template Data()), n, + GetConstOnes(m), 1, + &zero, reinterpret_cast(gemm_buffer.get()), n, device_prop)); + + // Gemm, note that CUDA assumes col-major, so result(N, M) = 1 * weights x input + 1 x B. + CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( + cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &one, + reinterpret_cast(weights->template Data()), n, + reinterpret_cast(input->template Data()), k, + &one, reinterpret_cast(gemm_buffer.get()), n, device_prop)); + + // TODO: calculate the exact value from global flags. + int max_num_global = sequence_length; + + // Fully connection for global projection. + // Note that Q only need handle global query tokens if we split GEMM to global Q/K/V separately. + // When there is no global token, need not run glboal GEMM. + auto global_gemm_buffer = GetScratchBuffer(max_num_global > 0 ? qkv_size : 0); + + if (max_num_global > 0) { + CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( + cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, 1, &one, + reinterpret_cast(global_bias->template Data()), n, + GetConstOnes(m), 1, + &zero, reinterpret_cast(global_gemm_buffer.get()), n, device_prop)); + + CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( + cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &one, + reinterpret_cast(global_weights->template Data()), n, + reinterpret_cast(input->template Data()), k, + &one, reinterpret_cast(global_gemm_buffer.get()), n, device_prop)); + } + + size_t workSpaceSize = GetLongformerAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, max_num_global); + auto workspace_buffer = GetScratchBuffer(workSpaceSize); + if (!LaunchLongformerAttentionKernel( + device_prop, + reinterpret_cast(gemm_buffer.get()), + reinterpret_cast(mask->template Data()), + reinterpret_cast(global_gemm_buffer.get()), + global_attention->template Data(), + output->template MutableData(), + batch_size, + sequence_length, + num_heads_, + head_size, + window_, + max_num_global, + workspace_buffer.get(), + cublas, + element_size)) { + // Get last error to reset it to cudaSuccess. + CUDA_CALL(cudaGetLastError()); + return Status(common::ONNXRUNTIME, common::FAIL); + } + + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h new file mode 100644 index 0000000000..c2ce91b8d2 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +#include "core/providers/cuda/cuda_common.h" +#include "contrib_ops/cpu/bert/longformer_attention_base.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +template +class LongformerAttention final : public CudaKernel, public LongformerAttentionBase { + public: + LongformerAttention(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu new file mode 100644 index 0000000000..fd9637dfc9 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu @@ -0,0 +1,859 @@ +/* +Copyright (c) NVIDIA Corporation and Microsoft 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. +*/ + +// Limitations of current Longformer Attention CUDA Kernels: +// (1) Does not support global tokens in the middle. All global tokens shall be in the beginning of sequence. +// (2) Batch size <= 128 (defined in MAX_LONGFORMER_BATCH_SIZE) + +#include +#include +#include +#include +#include +#include +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cuda_common.h" +#include "longformer_attention_impl.h" +#include "attention_impl.h" +#include "attention_softmax.h" + +using namespace onnxruntime::cuda; +using namespace cub; + +#define CHECK(status) \ + if (!CUBLAS_CALL(status)) { \ + return false; \ + } + +constexpr int MAX_LONGFORMER_BATCH_SIZE = 128; + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Denote: batch size (B), sequence length (S), number of heads (N), dimension per head (H), max number of global tokens (G) +// +// Workspace layout (by default, the data type T is float or half): +// [SoftmaxSpace: see below] [Q:BxNxSxH] [K:BxNxSxH] [V:BxNxSxH] [Global_Q:BxNxGxH] [Global_K:BxNxSxH] [Global_V:BxNxSxH] +// where Global_Q, Global_K and Global_V are optional. They are not allocated when there is no any global token. +// +// SoftmaxSpace layout (tmp_storage could use the space of scratch1, scratch2, Q and K): +// [Global_Idx: int BxS][batch_global_num: int BxS][sequence_index: int BxS][tmp_storage: int 1024x1] +// [scratch1: BxNxSxS ] [scratch2: BxNxSxS ] +// Allocated size could be slightly larger than needed: batch_global_num uses only Bx1 and allocated BxS. +// Scratch size is allocated as multiples of 256. + +size_t GetLongformerSoftmaxWorkspaceSize( + size_t element_size, + int batch_size, + int num_heads, + int sequence_length) { + size_t temp_size = sizeof(int) * 1024; + size_t scratch_size = 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length); + return 3 * batch_size * sequence_length * sizeof(int) + std::max(scratch_size, temp_size); +} + +size_t GetLongformerAttentionWorkspaceSize( + size_t element_size, + int batch_size, + int num_heads, + int head_size, + int sequence_length, + int max_num_global) { + size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length); + size_t qkv_size = 3 * batch_size * sequence_length * num_heads * head_size * element_size; + size_t global_qkv_size = max_num_global > 0 ? qkv_size : 0; + return softmax_size + qkv_size + global_qkv_size; +} + +__global__ void InitSequenceIndexKernel(int* sequence_index, int sequence_length) { + int batch_index = blockIdx.x; + for (int i = threadIdx.x; i < sequence_length; i += blockDim.x) { + sequence_index[batch_index * sequence_length + i] = i; + } +} + +// TODO: Move this to its own plugin that can be run once for all layers. +int* BuildGlobalIndex(cudaStream_t stream, const int* global_attention, int batch_size, int sequence_length, void* workspace, size_t softmax_workspace_size) { + int* global_idx = reinterpret_cast(workspace); + int* batch_global_num = global_idx + batch_size * sequence_length; // Number of global tokens in each batch, shape is (batch_size) + int* sequence_index = batch_global_num + batch_size * sequence_length; + int* tmp_storage = sequence_index + batch_size * sequence_length; + + InitSequenceIndexKernel<<>>(sequence_index, sequence_length); + + // Determine temporary device storage requirements + // Find the global attention indices and number of global attention tokens + size_t temp_storage_bytes = 0; + cub::DevicePartition::Flagged(NULL, temp_storage_bytes, sequence_index, + global_attention, global_idx, batch_global_num, sequence_length, stream); + assert(temp_storage_bytes <= softmax_workspace_size - static_cast(3 * batch_size * sequence_length)); + + for (int i = 0; i < batch_size; ++i) { + cub::DevicePartition::Flagged(reinterpret_cast(tmp_storage), temp_storage_bytes, sequence_index, + global_attention + i * sequence_length, global_idx + i * sequence_length, + batch_global_num + i, sequence_length, stream); + } + + return global_idx; +} + +template +__launch_bounds__(blockSize) + __global__ void LongformerSoftmaxKernel(const int* global_attention, + const int* global_idx, + const int* batch_global_num, + const T* input, + const T* attention_mask, + T* output, + float scaler, + int dim0, + int sequence_length, + int attention_window) { + typedef cub::BlockReduce BlockReduce; + __shared__ typename BlockReduce::TempStorage block_reduce_temp; + __shared__ float max_shared; + __shared__ float sum_shared; + + const T* input_block = input + sequence_length * blockIdx.x; + T* output_block = output + sequence_length * blockIdx.x; + const int batch_index = blockIdx.x / dim0; + const int row_index = blockIdx.x % sequence_length; + const int global_num = batch_global_num[batch_index]; + + // To be consistent with Huggingface Longformer, the row of maksed word are set as zero. + if ((float)attention_mask[batch_index * sequence_length + row_index] < 0.0f) { + for (int i = threadIdx.x; i < sequence_length; i += blockSize) { + output_block[i] = (T)(0); + } + return; + } + + // local attention token + int col_start = 0; + int col_end = sequence_length; + bool is_local_row = (global_attention[batch_index * sequence_length + row_index] == (int)0); + if (is_local_row) { + col_start = row_index - attention_window; + if (col_start < 0) { + col_start = 0; + } + + col_end = row_index + attention_window + 1; + if (col_end > sequence_length) { + col_end = sequence_length; + } + } + + const T* mask_block = attention_mask + sequence_length * batch_index; + int tid = threadIdx.x; + + // calculate max input + float max_input = -CUDART_INF_F; + // #pragma unroll 16 + for (int i = tid + col_start; i < col_end; i += blockSize) { + float x = input_block[i]; + x = x * scaler + (float)mask_block[i]; + if (max_input < x) { + max_input = x; + } + } + + if (is_local_row) { + for (int g = tid; g < global_num; g += blockSize) { + int i = global_idx[g]; + if (i < col_start || i > col_end) { + float x = input_block[i]; + x = x * scaler + (float)mask_block[i]; + if (max_input < x) { + max_input = x; + } + } + } + } + //__syncthreads(); + float max_block = BlockReduce(block_reduce_temp).Reduce(max_input, cub::Max()); + if (tid == 0) { + max_shared = max_block; + } + __syncthreads(); + + float sum_input = 0.f; + // #pragma unroll 16 + for (int i = tid + col_start; i < col_end; i += blockSize) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + sum_input += x; + } + + if (is_local_row) { + for (int g = tid; g < global_num; g += blockSize) { + int i = global_idx[g]; + if (i < col_start || i > col_end) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + sum_input += x; + } + } + } + + //__syncthreads(); + float sum_block = BlockReduce(block_reduce_temp).Reduce(sum_input, cub::Sum()); + if (tid == 0) { + sum_shared = sum_block; + } + __syncthreads(); + float recip_sum = 1.f / sum_shared; + + if (is_local_row) { + // We only need to fill in zeros for blocks that will be used in the matrix multiplication + // following the Softmax. + // + // For now zero-out only [row_index - 2*attention_window, row_index + 2*attention_window], + // we can even be more agressive and reduce the zeroing out window size since + // each row has entries in 3 blocks (3*attention_window size instead of 4*attention_window) + int zero_start = row_index - 2 * attention_window; + if (zero_start < 0) { + zero_start = 0; + } + + int zero_end = row_index + 2 * attention_window; + if (zero_end > sequence_length) { + zero_end = sequence_length; + } + + for (int i = tid + zero_start; i < zero_end; i += blockSize) { + output_block[i] = (T)(0.); + } + + for (int g = tid; g < global_num; g += blockSize) { + int i = global_idx[g]; + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + output_block[i] = (T)(recip_sum * x); + } + } + + // #pragma unroll 16 + for (int i = tid + col_start; i < col_end; i += blockSize) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + output_block[i] = (T)(recip_sum * x); + } +} + +bool launchSoftmaxKernel( + cudaStream_t stream, + cublasHandle_t cublas, + void* workspace, + size_t softmax_workspace_size, + const void* q, // transposed Q with shape (B, N, S, H) + const void* k, // transposed K with shape (B, N, S, H) + const void* v, // transposed V with shape (B, N, S, H) + const void* attention_mask, // attention mask with shape (B, S), with value 0 not masked and -10000 masked. + const void* global_q, // Q for global tokens with shape (B, N, G, H) + const void* global_k, // K for global tokens with shape (B, N, S, H) + const void* global_v, // V for global tokens with shape (B, N, S, H) + const int* global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention. + void* output, // output with shape (B, N, S, H) + float scaler, // scalar + int batch_size, // batch size + int sequence_length, // sequence length + int num_heads, // number of heads + int head_size, // hidden size per head + int attention_window, // one sided windows size + int max_num_global, // maximum number of global tokens (G) in all batches + size_t element_size) { // size of element: 2 for half, and 4 for float + if (batch_size > MAX_LONGFORMER_BATCH_SIZE) { + ORT_THROW("LongformerAttention CUDA operator does not support batch size > 128."); + } + + bool is_fp16 = (element_size == 2); + void* scratch1 = reinterpret_cast(workspace) + 3 * sizeof(int) * batch_size * sequence_length; + void* scratch2 = reinterpret_cast(scratch1) + GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length); + + // Build index for global tokens + int* global_idx = BuildGlobalIndex(stream, global_attention, batch_size, sequence_length, workspace, softmax_workspace_size); + int* batch_global_num = global_idx + batch_size * sequence_length; + + int num_global[MAX_LONGFORMER_BATCH_SIZE] = {-1}; + if (!CUDA_CALL(cudaMemcpyAsync(&num_global[0], batch_global_num, batch_size * sizeof(int), cudaMemcpyDeviceToHost, stream))) { + return false; + } + + // setup shared parameters for two strided batched matrix multiplies + cudaDataType_t Atype; + cudaDataType_t Btype; + cudaDataType_t Ctype; + cudaDataType_t resultType; + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT; + + __half one_fp16, zero_fp16; + float one_fp32, zero_fp32; + void *alpha, *beta_0, *beta_1; + + if (is_fp16) { + one_fp16 = __float2half(1.f); + zero_fp16 = __float2half(0.f); + alpha = static_cast(&one_fp16); + beta_0 = static_cast(&zero_fp16); + beta_1 = static_cast(&one_fp16); + Atype = CUDA_R_16F; + Btype = CUDA_R_16F; + Ctype = CUDA_R_16F; + resultType = CUDA_R_16F; + algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP; + } else { + one_fp32 = 1.f; + zero_fp32 = 0.f; + alpha = static_cast(&one_fp32); + beta_0 = static_cast(&zero_fp32); + beta_1 = static_cast(&one_fp32); + Atype = CUDA_R_32F; + Btype = CUDA_R_32F; + Ctype = CUDA_R_32F; + resultType = CUDA_R_32F; + } + + // Strided batch matrix multiply + // qk = q * k^T + // Shapes: q and k = B x N x S x H, qk = B x N x S x S + // Convert col-major to row-major by swapping q and k in Gemm + + // Local attention part + // S x S is calculated using sliding block WxW (W is one sided window size) like the following: + // [W][W] + // [W][W][W] + // [W][W][W] + // [W][W] + // The first and last rows have 2 blocks, and the remaining has 3 blocks per row. + // The calculation are splited into 3 parts. Firstly, fill the middle rows, then the first row and finally the last row. + // The results are stored in scratch1. + // TODO: Save space by not storing the whole matrix. Instead only allocate space for these blocks. + + int w = attention_window; + int x_offset = num_heads * sequence_length * head_size; + int y_offset = num_heads * sequence_length * sequence_length; + int last_block = (sequence_length / w) - 1; + int strideA = sequence_length * head_size; + int strideB = sequence_length * head_size; + int strideC = sequence_length * sequence_length; + + // When S == 2W, there is no middle rows of blocks: + // [W][W] + // [W][W] + // We can use normal matrix multiplication in this case. + if (sequence_length == 2 * w) { + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + sequence_length, + sequence_length, + head_size, + alpha, + k, + Atype, + head_size, + sequence_length * head_size, + q, + Btype, + head_size, + sequence_length * head_size, + beta_0, + scratch1, + Ctype, + sequence_length, + sequence_length * sequence_length, + batch_size * num_heads, + resultType, + algo)); + } else { // sequence_length > 2 * w + for (int i = 0; i < batch_size; ++i) { + for (int j = 0; j < num_heads; ++j) { + void* q_head = (char*)q + (i * x_offset + j * sequence_length * head_size + w * head_size) * element_size; + void* k_head = (char*)k + (i * x_offset + j * sequence_length * head_size) * element_size; + void* qk_head = (char*)scratch1 + (i * y_offset + j * sequence_length * sequence_length + w * sequence_length) * element_size; + int count = (sequence_length - 2 * w) / w; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + 3 * w, // m + w, // n + head_size, // k + alpha, // alpha + k_head, // A + Atype, // A type + head_size, // lda + w * head_size, // strideA + q_head, // B + Btype, // B type + head_size, // ldb + w * head_size, // strideB + beta_0, // beta + qk_head, // C + Ctype, // C type + sequence_length, // ldc + sequence_length * w + w, // strideC + count, // batch count + resultType, + algo)); + } + } + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + 2 * w, // m + w, // n + head_size, // k + alpha, // alpha + k, // A + Atype, // A type + head_size, // lda + strideA, // strideA + q, // B + Btype, // B type + head_size, // ldb + strideB, // strideB + beta_0, // beta + scratch1, // C + Ctype, // C type + sequence_length, // ldc + strideC, // strideC + batch_size * num_heads, // batch count + resultType, + algo)); + + void* q_head = (char*)q + (last_block * w * head_size) * element_size; + void* k_head = (char*)k + ((last_block - 1) * w * head_size) * element_size; + void* qk_head = (char*)scratch1 + (last_block * w * sequence_length + (last_block - 1) * w) * element_size; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + 2 * w, + w, + head_size, + alpha, + k_head, + Atype, + head_size, + strideA, + q_head, + Btype, + head_size, + strideB, + beta_0, + qk_head, + Ctype, + sequence_length, + strideC, + batch_size * num_heads, + resultType, + algo)); + } + + // Global attention part + for (int i = 0; i < batch_size; ++i) { + if (num_global[i] > 0) { + void* q_batch = (char*)q + (i * x_offset) * element_size; + void* k_batch = (char*)k + (i * x_offset) * element_size; + void* qk_batch = (char*)scratch1 + (i * y_offset) * element_size; + // Local tokens attending global tokens + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + num_global[i], + sequence_length, + head_size, + alpha, + k_batch, + Atype, + head_size, + strideA, + q_batch, + Btype, + head_size, + strideB, + beta_0, + qk_batch, + Ctype, + sequence_length, + strideC, + num_heads, + resultType, + algo)); + + void* global_q_batch = (char*)global_q + (i * num_heads * max_num_global * head_size) * element_size; + void* global_k_batch = (char*)global_k + (i * x_offset) * element_size; + int strideB_global = max_num_global * head_size; + + // Global tokens attending everything + // This GEMMs need to be last to make sure all global token entries are re-written. + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + sequence_length, + num_global[i], + head_size, + alpha, + global_k_batch, + Atype, + head_size, + strideA, + global_q_batch, + Btype, + head_size, + strideB_global, + beta_0, + qk_batch, + Ctype, + sequence_length, + strideC, + num_heads, + resultType, + algo)); + } + } + + int dim0 = sequence_length * num_heads; + int dim1 = sequence_length; + void* softmax_out = scratch2; + + const int blockSize = 64; + const int gridSize = batch_size * num_heads * sequence_length; + if (is_fp16) { + LongformerSoftmaxKernel<__half, blockSize><<>>( + global_attention, + global_idx, + batch_global_num, + static_cast(scratch1), + static_cast(attention_mask), + static_cast<__half*>(softmax_out), scaler, dim0, dim1, attention_window); + } else { + LongformerSoftmaxKernel<<>>( + global_attention, + global_idx, + batch_global_num, + static_cast(scratch1), + static_cast(attention_mask), + static_cast(softmax_out), scaler, dim0, dim1, attention_window); + } + + // Run the matrix multiply: output = softmax_out * v + // softmax_out: B x N x S x S + // v: B x N x S x H + // attn_out: B x N x S x H + // Calculation uses full Gemm (S == 2W) or sliding blocks (S > 2W) in a way similar to local attention part. + + if (sequence_length == 2 * w) { + // convert col-major to row-major by swapping softmax_out and v + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + sequence_length, + sequence_length, + alpha, + v, + Atype, + head_size, + sequence_length * head_size, + softmax_out, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, + output, + Ctype, + head_size, + sequence_length * head_size, + batch_size * num_heads, + resultType, + algo)); + } + else { // sequence_length > 2 * w + for (int i = 0; i < batch_size; ++i) { + for (int j = 0; j < num_heads; ++j) { + void* v_head = (char*)v + (i * x_offset + j * head_size * sequence_length) * element_size; + void* prob_head = (char*)softmax_out + (i * y_offset + j * sequence_length * sequence_length + w * sequence_length) * element_size; + void* out_head = (char*)output + (i * x_offset + j * head_size * sequence_length + w * head_size) * element_size; + int count = (sequence_length - 2 * w) / w; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + w, + 3 * w, + alpha, + v_head, + Atype, + head_size, + w * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * w + w, + beta_0, + out_head, + Ctype, + head_size, + w * head_size, + count, + resultType, + algo)); + } + } + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + w, + 2 * w, + alpha, + v, + Atype, + head_size, + sequence_length * head_size, + softmax_out, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, + output, + Ctype, + head_size, + sequence_length * head_size, + batch_size * num_heads, + resultType, + algo)); + + void* v_head = (char*)v + (last_block - 1) * w * head_size * element_size; + void* prob_head = (char*)softmax_out + (sequence_length * last_block * w + (last_block - 1) * w) * element_size; + void* out_head = (char*)output + last_block * w * head_size * element_size; + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + w, + 2 * w, + alpha, + v_head, + Atype, + head_size, + sequence_length * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, + out_head, + Ctype, + head_size, + sequence_length * head_size, + batch_size * num_heads, + resultType, + algo)); + } + + + for (int i = 0; i < batch_size; ++i) { + if (num_global[i] > 0) { + int glob_longdim_mm = (last_block - 1) * w; + + void* v_head = (char*)v + (i * x_offset) * element_size; + void* prob_head = (char*)softmax_out + (i * y_offset + 2 * w * sequence_length) * element_size; + void* out_head = (char*)output + (i * x_offset + 2 * w * head_size) * element_size; + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + glob_longdim_mm, + num_global[i], + alpha, + v_head, + Atype, + head_size, + sequence_length * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_1, + out_head, + Ctype, + head_size, + sequence_length * head_size, + num_heads, + resultType, + algo)); + + // Global tokens + v_head = (char*)global_v + (i * x_offset) * element_size; + prob_head = (char*)softmax_out + (i * y_offset) * element_size; + out_head = (char*)output + (i * x_offset) * element_size; + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + num_global[i], + sequence_length, // Re-write entries completely + alpha, + v_head, + Atype, + head_size, + sequence_length * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, // Use beta=0 to overwrite + out_head, // Here assumes global tokens are at the beginning of sequence. + Ctype, + head_size, + sequence_length * head_size, + num_heads, + resultType, + algo)); + } + } + + return true; +} + +template +bool LongformerQkvToContext( + const cudaDeviceProp& prop, cublasHandle_t& cublas, cudaStream_t stream, + const int batch_size, const int sequence_length, const int num_heads, const int head_size, + const int window, const size_t element_size, + const T* input, const T* attention_mask, + const T* global_input, const int* global_attention, const int max_num_global, + T* workspace, + T* output) { + size_t softmax_workspace_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length); + T* qkv = reinterpret_cast((char*)workspace + softmax_workspace_size); + + // Input should be BxSx3xNxH => qkv: 3xBxNxSxH + if (!LaunchTransQkv(stream, sequence_length, batch_size, head_size, num_heads, input, qkv)) { + return false; + } + + // Input 'global_input' should be BxSx3xNxH => global_qkv: 3xBxNxSxH + T* global_qkv = qkv + 3 * batch_size * sequence_length * num_heads * head_size * element_size; + + // When there is no global token, no need to process global Q, K and V + if (max_num_global > 0 && nullptr != global_input) { + if (!LaunchTransQkv(stream, sequence_length, batch_size, head_size, num_heads, global_input, global_qkv)) { + return false; + } + } + + // Now qkv has Q, K, V: each has size BxNxSxH + const int elements = batch_size * num_heads * sequence_length * head_size; + const T* q = qkv; + const T* k = q + elements; + const T* v = k + elements; + + const T* global_q = global_qkv; + const T* global_k = global_q + elements; + const T* global_v = global_k + elements; + + cublasSetStream(cublas, stream); + CublasMathModeSetter helper(prop, cublas, CUBLAS_TENSOR_OP_MATH); + + // Q*K' are scaled by 1/sqrt(H) + const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); + + cudaDeviceSynchronize(); + + T* temp_output = qkv; // Q will be overwritten + if (!launchSoftmaxKernel( + stream, + cublas, + workspace, + softmax_workspace_size, + q, // Transposed Q with shape B x N x S x H + k, // Transposed K with shape B x N x S x H + v, // Transposed V with shape B x N x S x H + attention_mask, // Attention mask flags with shape B x S + global_q, // Transposed global Q with shape B x N x G x H + global_k, // Transposed global K with shape B x N x S x H + global_v, // Transposed global V with shape B x N x S x H + global_attention, // Global attention flags with shape B x S + temp_output, // Output with shape B x N x S x H + rsqrt_head_size, // Scaler + batch_size, // Batch size + sequence_length, // Sequence length + num_heads, // Number of attention heads + head_size, // Hidden size per head + window, // Half (one-sided) windows size + max_num_global, // Maximum number of global tokens (G) + element_size)) { + return false; + } + + // The temp_output is BxNxSxH, transpose it to final output BxSxNxH + return LaunchTransCtx(stream, sequence_length, batch_size, head_size, num_heads, temp_output, output); +} + +bool LaunchLongformerAttentionKernel( + const cudaDeviceProp& prop, + const void* input, + const void* attention_mask, + const void* global_input, + const int* global_attention, + void* output, + int batch_size, + int sequence_length, + int num_heads, + int head_size, + int window, + int max_num_global, + void* workspace, + cublasHandle_t& cublas, + const size_t element_size) { + // use default stream + const cudaStream_t stream = nullptr; + + if (element_size == 2) { + return LongformerQkvToContext(prop, cublas, stream, + batch_size, sequence_length, num_heads, head_size, window, element_size, + reinterpret_cast(input), + reinterpret_cast(attention_mask), + reinterpret_cast(global_input), + global_attention, + max_num_global, + reinterpret_cast(workspace), + reinterpret_cast(output)); + } else { + return LongformerQkvToContext(prop, cublas, stream, + batch_size, sequence_length, num_heads, head_size, window, element_size, + reinterpret_cast(input), + reinterpret_cast(attention_mask), + reinterpret_cast(global_input), + global_attention, + max_num_global, + reinterpret_cast(workspace), + reinterpret_cast(output)); + } +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h new file mode 100644 index 0000000000..632f6d6e5c --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/providers/cuda/shared_inc/cuda_utils.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +size_t GetLongformerAttentionWorkspaceSize( + size_t element_size, + int batchsize, + int num_heads, + int head_size, + int sequence_length, + int max_num_global); + +bool LaunchLongformerAttentionKernel( + const cudaDeviceProp& device_prop, // Device Properties + const void* input, // Input tensor + const void* attention_mask, // Attention mask with shape (B, S) + const void* global_input, // Global attention input, or nullptr when max_num_global == 0. + const int* global_attention, // Global attention flags with shape (B, S) + void* output, // Output tensor + int batch_size, // Batch size (B) + int sequence_length, // Sequence length (S) + int num_heads, // Number of attention heads (N) + int head_size, // Hidden layer size per head (H) + int window, // One sided attention window (W) + int max_num_global, // Maximum number of global tokens (G) + void* workspace, // Temporary buffer + cublasHandle_t& cublas, // Cublas handle + const size_t element_size // Element size of input tensor +); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 6dc5fa4f1d..bc94cc266a 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -53,6 +53,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ImageScaler); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ImageScaler); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ImageScaler); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, LongformerAttention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, LongformerAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ParametricSoftplus); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ParametricSoftplus); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ParametricSoftplus); @@ -130,6 +132,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 5d18fd6892..43d730a29f 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -431,6 +431,35 @@ and present state are optional. Present state could appear in output even when p return; }); + static const char* Longformer_Attention_doc = R"DOC( +Longformer Self Attention with a local context and a global context. Tokens attend locally: Each token +attends to its W previous tokens and W succeding tokens with W being the window length. A selected few tokens +attend globally to all other tokens. + +The attention mask is of shape (batch_size, sequence_length), where sequence_length is a multiple of 2W after padding. +Mask value < 0 (like -10000.0) means the token is masked, 0 otherwise. + +Global attention flags have value 1 for the tokens attend globally and 0 otherwise. +)DOC"; + + ONNX_CONTRIB_OPERATOR_SCHEMA(LongformerAttention) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc(Longformer_Attention_doc) + .Attr("num_heads", "Number of attention heads", AttributeProto::INT) + .Attr("window", "One sided attention windows length W, or half of total window length", AttributeProto::INT) + .Input(0, "input", "3D input tensor with shape (batch_size, sequence_length, hidden_size), hidden_size = num_heads * head_size", "T") + .Input(1, "weight", "2D input tensor with shape (hidden_size, 3 * hidden_size)", "T") + .Input(2, "bias", "1D input tensor with shape (3 * hidden_size)", "T") + .Input(3, "mask", "Attention mask with shape (batch_size, sequence_length)", "T") + .Input(4, "global_weight", "2D input tensor with shape (hidden_size, 3 * hidden_size)", "T") + .Input(5, "global_bias", "1D input tensor with shape (3 * hidden_size)", "T") + .Input(6, "global", "Global attention flags with shape (batch_size, sequence_length)", "G") + .Output(0, "output", "3D output tensor with shape (batch_size, append_length, hidden_size)", "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("G", {"tensor(int32)"}, "Constrain to integer types") + .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); + static const char* EmbedLayerNormalization_ver1_doc = R"DOC( EmbedLayerNormalization is the fusion of embedding layer in BERT model, with optional mask processing. The embedding layer takes input_ids (word IDs) and segment_ids (sentence IDs) to look up word_embedding, position_embedding, diff --git a/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc b/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc new file mode 100644 index 0000000000..8ad0d2f44c --- /dev/null +++ b/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc @@ -0,0 +1,301 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/common/cuda_op_test_utils.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { + +static void RunAttentionTest( + const std::vector& input_data, // input: [batch_size, sequence_length, hidden_size] + const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] + const std::vector& bias_data, // bias: [3 * hidden_size] + const std::vector& mask_data, // mask: [batch_size, sequence_length] + const std::vector& global_weights_data, // global_weights: [hidden_size, 3 * hidden_size] + const std::vector& global_bias_data, // global_bias: [3 * hidden_size] + const std::vector& global_data, // global: [batch_size, sequence_length] + const std::vector& output_data, // output: [batch_size, sequence_length, hidden_size] + int batch_size, + int sequence_length, + int hidden_size, + int number_of_heads, + int window, + bool use_float16 = false) { + int min_cuda_architecture = use_float16 ? 530 : 0; + + bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + bool enable_cpu = false; + if (enable_cpu || enable_cuda) { + OpTester tester("LongformerAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(number_of_heads)); + tester.AddAttribute("window", static_cast(window)); + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector weights_dims = {hidden_size, 3 * hidden_size}; + std::vector bias_dims = {3 * hidden_size}; + std::vector mask_dims = {batch_size, sequence_length}; + std::vector global_weights_dims = {hidden_size, 3 * hidden_size}; + std::vector global_bias_dims = {3 * hidden_size}; + std::vector global_dims = {batch_size, sequence_length}; + std::vector output_dims = input_dims; + + if (use_float16) { + tester.AddInput("input", input_dims, ToFloat16(input_data)); + tester.AddInput("weight", weights_dims, ToFloat16(weights_data)); + tester.AddInput("bias", bias_dims, ToFloat16(bias_data)); + tester.AddInput("mask", mask_dims, ToFloat16(mask_data)); + tester.AddInput("global_weight", global_weights_dims, ToFloat16(global_weights_data)); + tester.AddInput("global_bias", global_bias_dims, ToFloat16(global_bias_data)); + tester.AddOutput("output", output_dims, ToFloat16(output_data)); + } else { + tester.AddInput("input", input_dims, input_data); + tester.AddInput("weight", weights_dims, weights_data); + tester.AddInput("bias", bias_dims, bias_data); + tester.AddInput("mask", mask_dims, mask_data); + tester.AddInput("global_weight", global_weights_dims, global_weights_data); + tester.AddInput("global_bias", global_bias_dims, global_bias_data); + tester.AddOutput("output", output_dims, output_data); + } + + tester.AddInput("global", global_dims, global_data); + + if (enable_cuda) { + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + + if (enable_cpu) { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + } +} + +static void GetTinyLongformerData( + std::vector& weight_data, + std::vector& bias_data, + std::vector& global_weight_data, + std::vector& global_bias_data) { + weight_data = { + -0.03736265f, -0.03700203f, -0.01265421f, -0.00005944f, 0.03058976f, -0.00278089f, -0.00340891f, -0.00879795f, + -0.02133591f, 0.00784427f, 0.03478181f, -0.03402156f, 0.02373371f, 0.0063644f, -0.01468624f, 0.0504882f, + 0.00077419f, 0.00028939f, -0.02853666f, -0.01804554f, 0.00576997f, 0.0259758f, -0.0051732f, 0.01098385f, + -0.03283566f, 0.0021469f, 0.00610109f, -0.03178682f, 0.00858954f, 0.0257727f, -0.02472394f, -0.02258367f, + -0.0397051f, 0.00637861f, -0.00492215f, 0.00614283f, 0.01885833f, 0.0072404f, -0.02029094f, 0.00384926f, + -0.01666811f, 0.01191303f, 0.00760845f, -0.01181425f, -0.00996346f, 0.00950979f, -0.01039405f, 0.01063836f, + -0.0038129f, -0.00482849f, -0.00006185f, 0.01160069f, 0.00319028f, 0.0009577f, 0.03295509f, 0.00164638f, + -0.0020025f, 0.01478013f, -0.01177825f, 0.04226499f, -0.0057663f, -0.00160333f, 0.00270868f, -0.00655655f, + 0.02530637f, 0.00078956f, -0.00534409f, 0.01388706f, -0.01226765f, 0.02728544f, 0.01600054f, -0.01379196f, + 0.01883291f, -0.01875546f, -0.01316095f, 0.01392878f, 0.0286913f, 0.02039422f, 0.0164386f, -0.01549883f, + -0.01735487f, -0.04088231f, -0.00565702f, 0.00084416f, 0.03280086f, -0.02137024f, 0.00550848f, 0.0178341f, + 0.01426786f, -0.03154857f, -0.02185805f, -0.00808465f, -0.01845198f, 0.03087913f, 0.01739769f, 0.01742294f, + -0.02740332f, -0.00040067f, -0.01135285f, 0.01528029f, -0.02968147f, 0.02274286f, -0.03160891f, 0.02072382f, + 0.01509116f, 0.03541054f, 0.01261792f, 0.02072236f, 0.00617444f, 0.01808203f, 0.00342582f, 0.03231942f, + 0.04922125f, 0.00663754f, -0.01104379f, 0.00624478f, -0.00897994f, 0.01370005f, -0.01877974f, -0.0193537f, + 0.02327156f, -0.0071105f, -0.01520648f, 0.00595162f, 0.00329103f, 0.00786545f, 0.0109383f, 0.0170538f, + -0.01224139f, -0.00105941f, 0.01048978f, -0.00976777f, -0.02140773f, 0.00742492f, 0.01248992f, 0.0075539f, + 0.00286943f, -0.01414381f, 0.03620159f, 0.01033463f, 0.01399594f, 0.01650806f, -0.00202679f, -0.00328803f, + -0.02329111f, -0.02204572f, -0.03433893f, 0.01819115f, 0.01876776f, 0.00352132f, -0.02261852f, -0.0021271f, + -0.01685239f, 0.00707082f, 0.00893882f, -0.03597079f, -0.02897478f, 0.01790397f, -0.00189054f, -0.00782812f, + 0.01679492f, 0.01131928f, -0.02338723f, -0.00732064f, 0.00425937f, -0.02528145f, -0.01408323f, -0.04177788f, + -0.00529976f, 0.01193773f, -0.00405255f, 0.02042645f, -0.03389123f, -0.02279385f, 0.01627762f, 0.02311901f, + 0.01063351f, -0.05164707f, 0.0231087f, 0.00633761f, 0.01774552f, 0.02063041f, -0.01823085f, -0.00058292f, + 0.05767727f, 0.03186193f, -0.01058707f, 0.00119207f, -0.00613069f, 0.00137152f, 0.00948873f, -0.01503408f}; + + bias_data = { + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + + global_weight_data = { + 0.0294202f, -0.01245436f, -0.03289853f, 0.02588101f, 0.01704302f, 0.0005645f, 0.01802721f, -0.02043311f, + 0.00065155f, 0.01990021f, -0.02261156f, -0.01354096f, 0.03066289f, 0.00368064f, 0.01333689f, -0.02616469f, + 0.03910812f, 0.00521647f, -0.01002843f, 0.00574944f, -0.00215477f, -0.03501118f, -0.00492278f, 0.02129062f, + 0.0198415f, 0.00558937f, 0.01016298f, -0.01226531f, 0.02398407f, 0.00804429f, 0.01118671f, -0.00166833f, + -0.01615552f, -0.01540161f, 0.00378566f, 0.01431978f, -0.00993885f, 0.00361956f, 0.0116084f, -0.00427342f, + -0.0176292f, -0.01806638f, 0.01677675f, -0.0068318f, 0.00709118f, -0.01304021f, 0.00310118f, -0.01056628f, + -0.03345605f, 0.03466916f, 0.0154178f, -0.00299957f, 0.00297282f, -0.03894032f, 0.01632687f, 0.00817282f, + 0.02143084f, -0.01206295f, -0.0027577f, 0.0047877f, 0.00112486f, 0.01005225f, 0.03119631f, -0.01274098f, + 0.03476755f, 0.02006538f, -0.02917822f, 0.00559175f, -0.01407396f, 0.03570889f, 0.03545335f, -0.00095669f, + -0.00425291f, 0.00163976f, -0.03921806f, 0.00625822f, -0.0246806f, 0.00815409f, 0.0092521f, 0.00433091f, + 0.00997109f, 0.00915498f, -0.02030378f, -0.00968541f, 0.00971882f, -0.00865303f, -0.05141142f, 0.02430373f, + -0.02086174f, -0.03922447f, 0.01982334f, -0.01131777f, 0.01826571f, 0.01116992f, -0.01033904f, 0.02003466f, + -0.01685323f, -0.00281926f, -0.01908645f, 0.00376858f, 0.03040931f, -0.01649855f, -0.00193368f, 0.02237837f, + -0.01741545f, -0.0163137f, -0.01750915f, -0.03061339f, -0.00099388f, 0.02248359f, -0.00808812f, -0.00336211f, + 0.00031454f, 0.04020427f, -0.03191524f, -0.01187126f, 0.01208827f, -0.00885536f, -0.02643824f, -0.00441225f, + -0.00770211f, -0.00784583f, 0.01867373f, -0.01057389f, -0.00694279f, 0.02644924f, -0.00346569f, -0.01113159f, + -0.01246041f, -0.02784149f, -0.00033919f, 0.00444981f, -0.0196088f, 0.00031822f, -0.00622559f, -0.0112043f, + -0.00666037f, 0.00290975f, 0.00625247f, 0.02937804f, 0.01013064f, -0.00624247f, -0.00471575f, 0.00657559f, + -0.03100131f, 0.03663468f, -0.00701449f, -0.02267093f, 0.00842566f, -0.01073638f, -0.01930492f, 0.01421315f, + 0.03018992f, 0.00599129f, 0.00773254f, 0.0330062f, 0.00622745f, 0.00355852f, -0.02076612f, -0.02770657f, + -0.0232522f, -0.01313687f, -0.0039819f, -0.00212091f, -0.01119681f, -0.02110178f, -0.02889993f, -0.01824146f, + 0.00702973f, -0.02526855f, 0.02776029f, 0.02142943f, -0.01821483f, 0.01552922f, 0.01018873f, -0.02794646f, + 0.0458297f, 0.01344144f, -0.00684268f, 0.00251583f, -0.00135352f, -0.01950933f, 0.00551906f, -0.00823819f, + -0.01059994f, 0.01091084f, -0.00090001f, -0.00103661f, 0.0260281f, -0.01601363f, -0.00342217f, 0.0198724f}; + + global_bias_data = { + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; +} + +// Run test on tiny longformer model with batch size 1 and sequence length 4 or 8 +static void RunTinyLongformerBatch1( + std::vector& mask_data, + std::vector& global_data, + std::vector& output_data, + bool use_float16, + bool window_cover_whole_sequence = false) { + int batch_size = 1; + int one_sided_attention_window_size = 2; + int hidden_size = 8; + int number_of_heads = 2; + + // Total windows size 4 will cover the whole sequence length 4 + int sequence_length = window_cover_whole_sequence ? 4 : 8; + std::vector input_data; + if (window_cover_whole_sequence) { + input_data = { + -1.1354f, -0.3925f, -0.4778f, -1.2506f, 0.0692f, 1.6338f, 0.1119f, 1.4415f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f, + -1.8157f, -0.6058f, 0.4016f, -0.4158f, -0.0343f, 1.0701f, -0.2686f, 1.6685f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f}; + } else { + input_data = { + -1.1354f, -0.3925f, -0.4778f, -1.2506f, 0.0692f, 1.6338f, 0.1119f, 1.4415f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f, + -1.8157f, -0.6058f, 0.4016f, -0.4158f, -0.0343f, 1.0701f, -0.2686f, 1.6685f, + -1.5789f, 1.1709f, -1.4334f, 0.7868f, -0.2459f, 0.3509f, -0.1686f, 1.1181f, + 0.7612f, -2.1558f, -0.2982f, 0.8303f, 0.8898f, -0.8073f, 0.0386f, 0.7414f, + -1.6978f, 0.7141f, 1.1334f, -1.1359f, 0.8673f, -0.8174f, 0.7631f, 0.1731f, + -1.0536f, -0.0425f, -1.1194f, -0.6423f, 2.1825f, 0.2547f, 0.6015f, -0.1809f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f}; + } + + std::vector weight_data; + std::vector bias_data; + std::vector global_weight_data; + std::vector global_bias_data; + GetTinyLongformerData(weight_data, bias_data, global_weight_data, global_bias_data); + + RunAttentionTest(input_data, weight_data, bias_data, mask_data, global_weight_data, global_bias_data, global_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, one_sided_attention_window_size, use_float16); +} + +TEST(LongformerAttentionTest, LongformerAttention_NoGlobal) { + // last word is masked. + std::vector mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f}; + + // no global attention. + std::vector global_data = {0, 0, 0, 0, 0, 0, 0, 0}; + + std::vector output_data = { + 0.0793f, 0.0678f, 0.0705f, 0.0456f, 0.0197f, -0.0301f, 0.0010f, -0.0511f, + 0.0607f, 0.0545f, 0.0659f, 0.0324f, 0.0102f, -0.0317f, 0.0017f, -0.0328f, + 0.0738f, 0.0416f, 0.0322f, 0.0258f, 0.0063f, -0.0222f, 0.0053f, -0.0309f, + 0.0725f, 0.0451f, 0.0149f, 0.0241f, -0.0071f, -0.0231f, 0.0028f, -0.0369f, + 0.0738f, 0.0303f, 0.0137f, 0.0230f, -0.0089f, -0.0342f, -0.0085f, -0.0450f, + 0.0658f, 0.0278f, -0.0035f, 0.0129f, -0.0117f, -0.0357f, -0.0190f, -0.0424f, + 0.0860f, 0.0322f, -0.0218f, 0.0194f, -0.0096f, -0.0354f, -0.0265f, -0.0639f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f}; + + RunTinyLongformerBatch1(mask_data, global_data, output_data, false); +} + +/* +* This following case is generated from the first self-attention of a tiny longformer model like the following: + import torch + from transformers import AutoModel + model = AutoModel.from_pretrained("patrickvonplaten/longformer-random-tiny") + input_ids = torch.arange(7).unsqueeze(0) + attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device) + global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=input_ids.device) + global_attention_mask[:, [0, 1,]] = 1 + outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask) +*/ + +TEST(LongformerAttentionTest, LongformerAttention_GlobalStart) { + std::vector mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f}; + + // Global at the start of the sequence + std::vector global_data = {1, 1, 0, 0, 0, 0, 0, 0}; + + std::vector output_data = { + -0.0356f, 0.0521f, -0.0198f, 0.0005f, 0.0245f, -0.0014f, -0.0170f, -0.0123f, + -0.0356f, 0.0521f, -0.0198f, 0.0005f, 0.0245f, -0.0014f, -0.0169f, -0.0122f, + 0.0738f, 0.0416f, 0.0322f, 0.0258f, 0.0063f, -0.0222f, 0.0053f, -0.0309f, + 0.0719f, 0.0473f, 0.0291f, 0.0279f, 0.0008f, -0.0289f, 0.0004f, -0.0416f, + 0.0716f, 0.0449f, 0.0283f, 0.0269f, 0.0017f, -0.0332f, -0.0104f, -0.0461f, + 0.0659f, 0.0457f, 0.0193f, 0.0208f, 0.0017f, -0.0341f, -0.0176f, -0.0445f, + 0.0780f, 0.0519f, 0.0128f, 0.0263f, 0.0056f, -0.0335f, -0.0219f, -0.0578f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f}; + + RunTinyLongformerBatch1(mask_data, global_data, output_data, false); +} + +TEST(LongformerAttentionTest, LongformerAttention_Float16) { + std::vector mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f}; + + // Global at the start of the sequence + std::vector global_data = {1, 1, 0, 0, 0, 0, 0, 0}; + + std::vector output_data = { + -0.0356f, 0.0521f, -0.0198f, 0.0005f, 0.0245f, -0.0014f, -0.0170f, -0.0123f, + -0.0356f, 0.0521f, -0.0198f, 0.0005f, 0.0245f, -0.0014f, -0.0169f, -0.0122f, + 0.0738f, 0.0416f, 0.0322f, 0.0258f, 0.0063f, -0.0222f, 0.0053f, -0.0309f, + 0.0719f, 0.0473f, 0.0291f, 0.0279f, 0.0008f, -0.0289f, 0.0004f, -0.0416f, + 0.0716f, 0.0449f, 0.0283f, 0.0269f, 0.0017f, -0.0332f, -0.0104f, -0.0461f, + 0.0659f, 0.0457f, 0.0193f, 0.0208f, 0.0017f, -0.0341f, -0.0176f, -0.0445f, + 0.0780f, 0.0519f, 0.0128f, 0.0263f, 0.0056f, -0.0335f, -0.0219f, -0.0578f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f}; + + RunTinyLongformerBatch1(mask_data, global_data, output_data, true); +} + +TEST(LongformerAttentionTest, LongformerAttention_FullWindow) { + // last word is masked. + std::vector mask_data = {0.0f, 0.0f, 0.0f, -10000.0f}; + + // no global attention. + std::vector global_data = {0, 0, 0, 0}; + + std::vector output_data = { + 0.0793f, 0.0678f, 0.0705f, 0.0456f, 0.0197f, -0.0301f, 0.0010f, -0.0511f, + 0.0793f, 0.0678f, 0.0705f, 0.0456f, 0.0197f, -0.0301f, 0.0010f, -0.0511f, + 0.0793f, 0.0678f, 0.0705f, 0.0456f, 0.0197f, -0.0301f, 0.0010f, -0.0511f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000}; + + // One doulbe-sided window will cover the whole sequence. + bool window_cover_whole_sequence = true; + RunTinyLongformerBatch1(mask_data, global_data, output_data, false, window_cover_whole_sequence); +} + +/* +// TODO: enable the following tests after removing the limitations of CUDA kernels. +TEST(LongformerAttentionTest, LongformerAttention_GlobalMiddle) { + std::vector mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f}; + + // Global in the middle of sequence + std::vector global_data = {0, 1, 0, 0, 1, 0, 0, 0}; + + std::vector output_data = { + 0.0909f, 0.0482f, 0.0273f, 0.0339f, 0.0124f, -0.0186f, 0.0056f, -0.0441f, + -0.0356f, 0.0521f, -0.0198f, 0.0005f, 0.0245f, -0.0014f, -0.0169f, -0.0122f, + 0.0738f, 0.0416f, 0.0322f, 0.0258f, 0.0063f, -0.0222f, 0.0053f, -0.0309f, + 0.0725f, 0.0451f, 0.0149f, 0.0241f, -0.0071f, -0.0231f, 0.0028f, -0.0369f, + -0.0355f, 0.0521f, -0.0198f, 0.0004f, 0.0245f, -0.0014f, -0.0169f, -0.0122f, + 0.0653f, 0.0431f, 0.0032f, 0.0156f, -0.0061f, -0.0293f, -0.0188f, -0.0404f, + 0.0803f, 0.0502f, -0.0089f, 0.0212f, -0.0030f, -0.0275f, -0.0244f, -0.0560f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f}; + + RunTinyLongformerBatch1(mask_data, global_data, output_data, false); +} +*/ + +} // namespace test +} // namespace onnxruntime diff --git a/tools/ci_build/amd_hipify.py b/tools/ci_build/amd_hipify.py index 7b1de841f7..bf26506cb0 100644 --- a/tools/ci_build/amd_hipify.py +++ b/tools/ci_build/amd_hipify.py @@ -20,6 +20,10 @@ contrib_ops_files = [ 'bert/embed_layer_norm_impl.h', 'bert/fast_gelu_impl.cu', 'bert/layer_norm.cuh', + 'bert/longformer_attention.cc', + 'bert/longformer_attention.h', + 'bert/longformer_attention_impl.cu', + 'bert/longformer_attention_impl.h', 'bert/skip_layer_norm.cc', 'bert/skip_layer_norm.h', 'bert/skip_layer_norm_impl.cu',