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

Support long sequence in attention (#14371)

Browse source 
Support long sequence in attention operator for (1) raw mask of 2/3/4-D,
(2) no mask.
Set longer greedy search max length.
This commit is contained in:
Zhang Lei 2023-01-27 09:39:09 -08:00 committed by GitHub
parent 368d2fc11e
commit f87dd408f6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 195 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/cpu/transformers/greedy_search_parameters.cc
View file

@ -6,7 +6,7 @@ namespace onnxruntime {
namespace contrib {
namespace transformers {
constexpr int kMaxSequenceLength = 4096;
constexpr int kMaxSequenceLength = 16384;
void GreedySearchParameters::ParseFromAttributes(const OpKernelInfo& info) {
model_type = static_cast<int>(info.GetAttrOrDefault<int64_t>("model_type", 0));

196
onnxruntime/contrib_ops/cuda/bert/attention_softmax.h
View file

@ -169,6 +169,189 @@ __device__ inline void SoftmaxSmall(const int all_sequence_length,
}
}
template <typename T, unsigned TPB>
__global__ void SoftmaxLargeKernel(const int all_sequence_length,
const int sequence_length,
const int valid_end,
const int valid_start,
const T* add_before_softmax,
const T* input,
T* output,
bool is_unidirectional) {
extern __shared__ float cached_data[]; // float[all_sequence_length]
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
__shared__ float sum_reverse_block;
__shared__ float max_block;
// Update end position for unidirectional.
int end = valid_end;
int end_unid = -1;
if (is_unidirectional) {
end_unid = all_sequence_length - sequence_length + (blockIdx.x % sequence_length) + 1;
if (end_unid <= valid_start) {
;
// In this situation, mask of [0, end_unid) and [valid_start, valid_end) has -10000,
// and [end_unid, valid_start) and [valid_end, all_seq_len) has -20000.
// So [0, end_unid) will also have value after softmax.
// KEEP SMALL KERNEL CODE LOGIC HERE as COMMENT
// is_valid = threadIdx.x < end_unid; // is_valid initialized with false
} else {
end = min(valid_end, end_unid);
}
}
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
float thread_data_max = -CUDART_INF_F;
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
const int index = offset + seq_idx;
bool is_valid = (seq_idx < end_unid) || (seq_idx >= valid_start && seq_idx < end);
// e^x is represented as infinity if x is large enough, like 100.f.
// Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough.
// a math transform as below is leveraged to get a stable softmax:
// e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max))
float input_data = is_valid
? (add_before_softmax ? float(input[index] + add_before_softmax[index]) : float(input[index]))
: float(-CUDART_INF_F);
cached_data[seq_idx] = input_data;
thread_data_max = max(thread_data_max, input_data);
}
const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end);
// Store max value
if (threadIdx.x == 0) {
max_block = max;
}
__syncthreads();
float thread_data_exp(0.f);
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
bool is_valid = (seq_idx < end_unid) || (seq_idx >= valid_start && seq_idx < end);
cached_data[seq_idx] = is_valid ? expf(cached_data[seq_idx] - max_block) : 0.0f;
thread_data_exp += cached_data[seq_idx];
}
const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end);
// Store value of 1.0/sum.
if (threadIdx.x == 0) {
sum_reverse_block = (1.f) / sum;
}
__syncthreads();
// threadIdx.x might be larger than all_sequence_length due to alignment to 32x.
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
output[offset + seq_idx] = T(cached_data[seq_idx] * sum_reverse_block);
}
}
template <typename T, int TPB>
__global__ void SoftmaxWithRawMaskLargeKernel(const int all_sequence_length,
const int sequence_length,
const int* attention_mask, // 2D, 3D or 4D attention mask
const bool* key_padding_mask,
const T* add_before_softmax,
const T* input,
T* output,
const bool is_unidirectional,
const float rsqrt_head_size,
const int mask_dimension,
const int max_sequence_length,
const bool skip_softmax,
const float mask_filter_value) {
extern __shared__ float cached_data[]; // float[all_sequence_length]
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
__shared__ float sum_reverse_block;
__shared__ float max_block;
float max_thread_data = -CUDART_INF_F;
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
int base_index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
float thread_data = -CUDART_INF_F;
int index = base_index + seq_idx;
if (add_before_softmax == nullptr) {
thread_data = float(input[index]) * rsqrt_head_size;
} else {
thread_data = float(input[index] + add_before_softmax[index]) * rsqrt_head_size;
}
const int sequence_index = blockIdx.x % sequence_length;
if (is_unidirectional) {
int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length.
if (seq_idx > from_index) {
thread_data = mask_filter_value;
}
}
int mask_offset = 0;
const int batch_index = blockIdx.y;
if (mask_dimension == 2) {
mask_offset = batch_index * all_sequence_length + seq_idx;
} else if (mask_dimension == 3) {
mask_offset = (batch_index * sequence_length + sequence_index) * all_sequence_length + seq_idx;
} else if (mask_dimension == 4) {
int from_index = all_sequence_length - sequence_length + sequence_index;
mask_offset = (batch_index * max_sequence_length + from_index) * max_sequence_length + seq_idx;
}
if (nullptr == key_padding_mask) {
const int& mask = attention_mask[mask_offset];
if (mask == 0)
thread_data += mask_filter_value;
} else {
const bool mask = key_padding_mask[mask_offset];
if (mask) {
thread_data = -CUDART_INF_F;
}
}
if (skip_softmax) {
output[index] = T(thread_data);
}
cached_data[seq_idx] = thread_data;
max_thread_data = max(max_thread_data, thread_data);
}
if (skip_softmax) {
return;
}
const float max = BlockReduce(tmp_storage).Reduce(max_thread_data, cub::Max(), all_sequence_length);
// Store max value
if (threadIdx.x == 0) {
max_block = max;
}
__syncthreads();
float sum_thread_data_exp = 0.0f;
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
auto ev = expf(cached_data[seq_idx] - max_block);
cached_data[seq_idx] = ev;
sum_thread_data_exp += ev;
}
const auto sum = BlockReduce(tmp_storage).Reduce(sum_thread_data_exp, cub::Sum(), TPB);
// Store value of 1.0/sum
if (threadIdx.x == 0) {
sum_reverse_block = (1.f) / sum;
}
__syncthreads();
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
output[base_index + seq_idx] = T(cached_data[seq_idx] * sum_reverse_block);
}
}
template <typename T, unsigned TPB>
__device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
const int sequence_length,
@ -315,7 +498,10 @@ Status ComputeSoftmax(cudaStream_t stream, const int all_sequence_length, const
SoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024.");
const int blockSize = 256;
const int sh_bytes = sizeof(float) * all_sequence_length;
SoftmaxLargeKernel<T, blockSize><<<grid, blockSize, sh_bytes, stream>>>(
all_sequence_length, sequence_length, all_sequence_length, 0, add_before_softmax, input, output, true);
}
return CUDA_CALL(cudaGetLastError());
@ -518,7 +704,13 @@ Status ComputeSoftmaxWithRawMask(cudaStream_t stream,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024.");
const int blockSize = 256;
const int sh_bytes = sizeof(float) * all_sequence_length;
SoftmaxWithRawMaskLargeKernel<T, blockSize>
<<<grid, blockSize, sh_bytes, stream>>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
}
if (use_persistent_softmax) {