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

Add Longformer Attention Cuda Op(#5932)

Browse source 
Limitation: Global tokens must be at the beginning of sequence.
This commit is contained in:
Tianlei Wu 2020-11-25 13:52:10 -08:00 committed by GitHub
parent e39e82b43a
commit 31a6be3d67
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 1552 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

123
onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.cc Normal file
View file

@ -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<int>(num_heads);
int64_t window = 0;
ORT_ENFORCE(info.GetAttr("window", &window).IsOK() && window > 0);
window_ = static_cast<int>(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<int>(dims[0]);
int sequence_length = static_cast<int>(dims[1]);
int hidden_size = static_cast<int>(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<int>(mask_dims[0]) != batch_size || static_cast<int>(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<int>(global_dims[0]) != batch_size || static_cast<int>(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

29
onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h Normal file
View file

@ -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

138
onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc Normal file
View file

@ -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<T>()), \
LongformerAttention<T>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(MLFloat16)
template <typename T>
LongformerAttention<T>::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {}
template <typename T>
Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
const Tensor* input = context->Input<Tensor>(0);
const Tensor* weights = context->Input<Tensor>(1);
const Tensor* bias = context->Input<Tensor>(2);
const Tensor* mask = context->Input<Tensor>(3);
const Tensor* global_weights = context->Input<Tensor>(4);
const Tensor* global_bias = context->Input<Tensor>(5);
const Tensor* global_attention = context->Input<Tensor>(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<int>(shape[0]);
int sequence_length = static_cast<int>(shape[1]);
int hidden_size = static_cast<int>(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<T>(qkv_size);
typedef typename ToCudaType<T>::MappedType CudaT;
CudaT one = ToCudaType<T>::FromFloat(1.0f);
CudaT zero = ToCudaType<T>::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<const CudaT*>(bias->template Data<T>()), n,
GetConstOnes<CudaT>(m), 1,
&zero, reinterpret_cast<CudaT*>(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<const CudaT*>(weights->template Data<T>()), n,
reinterpret_cast<const CudaT*>(input->template Data<T>()), k,
&one, reinterpret_cast<CudaT*>(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<T>(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<const CudaT*>(global_bias->template Data<T>()), n,
GetConstOnes<CudaT>(m), 1,
&zero, reinterpret_cast<CudaT*>(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<const CudaT*>(global_weights->template Data<T>()), n,
reinterpret_cast<const CudaT*>(input->template Data<T>()), k,
&one, reinterpret_cast<CudaT*>(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<void>(workSpaceSize);
if (!LaunchLongformerAttentionKernel(
device_prop,
reinterpret_cast<const CudaT*>(gemm_buffer.get()),
reinterpret_cast<const CudaT*>(mask->template Data<T>()),
reinterpret_cast<const CudaT*>(global_gemm_buffer.get()),
global_attention->template Data<int>(),
output->template MutableData<T>(),
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

26
onnxruntime/contrib_ops/cuda/bert/longformer_attention.h Normal file
View file

@ -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 <typename T>
class LongformerAttention final : public CudaKernel, public LongformerAttentionBase {
public:
LongformerAttention(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* context) const override;
};
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

859
onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu Normal file
View file

@ -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 <cub/cub.cuh>
#include <cublas_v2.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <library_types.h>
#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<int*>(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<<<batch_size, 128, 0, stream>>>(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<size_t>(3 * batch_size * sequence_length));
for (int i = 0; i < batch_size; ++i) {
cub::DevicePartition::Flagged(reinterpret_cast<void*>(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 <typename T, int blockSize>
__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<float, blockSize> 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<char*>(workspace) + 3 * sizeof(int) * batch_size * sequence_length;
void* scratch2 = reinterpret_cast<char*>(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<void*>(&one_fp16);
beta_0 = static_cast<void*>(&zero_fp16);
beta_1 = static_cast<void*>(&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<void*>(&one_fp32);
beta_0 = static_cast<void*>(&zero_fp32);
beta_1 = static_cast<void*>(&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><<<gridSize, blockSize, 0, stream>>>(
global_attention,
global_idx,
batch_global_num,
static_cast<const __half*>(scratch1),
static_cast<const __half*>(attention_mask),
static_cast<__half*>(softmax_out), scaler, dim0, dim1, attention_window);
} else {
LongformerSoftmaxKernel<float, blockSize><<<gridSize, blockSize, 0, stream>>>(
global_attention,
global_idx,
batch_global_num,
static_cast<const float*>(scratch1),
static_cast<const float*>(attention_mask),
static_cast<float*>(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 <typename T>
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<T*>((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<float>(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<const half*>(input),
reinterpret_cast<const half*>(attention_mask),
reinterpret_cast<const half*>(global_input),
global_attention,
max_num_global,
reinterpret_cast<half*>(workspace),
reinterpret_cast<half*>(output));
} else {
return LongformerQkvToContext(prop, cublas, stream,
batch_size, sequence_length, num_heads, head_size, window, element_size,
reinterpret_cast<const float*>(input),
reinterpret_cast<const float*>(attention_mask),
reinterpret_cast<const float*>(global_input),
global_attention,
max_num_global,
reinterpret_cast<float*>(workspace),
reinterpret_cast<float*>(output));
}
}
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

39
onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h Normal file
View file

@ -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

4
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ImageScaler)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ImageScaler)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ImageScaler)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, LongformerAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, LongformerAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ParametricSoftplus)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ParametricSoftplus)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ParametricSoftplus)>,

29
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -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,

301
onnxruntime/test/contrib_ops/longformer_attention_op_test.cc Normal file
View file

@ -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<float>& input_data, // input: [batch_size, sequence_length, hidden_size]
const std::vector<float>& weights_data, // weights: [hidden_size, 3 * hidden_size]
const std::vector<float>& bias_data, // bias: [3 * hidden_size]
const std::vector<float>& mask_data, // mask: [batch_size, sequence_length]
const std::vector<float>& global_weights_data, // global_weights: [hidden_size, 3 * hidden_size]
const std::vector<float>& global_bias_data, // global_bias: [3 * hidden_size]
const std::vector<int>& global_data, // global: [batch_size, sequence_length]
const std::vector<float>& 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<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
tester.AddAttribute<int64_t>("window", static_cast<int64_t>(window));
std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};
std::vector<int64_t> bias_dims = {3 * hidden_size};
std::vector<int64_t> mask_dims = {batch_size, sequence_length};
std::vector<int64_t> global_weights_dims = {hidden_size, 3 * hidden_size};
std::vector<int64_t> global_bias_dims = {3 * hidden_size};
std::vector<int64_t> global_dims = {batch_size, sequence_length};
std::vector<int64_t> output_dims = input_dims;
if (use_float16) {
tester.AddInput<MLFloat16>("input", input_dims, ToFloat16(input_data));
tester.AddInput<MLFloat16>("weight", weights_dims, ToFloat16(weights_data));
tester.AddInput<MLFloat16>("bias", bias_dims, ToFloat16(bias_data));
tester.AddInput<MLFloat16>("mask", mask_dims, ToFloat16(mask_data));
tester.AddInput<MLFloat16>("global_weight", global_weights_dims, ToFloat16(global_weights_data));
tester.AddInput<MLFloat16>("global_bias", global_bias_dims, ToFloat16(global_bias_data));
tester.AddOutput<MLFloat16>("output", output_dims, ToFloat16(output_data));
} else {
tester.AddInput<float>("input", input_dims, input_data);
tester.AddInput<float>("weight", weights_dims, weights_data);
tester.AddInput<float>("bias", bias_dims, bias_data);
tester.AddInput<float>("mask", mask_dims, mask_data);
tester.AddInput<float>("global_weight", global_weights_dims, global_weights_data);
tester.AddInput<float>("global_bias", global_bias_dims, global_bias_data);
tester.AddOutput<float>("output", output_dims, output_data);
}
tester.AddInput<int32_t>("global", global_dims, global_data);
if (enable_cuda) {
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
if (enable_cpu) {
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCpuExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
}
}
static void GetTinyLongformerData(
std::vector<float>& weight_data,
std::vector<float>& bias_data,
std::vector<float>& global_weight_data,
std::vector<float>& 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<float>& mask_data,
std::vector<int>& global_data,
std::vector<float>& 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<float> 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<float> weight_data;
std::vector<float> bias_data;
std::vector<float> global_weight_data;
std::vector<float> 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<float> mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f};
// no global attention.
std::vector<int> global_data = {0, 0, 0, 0, 0, 0, 0, 0};
std::vector<float> 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<float> 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<int> global_data = {1, 1, 0, 0, 0, 0, 0, 0};
std::vector<float> 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<float> 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<int> global_data = {1, 1, 0, 0, 0, 0, 0, 0};
std::vector<float> 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<float> mask_data = {0.0f, 0.0f, 0.0f, -10000.0f};
// no global attention.
std::vector<int> global_data = {0, 0, 0, 0};
std::vector<float> 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<float> 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<int> global_data = {0, 1, 0, 0, 1, 0, 0, 0};
std::vector<float> 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

4
tools/ci_build/amd_hipify.py
View file

@ -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',