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

Add PackedAttention for packing mode (#14858)

Browse source 
### Description
<!-- Describe your changes. -->
Transformer models can handle batch of inputs at once. However,
sequences in a batch usually have different length. Then we have to pad
the short one to have same length as the longest. This is not efficient
especially for large batch with high variance.

This PR introduces a PackedAttention operator which can take in packed
sequences (no padding) and also produces output in packing mode.

There will be another PR to use the PackedAttention to implement the
encoder in packing mode.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
Yufeng Li 2023-03-21 12:59:29 -07:00 committed by GitHub
parent ef76b3aeb8
commit c7ced7a5e9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
17 changed files with 1915 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
docs/ContribOperators.md
View file

@ -51,6 +51,7 @@ Do not modify directly.*
* <a href="#com.microsoft.NGramRepeatBlock">com.microsoft.NGramRepeatBlock</a>
* <a href="#com.microsoft.NhwcConv">com.microsoft.NhwcConv</a>
* <a href="#com.microsoft.NhwcMaxPool">com.microsoft.NhwcMaxPool</a>
* <a href="#com.microsoft.PackedAttention">com.microsoft.PackedAttention</a>
* <a href="#com.microsoft.Pad">com.microsoft.Pad</a>
* <a href="#com.microsoft.QAttention">com.microsoft.QAttention</a>
* <a href="#com.microsoft.QGemm">com.microsoft.QGemm</a>
@ -2607,6 +2608,78 @@ This version of the operator has been available since version 1 of the 'com.micr
</dl>
### <a name="com.microsoft.PackedAttention"></a><a name="com.microsoft.packedattention">**com.microsoft.PackedAttention**</a>
This is the packed version of Attention.
Sequences in one batch usually don't have same length and they are padded to have same length,
e.g., below is a batch with 3 sequences and tokens* are padded.
Sequence_0: 0, 1*, 2*, 3*
Sequence_1: 4, 5, 6*, 7*
Sequence_2: 8, 9, 10, 11
PackedAttention is designed to takes in packed input, i.e., only the real tokens without padding.
An input as above will be packed into 3 tensors like below:
- input ([h0, h4, h5, h8, h9, h10, h11])
- token_offset: 0, 4, 5, 8, 9, 10, 11, 1*, 2*, 3*, 6*, 7*
- cumulated_token_count: 0, 1, 1+2, 1+2+4
Input tensors contains the hidden embedding of real tokens.
Token_offset records the offset of token in the unpacked input.
cumulated_token_count records cumulated length of each sequnces length.
The operator only supports BERT like model with padding on right now.
#### Version
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
#### Attributes
<dl>
<dt><tt>num_heads</tt> : int (required)</dt>
<dd>Number of attention heads</dd>
<dt><tt>qkv_hidden_sizes</tt> : list of ints</dt>
<dd>Hidden dimension of Q, K, V: hidden_size, hidden_size and v_hidden_size</dd>
<dt><tt>scale</tt> : float</dt>
<dd>Custom scale will be used if specified. Default value is 1/sqrt(head_size)</dd>
</dl>
#### Inputs (5 - 6)
<dl>
<dt><tt>input</tt> : T</dt>
<dd>Input tensor with shape (token_count, input_hidden_size)</dd>
<dt><tt>weights</tt> : T</dt>
<dd>Merged Q/K/V weights with shape (input_hidden_size, hidden_size + hidden_size + v_hidden_size)</dd>
<dt><tt>bias</tt> : T</dt>
<dd>Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) for input projection</dd>
<dt><tt>token_offset</tt> : M</dt>
<dd>In packing mode, it specifies the offset of each token(batch_size, sequence_length).</dd>
<dt><tt>cumulative_sequence_length</tt> : M</dt>
<dd>A tensor with shape (batch_size + 1). It specifies the cumulative sequence length.</dd>
<dt><tt>relative_position_bias</tt> (optional) : T</dt>
<dd>A tensor with shape (batch_size, num_heads, sequence_length, sequence_length)or (1, num_heads, sequence_length, sequence_length).It specifies the additional bias to QxK'</dd>
</dl>
#### Outputs
<dl>
<dt><tt>output</tt> : T</dt>
<dd>2D output tensor with shape (token_count, v_hidden_size)</dd>
</dl>
#### Type Constraints
<dl>
<dt><tt>T</tt> : tensor(float), tensor(float16)</dt>
<dd>Constrain input and output types to float tensors.</dd>
<dt><tt>M</tt> : tensor(int32)</dt>
<dd>Constrain mask index to integer types</dd>
</dl>
### <a name="com.microsoft.Pad"></a><a name="com.microsoft.pad">**com.microsoft.Pad**</a>
Given `data` tensor, pads, mode, and value.

1
docs/OperatorKernels.md
View file

@ -816,6 +816,7 @@ Do not modify directly.*
|MultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* bias:**T**<br> *in* key_padding_mask:**M**<br> *in* relative_position_bias:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**T** = tensor(float), tensor(float16)|
|NGramRepeatBlock|*in* input_ids:**Tid**<br> *in* scores:**T**<br> *out* scores_out:**T**|1+|**T** = tensor(float)<br/> **Tid** = tensor(int64)|
|NhwcConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|PackedAttention|*in* input:**T**<br> *in* weights:**T**<br> *in* bias:**T**<br> *in* token_offset:**M**<br> *in* cumulative_sequence_length:**M**<br> *in* relative_position_bias:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(int8)<br/> **T2** = tensor(int8)<br/> **T3** = tensor(float), tensor(float16)<br/> **T4** = tensor(int32)|
|QOrderedAttention|*in* input:**Q**<br> *in* scale_input:**S**<br> *in* scale_Q_gemm:**S**<br> *in* scale_K_gemm:**S**<br> *in* scale_V_gemm:**S**<br> *in* Q_weight:**Q**<br> *in* K_weight:**Q**<br> *in* V_weight:**Q**<br> *in* scale_Q_weight:**S**<br> *in* scale_K_weight:**S**<br> *in* scale_V_weight:**S**<br> *in* Q_bias:**S**<br> *in* K_bias:**S**<br> *in* V_bias:**S**<br> *in* scale_QKT_gemm:**S**<br> *in* scale_QKT_softmax:**S**<br> *in* scale_values_gemm:**S**<br> *in* mask_index:**G**<br> *in* past:**Q**<br> *in* relative_position_bias:**S**<br> *out* output:**Q**|1+|**G** = tensor(int32)<br/> **Q** = tensor(int8)<br/> **S** = tensor(float)|
|QOrderedGelu|*in* X:**Q**<br> *in* scale_X:**S**<br> *in* scale_Y:**S**<br> *out* Y:**Q**|1+|**Q** = tensor(int8)<br/> **S** = tensor(float)|

16
onnxruntime/contrib_ops/cpu/bert/attention_common.h
View file

@ -59,6 +59,22 @@ struct AttentionParameters {
AttentionMaskType mask_type;
};
// Parameters deduced from node attributes and inputs/outputs.
struct PackedAttentionParameters {
int batch_size;
int sequence_length;
int input_hidden_size; // hidden size of input
int hidden_size; // hidden size of Q or K
int head_size; // hidden size per head of Q or K
int v_hidden_size; // hidden size of V
int v_head_size; // hidden size per head of V
int num_heads;
float scale;
int token_count;
bool has_relative_position_bias;
bool broadcast_res_pos_bias;
};
namespace attention {
// Environment variable to enable or disable TRT fused self attention kernel. Default is 0 (enabled).
constexpr const char* kDisableFusedSelfAttention = "ORT_DISABLE_FUSED_ATTENTION";

2
onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
View file

@ -297,7 +297,7 @@ Status PrepareQkv(contrib::AttentionParameters& parameters,
bool use_fused_causal = (nullptr != fused_runner && parameters.is_unidirectional);
// Default format for memory efficient attention.
// When there is past state, the format shal be BxNxSxH, so we disable memory efficient attention when there is past.
// When there is past state, the format shall be BxNxSxH, so we disable memory efficient attention when there is past.
DUMP_TENSOR_INIT();
if (nullptr != data.gemm_buffer) {
if (data.bias == nullptr) {

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

@ -36,7 +36,6 @@ namespace cuda {
template <typename T, unsigned TPB>
__device__ inline void Softmax(const int all_sequence_length,
const int sequence_length,
const int valid_end,
const int valid_start,
const T* rel_pos_bias,
@ -483,12 +482,11 @@ __global__ void SoftmaxKernelSmall(const int all_sequence_length,
template <typename T, unsigned TPB>
__global__ void SoftmaxKernel(const int all_sequence_length,
const int sequence_length,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output) {
Softmax<T, TPB>(all_sequence_length, sequence_length, all_sequence_length, 0,
Softmax<T, TPB>(all_sequence_length, all_sequence_length, 0,
rel_pos_bias, broadcast_rel_pos_bias, input, output);
}
@ -524,7 +522,7 @@ Status ComputeSoftmax(cudaStream_t stream, const int all_sequence_length, const
} else if (!is_unidirectional) {
const int blockSize = 1024;
SoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output);
all_sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output);
} else {
const int blockSize = 256;
const int sh_bytes = sizeof(float) * all_sequence_length;
@ -566,9 +564,102 @@ __global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length,
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
}
template <typename T, unsigned TPB>
__device__ inline void SoftmaxSmallPacked(const int sequence_length,
const int end,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
__shared__ float sum_reverse_block;
__shared__ float max_block;
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length;
const int index = offset + threadIdx.x;
bool is_valid = threadIdx.x < end;
// e^x is represented as infinity if x is large enough, like 100.f.
// Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough.
// a math transform as below is leveraged to get a stable softmax:
// e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max))
const bool no_rpb = (rel_pos_bias == nullptr);
const int size_per_batch = gridDim.x * sequence_length;
float input_data = no_rpb
? float(input[index])
: float(input[index] + (broadcast_rel_pos_bias
? rel_pos_bias[index % size_per_batch]
: rel_pos_bias[index]));
float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F);
const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end);
// Store max value
if (threadIdx.x == 0) {
max_block = max;
}
__syncthreads();
float thread_data_exp(0.f);
if (is_valid) {
thread_data_exp = expf(input_data - max_block);
}
const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end);
// Store value of 1.0/sum.
if (threadIdx.x == 0) {
sum_reverse_block = (1.f) / sum;
}
__syncthreads();
// threadIdx.x might be larger than all_sequence_length due to alignment to 32x.
if (threadIdx.x < sequence_length) {
output[index] = T(thread_data_exp * sum_reverse_block);
}
}
template <typename T, unsigned TPB>
__global__ void SoftmaxKernelSmallWithCumSeqLen(const T* input,
const T* rel_pos_bias, const bool broadcast_rel_pos_bias,
const int* cum_seq_length, const int sequence_length,
T* output) {
__shared__ int end_position;
if (threadIdx.x == 0) {
const int batch = blockIdx.y;
end_position = cum_seq_length[batch + 1] - cum_seq_length[batch];
}
__syncthreads();
SoftmaxSmallPacked<T, TPB>(sequence_length, end_position,
rel_pos_bias, broadcast_rel_pos_bias,
input, output);
}
template <typename T, unsigned TPB>
__global__ void SoftmaxKernelWithCumSeqLen(const T* input,
const T* rel_pos_bias, const bool broadcast_rel_pos_bias,
const int* cum_seq_length, const int sequence_length,
T* output) {
__shared__ int end_position;
if (threadIdx.x == 0) {
const int batch = blockIdx.y;
end_position = cum_seq_length[batch + 1] - cum_seq_length[batch];
}
__syncthreads();
Softmax<T, TPB>(sequence_length, end_position, 0 /*start_position*/,
rel_pos_bias, broadcast_rel_pos_bias, input, output);
}
template <typename T, unsigned TPB>
__global__ void MaskedSoftmaxKernel(const int all_sequence_length,
const int sequence_length,
const int* mask_end,
const int* mask_start,
const T* rel_pos_bias,
@ -590,7 +681,7 @@ __global__ void MaskedSoftmaxKernel(const int all_sequence_length,
}
__syncthreads();
Softmax<T, TPB>(all_sequence_length, sequence_length, end_position, start_position,
Softmax<T, TPB>(all_sequence_length, end_position, start_position,
rel_pos_bias, broadcast_rel_pos_bias, input, output);
}
@ -616,6 +707,58 @@ __global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length,
skip_softmax, mask_filter_value);
}
template <typename T>
Status ComputeSoftmaxWithCumSeqLength(
const T* input,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const int32_t* cum_seq_length,
const int batch_size,
const int sequence_length,
const int num_heads,
T* output, cudaStream_t stream) {
const dim3 grid(sequence_length * num_heads, batch_size, 1);
if (sequence_length <= 32) {
const int blockSize = 32;
SoftmaxKernelSmallWithCumSeqLen<T, blockSize>
<<<grid, blockSize, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
} else if (sequence_length <= 64) {
const int blockSize = 64;
SoftmaxKernelSmallWithCumSeqLen<T, blockSize>
<<<grid, blockSize, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
} else if (sequence_length <= 128) {
const int blockSize = 128;
SoftmaxKernelSmallWithCumSeqLen<T, blockSize>
<<<grid, blockSize, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
} else if (sequence_length <= 256) {
const int blockSize = 256;
SoftmaxKernelSmallWithCumSeqLen<T, blockSize>
<<<grid, blockSize, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
} else if (sequence_length <= 512) {
const int blockSize = 512;
SoftmaxKernelSmallWithCumSeqLen<T, blockSize>
<<<grid, blockSize, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
} else if (sequence_length <= 1024) {
const int blockSize = 1024;
SoftmaxKernelSmallWithCumSeqLen<T, blockSize>
<<<grid, blockSize, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
} else {
SoftmaxKernelWithCumSeqLen<T, 1024>
<<<grid, 1024, 0, stream>>>(input, rel_pos_bias, broadcast_rel_pos_bias,
cum_seq_length, sequence_length, output);
}
return CUDA_CALL(cudaGetLastError());
}
template <typename T>
Status ComputeSoftmaxWithMask1D(cudaStream_t stream,
const int all_sequence_length,
@ -664,7 +807,7 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream,
} else if (!is_unidirectional) {
const int blockSize = 1024;
MaskedSoftmaxKernel<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
<<<grid, blockSize, 0, stream>>>(all_sequence_length, mask_index, mask_start,
rel_pos_bias, broadcast_rel_pos_bias, input, output);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024.");

324
onnxruntime/contrib_ops/cuda/bert/packed_attention.cc Normal file
View file

@ -0,0 +1,324 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/cuda/bert/packed_attention.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/shared_inc/fpgeneric.h"
#include "core/platform/env_var_utils.h"
#include "contrib_ops/cuda/bert/packed_attention_impl.h"
#include "contrib_ops/cuda/bert/bert_padding.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.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( \
PackedAttention, \
kMSDomain, \
1, \
T, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
PackedAttention<T>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(MLFloat16)
template <typename T>
PackedAttention<T>::PackedAttention(const OpKernelInfo& info) : CudaKernel(info) {
int64_t num_heads = 0;
ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0);
num_heads_ = static_cast<int32_t>(num_heads);
scale_ = info.GetAttrOrDefault<float>("scale", 0.0f);
if (!info.GetAttrs<int64_t>("qkv_hidden_sizes", qkv_hidden_sizes_).IsOK()) {
qkv_hidden_sizes_.clear();
}
disable_fused_runner_ = sizeof(T) != 2 ||
ParseEnvironmentVariableWithDefault<bool>(attention::kDisableFusedSelfAttention, false);
enable_trt_flash_attention_ = sizeof(T) == 2 &&
!ParseEnvironmentVariableWithDefault<bool>(attention::kDisableTrtFlashAttention, false);
}
template <typename T>
Status PackedAttention<T>::CheckInputs(const TensorShape& input_shape,
const TensorShape& weights_shape,
const TensorShape& bias_shape,
const TensorShape& token_offset_shape,
const TensorShape& cu_seq_len_shape,
const Tensor* relative_position_bias,
PackedAttentionParameters& parameters) const {
// Abbreviation and Meanings:
// T: token_count
// B: batch_size
// S: sequence_length (input sequence length of query)
// N: num_heads
// H: head size for Q and K, aka q_head_size or v_head_size or qk_head_size
// H_v: v_head_size
// D_i: input hidden size
// D: hidden size for Q and K (D = N * H), aka q_hidden_size or k_hidden_size or qk_hidden_size
// D_v: v_hidden_size = num_heads * v_head_size
// Input shapes:
// input: : (T, D_i)
// weights (Q/K/V) : (D_i, D + D + D_v)
// bias (Q/K/V) : (D + D + D_v)
// token_offset : (B, S)
// cu_seq_len_shape : (B + 1)
// relative_position_bias : (B, N, S, S), (1, N, S, S) or NULL
const auto& input_dims = input_shape.GetDims();
if (input_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'input' is expected to have 2 dimensions in packing mode, got ",
input_dims.size());
}
int64_t token_count = input_dims[0];
int64_t input_hidden_size = input_dims[1];
const auto& token_offset_dims = token_offset_shape.GetDims();
if (token_offset_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'packing_token_offset' is expected to have 2 dimensions in packing mode, got ",
token_offset_dims.size());
}
int64_t batch_size = token_offset_dims[0];
int64_t sequence_length = token_offset_dims[1];
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());
}
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] != input_hidden_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 1 dimension 0 should have same length as dimension 2 of input 0");
}
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& cu_seq_len_dims = cu_seq_len_shape.GetDims();
if (cu_seq_len_dims.size() != 1 || cu_seq_len_dims[0] != batch_size + 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'cumulative_sequence_length' should have 1 dimension with size equal to batch_size + 1");
}
int64_t q_hidden_size = bias_dims[0] / static_cast<int64_t>(3);
int64_t k_hidden_size = q_hidden_size;
int64_t v_hidden_size = k_hidden_size;
if (qkv_hidden_sizes_.size() != 0) {
if (qkv_hidden_sizes_.size() != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"qkv_hidden_sizes attribute should have 3 elements");
}
for (size_t i = 0; i < qkv_hidden_sizes_.size(); i++) {
if (qkv_hidden_sizes_[i] % num_heads_ != 0) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"hidden_size should be divisible by num_heads:", qkv_hidden_sizes_[i]);
}
}
q_hidden_size = qkv_hidden_sizes_[0];
k_hidden_size = qkv_hidden_sizes_[1];
v_hidden_size = qkv_hidden_sizes_[2];
}
if (q_hidden_size != k_hidden_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"qkv_hidden_sizes first element should be same as the second");
}
if (bias_dims[0] != q_hidden_size + k_hidden_size + v_hidden_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'bias' dimension 0 should have same length as sum of Q/K/V hidden sizes:",
" q_hidden_size=", q_hidden_size, " k_hidden_size=", k_hidden_size, " v_hidden_size=",
v_hidden_size, "bias_dims[0]=", bias_dims[0]);
}
bool broadcast_res_pos_bias = false;
if (relative_position_bias != nullptr) {
const auto& relative_position_bias_dims = relative_position_bias->Shape().GetDims();
if (relative_position_bias_dims.size() != 4) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'relative_position_bias' is expected to have 4 dimensions, got ",
relative_position_bias_dims.size());
}
if (relative_position_bias_dims[0] != batch_size && relative_position_bias_dims[0] != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'relative_position_bias' dimension 0 should be same as batch_size or 1, got ",
relative_position_bias_dims[0]);
}
if (relative_position_bias_dims[0] == 1) {
broadcast_res_pos_bias = true;
}
if (relative_position_bias_dims[1] != num_heads_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'relative_position_bias' dimension 1 should be same as number of heads, got ",
relative_position_bias_dims[1]);
}
if (relative_position_bias_dims[2] != sequence_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'relative_position_bias' dimension 2 should be same as sequence_length, got ",
relative_position_bias_dims[2]);
}
if (relative_position_bias_dims[3] != sequence_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'relative_position_bias' dimension 3 should be same as sequence_length, got ",
relative_position_bias_dims[3]);
}
}
parameters.batch_size = static_cast<int>(batch_size);
parameters.sequence_length = static_cast<int>(sequence_length);
parameters.input_hidden_size = static_cast<int>(input_hidden_size);
parameters.hidden_size = static_cast<int>(q_hidden_size);
parameters.v_hidden_size = static_cast<int>(v_hidden_size);
parameters.head_size = static_cast<int>(q_hidden_size) / num_heads_;
parameters.v_head_size = static_cast<int>(v_hidden_size) / num_heads_;
parameters.num_heads = num_heads_;
parameters.scale = scale_;
parameters.token_count = static_cast<int32_t>(token_count);
parameters.has_relative_position_bias = nullptr != relative_position_bias;
parameters.broadcast_res_pos_bias = broadcast_res_pos_bias;
return Status::OK();
}
template <typename T>
MHARunner* PackedAttention<T>::TryGettingFusedRunner(const PackedAttentionParameters& parameters) const {
MHARunner* fused_runner = nullptr;
bool use_fused_runner = !disable_fused_runner_ &&
!parameters.has_relative_position_bias &&
parameters.hidden_size == parameters.v_hidden_size;
if(!use_fused_runner) {
return fused_runner;
}
// Check whether we can use fused kernel
auto& device_prop = GetDeviceProp();
int sm = device_prop.major * 10 + device_prop.minor;
bool is_fMHA_supported = FusedMHARunnerFP16v2::is_supported(sm,
parameters.head_size,
parameters.sequence_length,
enable_trt_flash_attention_,
false);
if(!is_fMHA_supported) {
return fused_runner;
}
// Assuming that num_heads and head_size do not change.
if (nullptr == fused_fp16_runner_.get()) {
fused_fp16_runner_.reset(new FusedMHARunnerFP16v2(num_heads_, parameters.head_size, sm, false /* causal_mask*/,
enable_trt_flash_attention_, parameters.scale));
}
// In case some kernel not loaded due to shared memory limit, we need to double check here.
const int S = fused_fp16_runner_->getSFromMaxSeqLen(parameters.sequence_length);
if (fused_fp16_runner_->isValid(S)) {
fused_runner = fused_fp16_runner_.get();
}
return fused_runner;
}
template <typename T>
Status PackedAttention<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* token_offset = context->Input<Tensor>(3);
const Tensor* cumulative_sequence_length = context->Input<Tensor>(4);
const Tensor* relative_position_bias = context->Input<Tensor>(5);
PackedAttentionParameters parameters;
ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(),
weights->Shape(),
bias->Shape(),
token_offset->Shape(),
cumulative_sequence_length->Shape(),
relative_position_bias,
parameters));
TensorShapeVector output_shape{parameters.token_count, parameters.v_hidden_size};
Tensor* output = context->Output(0, output_shape);
MHARunner* fused_runner = TryGettingFusedRunner(parameters);
typedef typename ToCudaType<T>::MappedType CudaT;
CudaT one = ToCudaType<T>::FromFloat(1.0f);
CudaT zero = ToCudaType<T>::FromFloat(0.0f);
IAllocatorUniquePtr<T> gemm_buffer;
int m = parameters.token_count;
int n = parameters.hidden_size + parameters.hidden_size + parameters.v_hidden_size;
int k = parameters.input_hidden_size;
gemm_buffer = GetScratchBuffer<T>(static_cast<size_t>(m) * n, context->GetComputeStream());
auto& device_prop = GetDeviceProp();
cublasHandle_t cublas = GetCublasHandle(context);
// Gemm, note that CUDA assumes col-major, so result(N, M) = 1 * weights x input + 1 x bias
// The bias part is not included here since we fuse bias, transpose and output 3 matrice into one cuda kernel.
CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(
cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &one,
reinterpret_cast<const CudaT*>(weights->Data<T>()), n,
reinterpret_cast<const CudaT*>(input->Data<T>()), k,
&zero, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop));
constexpr size_t element_size = sizeof(T);
size_t workSpaceSize = GetAttentionWorkspaceSize(element_size,
parameters.batch_size,
parameters.num_heads,
parameters.head_size,
parameters.v_head_size,
parameters.sequence_length,
fused_runner);
auto work_space = GetScratchBuffer<void>(workSpaceSize, context->GetComputeStream());
typedef typename ToCudaType<T>::MappedType CudaT;
PackedAttentionData<CudaT> data;
data.gemm_buffer = reinterpret_cast<CudaT*>(gemm_buffer.get());
data.bias = reinterpret_cast<const CudaT*>(bias->Data<T>());
data.relative_position_bias = (nullptr == relative_position_bias) ? nullptr : reinterpret_cast<const CudaT*>(relative_position_bias->Data<T>());
data.workspace = reinterpret_cast<CudaT*>(work_space.get());
data.token_offset = token_offset->Data<int32_t>();
data.cumulative_sequence_length = cumulative_sequence_length->Data<int32_t>();
data.output = reinterpret_cast<CudaT*>(output->MutableData<T>());
data.fused_runner = reinterpret_cast<void*>(fused_runner);
return QkvToContext<CudaT>(device_prop, cublas, Stream(context), parameters, data);
}
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

47
onnxruntime/contrib_ops/cuda/bert/packed_attention.h Normal file
View file

@ -0,0 +1,47 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <memory>
#include <vector>
#include "core/providers/cuda/cuda_kernel.h"
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h"
#include "contrib_ops/cpu/bert/attention_common.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
using namespace onnxruntime::cuda;
template <typename T>
class PackedAttention final : public CudaKernel {
public:
PackedAttention(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* context) const override;
private:
Status CheckInputs(const TensorShape& input_shape,
const TensorShape& weights_shape,
const TensorShape& bias_shape,
const TensorShape& packing_token_offset_shape,
const TensorShape& cu_seq_len_shape,
const Tensor* relative_position_bias,
PackedAttentionParameters& parameters) const;
MHARunner* TryGettingFusedRunner(const PackedAttentionParameters& parameters) const;
private:
int32_t num_heads_; // number of attention heads
std::vector<int64_t> qkv_hidden_sizes_; // Q, K, V hidden sizes parsed from the qkv_hidden_sizes attribute.
float scale_; // the scale to be used for softmax
bool disable_fused_runner_;
bool enable_trt_flash_attention_;
mutable std::unique_ptr<MHARunner> fused_fp16_runner_;
};
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

610
onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu Normal file
View file

@ -0,0 +1,610 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <cassert>
#include <cuda_fp16.h>
#include <cub/cub.cuh>
#include "core/providers/cuda/cu_inc/common.cuh"
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/shared_inc/fpgeneric.h"
#include "contrib_ops/cuda/bert/packed_attention_impl.h"
#include "contrib_ops/cuda/bert/attention_softmax.h"
#include "contrib_ops/cuda/bert/transformer_common.h"
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h"
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h"
#include "contrib_ops/cuda/bert/bert_padding.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
#include "contrib_ops/cuda/bert/rotary_embedding_util.h"
using namespace onnxruntime::cuda;
using namespace cub;
#define CHECK_CUDA(expr) CUDA_RETURN_IF_ERROR(expr)
namespace onnxruntime {
namespace contrib {
namespace cuda {
constexpr size_t kCUDAMemoryAlignment = 256;
constexpr int32_t kMAX_THREADS_PER_BLOCK = 256;
size_t GetAttentionScratchSize(
size_t element_size,
size_t batch_size,
size_t num_heads,
size_t sequence_length) {
const size_t bytes = element_size * batch_size * num_heads * sequence_length * sequence_length;
return ((bytes + kCUDAMemoryAlignment - 1) / kCUDAMemoryAlignment) * kCUDAMemoryAlignment;
}
size_t GetAttentionWorkspaceSize(
size_t element_size,
size_t batch_size,
size_t num_heads,
size_t qk_head_size,
size_t v_head_size,
size_t sequence_length,
void* fused_runner) {
// Note that q, k and v might need alignment for fused attention kernels.
const size_t qkv_bytes = element_size * batch_size * num_heads * sequence_length * (qk_head_size + qk_head_size + v_head_size);
if (fused_runner != nullptr) {
return qkv_bytes;
}
return qkv_bytes + 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length);
}
template <typename T, AttentionQkvFormat format>
__global__ void AddBiasTransposeQKVPacked(const T* input,
const T* biases,
int32_t N,
int32_t H_QK,
int32_t H_V,
T* q,
T* k,
T* v,
const int32_t* token_offset,
int32_t token_count);
// Grid: (S, B)
// Block: 256
// For unfused PackedAttention
// Input: Tx3xNxH
// Output: 3xBxNxSxH
// Where:
// T is token_count
// B is batch_size
// S is sequence_length
// N is num_heads
// H is head_size
template <typename T>
__global__ void AddBiasTransposeQKVPacked(
const T* input,
const T* biases,
int32_t N,
int32_t H_QK,
int32_t H_V,
T* q,
T* k,
T* v,
const int32_t* token_offset,
int32_t token_count) {
int s = blockIdx.x;
int b = blockIdx.y;
int S = gridDim.x;
const int packing_token_idx = b * S + s;
const int padding_token_idx = token_offset[packing_token_idx];
b = padding_token_idx / S;
s = padding_token_idx - b * S;
input += packing_token_idx * N * (H_QK + H_QK + H_V);
int k_offset = N * H_QK;
int v_offset = N * H_QK + N * H_QK;
q += (b * N * S + s) * H_QK;
k += (b * N * S + s) * H_QK;
v += (b * N * S + s) * H_V;
if (packing_token_idx < token_count) {
for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) {
int h = i % H_QK;
int n = i / H_QK;
q[n * S * H_QK + h] = input[i] + biases[i];
k[n * S * H_QK + h] = input[i + k_offset] + biases[i + k_offset];
}
for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) {
int h = i % H_V;
int n = i / H_V;
v[n * S * H_V + h] = input[i + v_offset] + biases[i + v_offset];
}
} else {
for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) {
int h = i % H_QK;
int n = i / H_QK;
q[n * S * H_QK + h] = biases[i];
k[n * S * H_QK + h] = biases[i + k_offset];
}
for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) {
int h = i % H_V;
int n = i / H_V;
v[n * S * H_V + h] = biases[i + v_offset];
}
}
}
// Grid: (S, B)
// Block: 256
// For memory efficient fMHA from CUTLASS. For future use, doesn't support fMHA from CUTLASS yet.
// Input: Tx3xNxH
// Output: 3xBxNxSxH
// T is token_count
// B is batch_size
// S is sequence_length
// N is num_heads
// H is head_size
template <typename T>
__global__ void AddBiasTransposeQKVPackedCutlass(
const T* input,
const T* biases,
int32_t N,
int32_t H_QK,
int32_t H_V,
T* q,
T* k,
T* v,
const int32_t* token_offset,
int32_t token_count) {
int s = blockIdx.x;
int b = blockIdx.y;
int S = gridDim.x;
const int packing_token_idx = b * S + s;
const int padding_token_idx = token_offset[packing_token_idx];
b = padding_token_idx / S;
s = padding_token_idx - b % S;
input += packing_token_idx * N * (H_QK + H_QK + H_V);
int k_offset = N * H_QK;
int v_offset = N * H_QK + N * H_QK;
q += (b * S * N + s * N) * H_QK;
k += (b * S * N + s * N) * H_QK;
v += (b * S * N + s * N) * H_V;
if (packing_token_idx < token_count) {
for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) {
q[i] = input[i] + biases[i];
k[i] = input[i + k_offset] + biases[i + k_offset];
}
for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) {
v[i] = input[i + v_offset] + biases[i + v_offset];
}
} else {
for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) {
q[i] = biases[i];
k[i] = biases[i + k_offset];
}
for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) {
v[i] = biases[i + v_offset];
}
}
}
// Grid: (S, B)
// Block: 256
// For fMHA from TRT
// Input: Tx3xNxH
// Output: TxNx3xH
// T is token_count
// B is batch_size
// S is sequence_length
// N is num_heads
// H is head_size
template <typename T>
__global__ void AddBiasTransposeQKVPackedTRT(
const T* input,
const T* biases,
int32_t N,
int32_t H,
T* output) {
int token_idx = blockIdx.x;
int Hx3 = H * 3;
int NxH = N * H;
int NxHx2 = N * H + N * H;
int offset = token_idx * N * Hx3;
input += offset;
output += offset;
for (int i = threadIdx.x; i < N * H; i += blockDim.x) {
int n = i / H;
int h = i % H;
output[n * Hx3 + h] = input[i] + biases[i];
output[n * Hx3 + H + h] = input[i + NxH] + biases[i + NxH];
output[n * Hx3 + H + H + h] = input[i + NxHx2] + biases[i + NxHx2];
}
}
template <typename T>
void InvokeAddBiasTranspose(
const T* input, const T* biases, T* output,
const int batch_size, const int sequence_length,
const int num_heads, const int qk_head_size, const int v_head_size,
AttentionQkvFormat format, const int32_t* token_offset, int32_t token_count,
cudaStream_t stream) {
if (format == AttentionQkvFormat::Q_K_V_BNSH) {
const dim3 grid(sequence_length, batch_size);
AddBiasTransposeQKVPacked<T><<<grid, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
input,
biases,
num_heads,
qk_head_size,
v_head_size,
output,
output + batch_size * sequence_length * num_heads * qk_head_size,
output + 2 * batch_size * sequence_length * num_heads * qk_head_size,
token_offset,
token_count);
} else if (format == AttentionQkvFormat::Q_K_V_BSNH) { // TODO: add memory efficient support
const dim3 grid(sequence_length, batch_size);
AddBiasTransposeQKVPackedCutlass<T><<<grid, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
input,
biases,
num_heads,
qk_head_size,
v_head_size,
output,
output + batch_size * sequence_length * num_heads * qk_head_size,
output + 2 * batch_size * sequence_length * num_heads * qk_head_size,
token_offset,
token_count);
} else {
ORT_ENFORCE(format == AttentionQkvFormat::QKV_BSN3H);
const dim3 grid(token_count);
AddBiasTransposeQKVPackedTRT<T><<<grid, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
input,
biases,
num_heads,
qk_head_size,
output);
}
}
template <typename T>
struct T4;
template <>
struct T4<float> {
using Type = float4;
};
template <>
struct T4<half> {
using Type = Half4;
};
template <typename T>
struct T2;
template <>
struct T2<float> {
using Type = float2;
};
template <>
struct T2<half> {
using Type = half2;
};
template <typename T>
void LaunchAddBiasTranspose(
const T* input, const T* biases, T* output,
const int batch_size, const int sequence_length,
const int num_heads, const int qk_head_size, const int v_head_size,
AttentionQkvFormat format, const int32_t* token_offset, int32_t token_count,
cudaStream_t stream) {
if (0 == (qk_head_size & 3) && 0 == (v_head_size & 3)) {
using T4Type = typename T4<T>::Type;
const int H = qk_head_size / 4;
const int H_v = v_head_size / 4;
const T4Type* input2 = reinterpret_cast<const T4Type*>(input);
const T4Type* biases2 = reinterpret_cast<const T4Type*>(biases);
T4Type* output2 = reinterpret_cast<T4Type*>(output);
InvokeAddBiasTranspose<T4Type>(
input2, biases2, output2,
batch_size, sequence_length,
num_heads, H, H_v,
format, token_offset, token_count, stream);
} else if (0 == (qk_head_size & 1) && 0 == (v_head_size & 1)) {
using T2Type = typename T2<T>::Type;
const int H = qk_head_size / 2;
const int H_v = v_head_size / 2;
const T2Type* input2 = reinterpret_cast<const T2Type*>(input);
const T2Type* biases2 = reinterpret_cast<const T2Type*>(biases);
T2Type* output2 = reinterpret_cast<T2Type*>(output);
InvokeAddBiasTranspose<T2Type>(
input2, biases2, output2,
batch_size, sequence_length,
num_heads, H, H_v,
format, token_offset, token_count, stream);
} else {
InvokeAddBiasTranspose<T>(
input, biases, output,
batch_size, sequence_length,
num_heads, qk_head_size, v_head_size,
format, token_offset, token_count, stream);
}
}
// Input: BxNxSxH
// Output: TxNxH
// where:
// T is token_count
// B is batch_size
// S is sequence_length
// N is num_heads
// H is head_size
// Grid: T
// Block: 256
template <typename T>
__global__ void __launch_bounds__(kMAX_THREADS_PER_BLOCK)
TransposeRemovePadding(T* target, const T* source, const int* token_offset,
const int B, const int S, const int N, const int H) {
int token_idx = blockIdx.x;
int source_idx = token_offset[token_idx];
int b = source_idx / S;
int s = source_idx - b * S;
target += token_idx * N * H;
source += b * N * S * H + s * H;
for (int i = threadIdx.x; i < N * H; i += blockDim.x) {
int n = i / H;
int h = i - n * H;
target[i] = source[n * S * H + h];
}
}
template <typename T>
Status LaunchTransposeRemovePadding(
T* output, const T* input,
const int* token_offset, const int token_count,
const int batch_size, const int seq_len, const int number_heads, const int head_size,
cudaStream_t stream);
// input: [batch_size, number_heads, seq_len, head_size]
// output: [token_count, number_heads * head_size]
template <>
Status LaunchTransposeRemovePadding(
half* output, const half* input,
const int* token_offset, const int token_count,
const int batch_size, const int seq_len, const int number_heads, const int head_size,
cudaStream_t stream) {
// Make sure memory is aligned to 128 bit
ORT_ENFORCE(!(reinterpret_cast<size_t>(input) & 0xF) && !(reinterpret_cast<size_t>(output) & 0xF), "alignment");
if (head_size % 8 == 0) {
const int4* input2 = reinterpret_cast<const int4*>(input);
int4* output2 = reinterpret_cast<int4*>(output);
TransposeRemovePadding<int4><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 8);
} else if (head_size % 4 == 0) {
const int64_t* input2 = reinterpret_cast<const int64_t*>(input);
int64_t* output2 = reinterpret_cast<int64_t*>(output);
TransposeRemovePadding<int64_t><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 4);
} else if (head_size % 2 == 0) {
const int32_t* input2 = reinterpret_cast<const int32_t*>(input);
int32_t* output2 = reinterpret_cast<int32_t*>(output);
TransposeRemovePadding<int32_t><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 2);
} else {
const int16_t* input2 = reinterpret_cast<const int16_t*>(input);
int16_t* output2 = reinterpret_cast<int16_t*>(output);
TransposeRemovePadding<int16_t><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size);
}
return CUDA_CALL(cudaGetLastError());
}
// input: [batch_size, number_heads, seq_len, head_size]
// output: [token_count, number_heads * head_size]
template <>
Status LaunchTransposeRemovePadding(
float* output, const float* input,
const int* token_offset, const int token_count,
const int batch_size, const int seq_len, const int number_heads, const int head_size,
cudaStream_t stream) {
ORT_ENFORCE(!(reinterpret_cast<size_t>(input) & 0xF) && !(reinterpret_cast<size_t>(output) & 0xF), "alignment");
if (head_size % 4 == 0) {
const int4* input2 = reinterpret_cast<const int4*>(input);
int4* output2 = reinterpret_cast<int4*>(output);
TransposeRemovePadding<int4><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 4);
} else if (head_size % 2 == 0) {
const int64_t* input2 = reinterpret_cast<const int64_t*>(input);
int64_t* output2 = reinterpret_cast<int64_t*>(output);
TransposeRemovePadding<int64_t><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 2);
} else {
const int32_t* input2 = reinterpret_cast<const int32_t*>(input);
int32_t* output2 = reinterpret_cast<int32_t*>(output);
TransposeRemovePadding<int32_t><<<token_count, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
output2, input2, token_offset, batch_size, seq_len, number_heads, head_size);
}
return CUDA_CALL(cudaGetLastError());
}
template <typename T>
Status FusedScaledDotProductAttention(
const cudaDeviceProp& device_prop,
cudaStream_t stream,
PackedAttentionParameters& parameters,
PackedAttentionData<T>& data) {
const int batch_size = parameters.batch_size;
const int sequence_length = parameters.sequence_length;
const int num_heads = parameters.num_heads;
const int qk_head_size = parameters.head_size;
const int v_head_size = parameters.v_head_size;
void* fused_runner = data.fused_runner;
assert(nullptr != fused_runner);
LaunchAddBiasTranspose(data.gemm_buffer, data.bias, data.workspace,
batch_size, sequence_length,
num_heads, qk_head_size, v_head_size,
AttentionQkvFormat::QKV_BSN3H, data.token_offset,
parameters.token_count, stream);
FusedMHARunnerFP16v2* fused_fp16_runner = reinterpret_cast<FusedMHARunnerFP16v2*>(fused_runner);
const int S = fused_fp16_runner->getSFromMaxSeqLen(sequence_length);
fused_fp16_runner->setup(S, batch_size);
fused_fp16_runner->run(data.workspace, data.cumulative_sequence_length, data.output, stream);
return Status::OK();
}
template <typename T>
Status UnfusedScaledDotProductAttention(
const cudaDeviceProp& device_prop,
cublasHandle_t& cublas,
cudaStream_t stream,
PackedAttentionParameters& parameters,
PackedAttentionData<T>& data) {
constexpr size_t element_size = sizeof(T);
const int batch_size = parameters.batch_size;
const int sequence_length = parameters.sequence_length;
const int num_heads = parameters.num_heads;
const int qk_head_size = parameters.head_size;
const int v_head_size = parameters.v_head_size;
const int batches = batch_size * num_heads;
const int size_per_batch_q = sequence_length * qk_head_size;
const int size_per_batch_k = sequence_length * qk_head_size;
const int size_per_batch_v = sequence_length * v_head_size;
const size_t elements_q = static_cast<size_t>(batches) * static_cast<size_t>(size_per_batch_q);
const size_t elements_k = static_cast<size_t>(batches) * static_cast<size_t>(size_per_batch_k);
const size_t elements_v = static_cast<size_t>(batches) * static_cast<size_t>(size_per_batch_v);
// Q, K and V pointers when fused attention is not used
T* qkv = data.workspace;
T* q = qkv;
T* k = q + elements_q;
T* v = k + elements_k;
LaunchAddBiasTranspose(data.gemm_buffer, data.bias, data.workspace,
batch_size, sequence_length,
num_heads, qk_head_size, v_head_size,
AttentionQkvFormat::Q_K_V_BNSH, data.token_offset,
parameters.token_count, stream);
T* scaled_qk = qkv + elements_q + elements_k + elements_v;
// Q, K and V are ready now
DUMP_TENSOR_INIT();
DUMP_TENSOR_D("gemm_buffer", data.gemm_buffer, parameters.token_count, (num_heads * (qk_head_size * 2 + v_head_size)));
DUMP_TENSOR_D("data.workspace", data.workspace, 3 * batch_size, num_heads, sequence_length, qk_head_size);
// Compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scaled_qk: BxNxSxT
// Q: BxNxSxH, K: BxNxSxH, Q*K': BxNxSxS
float one = 1.0f;
float zero = 0.f;
float scale = parameters.scale == 0.0f ? 1.f / sqrt(static_cast<float>(qk_head_size))
: parameters.scale;
cublasSetStream(cublas, stream);
CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(
cublas, CUBLAS_OP_T, CUBLAS_OP_N,
sequence_length, sequence_length, qk_head_size,
&scale,
k, qk_head_size, sequence_length * qk_head_size,
q, qk_head_size, sequence_length * qk_head_size,
&zero,
scaled_qk, sequence_length, sequence_length * sequence_length,
batches, device_prop));
DUMP_TENSOR_D("QK", scaled_qk, batch_size * num_heads, sequence_length, sequence_length);
const size_t bytes = GetAttentionScratchSize(element_size, batch_size, num_heads,
sequence_length);
T* attention_score = scaled_qk + (bytes / element_size);
// Apply softmax and store result R to attention_score: BxNxSxS
ORT_RETURN_IF_ERROR(ComputeSoftmaxWithCumSeqLength<T>(
scaled_qk,
data.relative_position_bias,
parameters.broadcast_res_pos_bias,
data.cumulative_sequence_length,
batch_size,
sequence_length,
num_heads,
attention_score, stream));
DUMP_TENSOR_D("Softmax", attention_score, batch_size * num_heads, sequence_length, sequence_length);
// compute R*V (as V*R), and store in temp_output (space used by Q): BxNxSxH_v
T* temp_output = qkv;
CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(
cublas, CUBLAS_OP_N, CUBLAS_OP_N,
v_head_size, sequence_length, sequence_length,
&one, v, v_head_size, sequence_length * v_head_size,
attention_score, sequence_length, sequence_length * sequence_length,
&zero, temp_output, v_head_size, sequence_length * v_head_size, batches, device_prop));
// Temp_output is BxNxSxH_v, transpose and remove padding to output token_countxNxH_v
Status result = LaunchTransposeRemovePadding(
data.output, temp_output,
data.token_offset, parameters.token_count,
batch_size, sequence_length, num_heads, v_head_size,
stream);
DUMP_TENSOR("unfused output", data.output, parameters.token_count, num_heads, v_head_size);
return result;
}
template <typename T>
Status QkvToContext(
const cudaDeviceProp& device_prop,
cublasHandle_t& cublas,
cudaStream_t stream,
PackedAttentionParameters& parameters,
PackedAttentionData<T>& data) {
void* fused_runner = data.fused_runner;
if (nullptr != fused_runner) {
return FusedScaledDotProductAttention<T>(device_prop, stream, parameters, data);
} else {
return UnfusedScaledDotProductAttention<T>(device_prop, cublas, stream, parameters, data);
}
}
template Status QkvToContext<float>(
const cudaDeviceProp& device_prop,
cublasHandle_t& cublas,
cudaStream_t stream,
PackedAttentionParameters& parameters,
PackedAttentionData<float>& data);
template Status QkvToContext<half>(
const cudaDeviceProp& device_prop,
cublasHandle_t& cublas,
cudaStream_t stream,
PackedAttentionParameters& parameters,
PackedAttentionData<half>& data);
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

53
onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h Normal file
View file

@ -0,0 +1,53 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/cuda/shared_inc/cuda_utils.h"
#include <cuda_fp16.h>
#include <cublas_v2.h>
#include "contrib_ops/cpu/bert/attention_common.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
size_t GetAttentionScratchSize(
size_t element_size,
size_t batch_size,
size_t num_heads,
size_t sequence_length);
size_t GetAttentionWorkspaceSize(
size_t element_size,
size_t batch_size,
size_t num_heads,
size_t qk_head_size,
size_t v_head_size,
size_t sequence_length,
void* fused_runner);
template <typename T>
struct PackedAttentionData {
T* gemm_buffer;
const T* bias;
const T* relative_position_bias;
const int32_t* token_offset;
const int32_t* cumulative_sequence_length;
T* workspace;
T* output;
void* fused_runner;
};
template <typename T>
Status QkvToContext(
const cudaDeviceProp& device_prop,
cublasHandle_t& cublas,
cudaStream_t stream,
contrib::PackedAttentionParameters& parameters,
PackedAttentionData<T>& data);
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

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

@ -63,6 +63,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, Affine);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Attention);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Attention);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, PackedAttention);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, PackedAttention);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BeamSearch);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, Crop);
@ -201,6 +203,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, Affine)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Attention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Attention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, PackedAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, PackedAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BeamSearch)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, Crop)>,

127
onnxruntime/core/graph/contrib_ops/bert_defs.cc
View file

@ -315,6 +315,133 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
AttentionTypeAndShapeInference(ctx, past_input_index);
}));
constexpr const char* PackingAttention_ver1_doc = R"DOC(
This is the packed version of Attention.
Sequences in one batch usually don't have same length and they are padded to have same length,
e.g., below is a batch with 3 sequences and tokens* are padded.
Sequence_0: 0, 1*, 2*, 3*
Sequence_1: 4, 5, 6*, 7*
Sequence_2: 8, 9, 10, 11
PackedAttention is designed to takes in packed input, i.e., only the real tokens without padding.
An input as above will be packed into 3 tensors like below:
- input ([h0, h4, h5, h8, h9, h10, h11])
- token_offset: 0, 4, 5, 8, 9, 10, 11, 1*, 2*, 3*, 6*, 7*
- cumulated_token_count: 0, 1, 1+2, 1+2+4
Input tensors contains the hidden embedding of real tokens.
Token_offset records the offset of token in the unpacked input.
cumulated_token_count records cumulated length of each sequnces length.
The operator only supports BERT like model with padding on right now.
)DOC";
// Shape inference for PackedAttention. Here are the shapes of inputs and output:
// Input 'input': (token_count, input_hidden_size)
// Input 'weights': (input_hidden_size, hidden_size + hidden_size + v_hidden_size)
// Input 'bias': (hidden_size + hidden_size + v_hidden_size)
// Input 'token_offset': (batch_size, sequence_length)
// Input 'cumulative_sequence_length': (batch_size + 1)
// Input 'relative_position_bias': (batch_size, num_heads, sequence_length, sequence_length)
// Output 'output': (token_count, v_hidden_size)
void PackedAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
// Type inference
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0);
// Shape inference
if (hasInputShape(ctx, 0) && hasInputShape(ctx, 2)) {
auto& input_shape = getInputShape(ctx, 0);
auto& input_dims = input_shape.dim();
int input_dim_size = input_dims.size();
if (input_dim_size != 2) {
fail_shape_inference("Inputs 0 shall be 2 dimensions");
}
auto& bias_shape = getInputShape(ctx, 2);
auto& bias_dims = bias_shape.dim();
if (bias_dims.size() != 1) {
fail_shape_inference("Invalid bias shape");
}
int64_t v_hidden_size = -1;
std::vector<int64_t> qkv_hidden_sizes;
getRepeatedAttribute(ctx, "qkv_hidden_sizes", qkv_hidden_sizes);
if (qkv_hidden_sizes.size() != 0) {
if (qkv_hidden_sizes.size() != 3) {
fail_shape_inference("qkv_hidden_sizes should have 3 elements")
}
v_hidden_size = qkv_hidden_sizes[2];
} else {
v_hidden_size = bias_shape.dim(0).dim_value() / 3;
}
ONNX_NAMESPACE::TensorShapeProto output_shape;
for (auto& dim : input_dims) {
*output_shape.add_dim() = dim;
}
output_shape.mutable_dim(input_dim_size - 1)->set_dim_value(v_hidden_size);
updateOutputShape(ctx, 0, output_shape);
}
}
ONNX_MS_OPERATOR_SET_SCHEMA(
PackedAttention, 1,
OpSchema()
.SetDoc(PackingAttention_ver1_doc)
.Attr("num_heads", "Number of attention heads", AttributeProto::INT)
.Attr("qkv_hidden_sizes",
"Hidden dimension of Q, K, V: hidden_size, hidden_size and v_hidden_size",
AttributeProto::INTS,
OPTIONAL_VALUE)
.Attr("scale",
"Custom scale will be used if specified. Default value is 1/sqrt(head_size)",
AttributeProto::FLOAT,
OPTIONAL_VALUE)
.Input(0,
"input",
"Input tensor with shape (token_count, input_hidden_size)",
"T")
.Input(1,
"weights",
"Merged Q/K/V weights with shape (input_hidden_size, hidden_size + hidden_size + v_hidden_size)",
"T")
.Input(2,
"bias",
"Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) for input projection",
"T")
.Input(3,
"token_offset",
"In packing mode, it specifies the offset of each token(batch_size, sequence_length).",
"M")
.Input(4,
"cumulative_sequence_length",
"A tensor with shape (batch_size + 1). It specifies the cumulative sequence length.",
"M")
.Input(5,
"relative_position_bias",
"A tensor with shape (batch_size, num_heads, sequence_length, sequence_length)"
"or (1, num_heads, sequence_length, sequence_length)."
"It specifies the additional bias to QxK'",
"T",
OpSchema::Optional)
.Output(0,
"output",
"2D output tensor with shape (token_count, v_hidden_size)",
"T")
.TypeConstraint("T",
{"tensor(float)", "tensor(float16)"},
"Constrain input and output types to float tensors.")
.TypeConstraint("M",
{"tensor(int32)"},
"Constrain mask index to integer types")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
PackedAttentionTypeAndShapeInference(ctx);
}));
constexpr const char* DecoderMaskedMultiheadAttention_ver1_doc = R"DOC(
Uni-directional attention that supports input sequence length of 1.

2
onnxruntime/core/graph/contrib_ops/ms_opset.h
View file

@ -81,6 +81,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MurmurHash3);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramRepeatBlock);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Pad);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PackedAttention);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RelativePositionBias);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRelativePositionBias);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RemovePadding);
@ -172,6 +173,7 @@ class OpSet_Microsoft_ver1 {
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MurmurHash3)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramRepeatBlock)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Pad)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PackedAttention)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QAttention)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QEmbedLayerNormalization)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RelativePositionBias)>());

502
onnxruntime/test/contrib_ops/packed_attention_op_test.cc Normal file
View file

@ -0,0 +1,502 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/platform/env_var_utils.h"
#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"
#include "test/util/include/scoped_env_vars.h"
#include "contrib_ops/cpu/bert/attention_common.h"
#include "test/contrib_ops/attention_op_test_helper.h"
namespace onnxruntime {
using contrib::AttentionMaskType;
namespace test {
static void RunPackedAttentionTest(
const std::vector<float>& input_data, // input: [token_count, 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<int32_t>& token_offset, // token_offset: [batch_size, sequence_length]
const std::vector<int32_t>& cumulative_sequence_length, // cum_seq_len: [batch_size + 1]
const std::vector<float>& output_data, // output: [token_count, hidden_size]
int batch_size,
int sequence_length,
int hidden_size,
int number_of_heads,
int token_count,
bool use_float16,
bool use_scale,
std::vector<int32_t> qkv_sizes,
const std::vector<float>& relative_position_bias_data) {
int min_cuda_architecture = use_float16 ? 530 : 0;
bool enable_cuda = HasCudaEnvironment(min_cuda_architecture);
if (enable_cuda) {
OpTester tester("PackedAttention", 1, onnxruntime::kMSDomain);
tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
int32_t qkv_hidden_size_sum;
int32_t v_hidden_size;
int32_t head_size_k;
if (qkv_sizes.size() != 0) {
qkv_hidden_size_sum = qkv_sizes[0] + qkv_sizes[1] + qkv_sizes[2];
std::vector<int64_t> sizes_attribute{qkv_sizes[0], qkv_sizes[1], qkv_sizes[2]};
tester.AddAttribute<std::vector<int64_t>>("qkv_hidden_sizes", sizes_attribute);
v_hidden_size = qkv_sizes[2];
head_size_k = qkv_sizes[1] / number_of_heads;
} else {
qkv_hidden_size_sum = 3 * hidden_size;
v_hidden_size = hidden_size;
head_size_k = hidden_size / number_of_heads;
}
if (use_scale) {
tester.AddAttribute<float>("scale", static_cast<float>(1.f / sqrt(head_size_k)));
}
std::vector<int64_t> input_dims = {token_count, hidden_size};
std::vector<int64_t> weights_dims = {hidden_size, qkv_hidden_size_sum};
std::vector<int64_t> bias_dims = {qkv_hidden_size_sum};
std::vector<int64_t> token_offset_dims = {batch_size, sequence_length};
std::vector<int64_t> cum_seq_len_dims = {batch_size + 1};
std::vector<int64_t> relative_position_bias_data_dims = {batch_size, number_of_heads, sequence_length, sequence_length};
std::vector<int64_t> output_dims = {token_count, v_hidden_size};
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<int32_t>("token_offset", token_offset_dims, token_offset);
tester.AddInput<int32_t>("cumulative_sequence_length", cum_seq_len_dims, cumulative_sequence_length);
if (relative_position_bias_data.size() > 0) {
tester.AddInput<MLFloat16>("relative_position_bias", relative_position_bias_data_dims, ToFloat16(relative_position_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<int32_t>("token_offset", token_offset_dims, token_offset);
tester.AddInput<int32_t>("cumulative_sequence_length", cum_seq_len_dims, cumulative_sequence_length);
if (relative_position_bias_data.size() > 0) {
tester.AddInput<float>("relative_position_bias", relative_position_bias_data_dims, relative_position_bias_data);
}
tester.AddOutput<float>("output", output_dims, output_data);
}
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
}
static void RunPackedAttentionTest(
const std::vector<float>& input_data, // input: [token_count, 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<int32_t>& token_offset, // token_offset: [batch_size, sequence_length]
const std::vector<int32_t>& cumulative_sequence_length, // cum_seq_len: [batch_size + 1]
const std::vector<float>& output_data, // output: [token_count, hidden_size]
int batch_size,
int sequence_length,
int hidden_size,
int number_of_heads,
int token_count,
std::vector<int32_t> qkv_sizes = {},
const std::vector<float>& relative_position_bias_data = {}) {
#define InvokePackedAttentionTest(use_float16, use_scale) \
RunPackedAttentionTest( \
input_data, \
weights_data, \
bias_data, \
token_offset, \
cumulative_sequence_length, \
output_data, \
batch_size, \
sequence_length, \
hidden_size, \
number_of_heads, \
token_count, \
use_float16, \
use_scale, \
qkv_sizes, \
relative_position_bias_data);
InvokePackedAttentionTest(true, true);
InvokePackedAttentionTest(true, false);
InvokePackedAttentionTest(false, true);
InvokePackedAttentionTest(false, false);
}
TEST(PackedAttentionTest, NoPack) {
int batch_size = 1;
int sequence_length = 2;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> weight_data = {
0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
std::vector<float> bias_data = {
-0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
std::vector<int32_t> token_offset{0, 1};
std::vector<int32_t> cum_seq_len{0, 2};
std::vector<float> output_data = {
3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f};
RunPackedAttentionTest(
input_data,
weight_data,
bias_data,
token_offset,
cum_seq_len,
output_data,
batch_size,
sequence_length,
hidden_size,
number_of_heads,
batch_size * sequence_length);
}
TEST(PackedAttentionTest, NoPackWithRelativePositionBias) {
int batch_size = 2;
int sequence_length = 2;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f,
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> weight_data = {
0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
std::vector<float> bias_data = {
-0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f,
0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
std::vector<int32_t> token_offset{0, 1, 2, 3};
std::vector<int32_t> cum_seq_len{0, 2, 4};
std::vector<float> relative_position_bias = {
0.2f, -0.1f, 0.4f, 2.5f, 1.6f, -1.1f, 0.4f, -2.5f,
0.2f, -0.1f, 0.4f, 2.5f, 1.6f, -1.1f, 0.4f, -2.5f};
std::vector<float> output_data = {
4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f,
-1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f,
4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f,
-1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f};
RunPackedAttentionTest(
input_data,
weight_data,
bias_data,
token_offset,
cum_seq_len,
output_data,
batch_size,
sequence_length,
hidden_size,
number_of_heads,
batch_size * sequence_length,
{},
relative_position_bias);
}
TEST(PackedAttentionTest, PackedWithRelativePositionBias) {
int batch_size = 2;
int sequence_length = 4;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f, // b0:s0
0.5f, 0.2f, 0.3f, -0.6f, // b0:s1
0.8f, -0.5f, 0.0f, 1.f, // b1:s0
0.5f, 0.2f, 0.3f, -0.6f // b1:s1
};
std::vector<float> weight_data = {
0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
std::vector<float> bias_data = {
-0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f,
0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
std::vector<int32_t> token_offset{0, 1, 4, 5, 2, 3, 6, 7};
std::vector<int32_t> cum_seq_len{0, 2, 4};
std::vector<float> relative_position_bias = {
0.2f, -0.1f, 0.f, 0.f, 0.4f, 2.5f, 0.f, 0.f,
1.6f, -1.1f, 0.f, 0.f, 0.4f, -2.5f, 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.2f, -0.1f, 0.f, 0.f, 0.4f, 2.5f, 0.f, 0.f,
1.6f, -1.1f, 0.f, 0.f, 0.4f, -2.5f, 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};
std::vector<float> output_data = {
4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f,
-1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f,
4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f,
-1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f};
RunPackedAttentionTest(
input_data,
weight_data,
bias_data,
token_offset,
cum_seq_len,
output_data,
batch_size,
sequence_length,
hidden_size,
number_of_heads,
4,
{},
relative_position_bias);
}
TEST(PackedAttentionTest, PackedBatch) {
int batch_size = 2;
int sequence_length = 4;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f, // b0:s0
0.5f, 0.2f, 0.3f, -0.6f, // b0:s1
0.8f, -0.5f, 0.0f, 1.f, // b1:s0
0.5f, 0.2f, 0.3f, -0.6f // b1:s1
};
std::vector<float> weight_data = {
0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
std::vector<float> bias_data = {
-0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
std::vector<int32_t> token_offset{0, 1, 4, 5, 2, 3, 6, 7};
std::vector<int32_t> cum_seq_len{0, 2, 4};
std::vector<float> output_data = {
3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f,
3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f};
RunPackedAttentionTest(
input_data,
weight_data,
bias_data,
token_offset,
cum_seq_len,
output_data,
batch_size,
sequence_length,
hidden_size,
number_of_heads,
4);
}
TEST(PackedAttentionTest, PackedBatchWithQKV) {
int batch_size = 2;
int sequence_length = 4;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f,
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<int32_t> qkv_sizes = {6, 6, 4};
std::vector<float> weight_data = {
0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f,
0.3f, 0.2f, 4.0f, 2.2f, 2.4f, 3.3f, 2.1f, 4.2f, 0.5f, 0.1f, 0.4f, 1.6f,
0.4f, 0.8f, 0.9f, 0.1f};
std::vector<float> bias_data = {
-0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f,
0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f,
0.5f, 0.7f, 0.2f, 1.2f};
std::vector<int32_t> token_offset{0, 1, 4, 5, 2, 3, 6, 7};
std::vector<int32_t> cum_seq_len{0, 2, 4};
std::vector<float> output_data = {
3.1967618465423584f, 0.51903456449508667f, 0.63051539659500122f, 2.9394614696502686f,
0.65332180261611938f, 1.000949501991272f, 0.74175024032592773f, 2.8231701850891113f,
3.1967618465423584f, 0.51903456449508667f, 0.63051539659500122f, 2.9394614696502686f,
0.65332180261611938f, 1.000949501991272f, 0.74175024032592773f, 2.8231701850891113f};
RunPackedAttentionTest(
input_data,
weight_data,
bias_data,
token_offset,
cum_seq_len,
output_data,
batch_size,
sequence_length,
hidden_size,
number_of_heads,
4,
qkv_sizes);
}
static void RunModelWithRandomInput(
int64_t batch_size,
int64_t sequence_length,
std::string& onnx_model,
bool is_float16) {
// ORT enables TF32 in GEMM for A100. TF32 will cause precsion loss and fail this test.
// Do not run this test unless TF32 is disabled explicitly.
if (HasCudaEnvironment(800) && ParseEnvironmentVariableWithDefault<int>("NVIDIA_TF32_OVERRIDE", 1) != 0) {
GTEST_SKIP() << "Skipping RunModelWithRandomInput in A100 since TF32 is enabled";
return;
}
RandomValueGenerator random{234};
constexpr int hidden_size = 768;
constexpr int num_heads = 12;
int token_count = 0;
std::vector<int32_t> cum_seq_len(batch_size + 1);
cum_seq_len[0] = 0;
int original_offset = 0;
int token_offset_idx = 0;
std::vector<int32_t> token_offset(batch_size * sequence_length);
for (int b = 0; b < batch_size; b++) {
int actual_seq_len = (sequence_length / (b + 1));
token_count += actual_seq_len;
cum_seq_len[b + 1] = token_count;
original_offset = b * sequence_length;
for (int s = 0; s < actual_seq_len; s++) {
token_offset[token_offset_idx++] = original_offset++;
}
}
for (int b = 0; b < batch_size; b++) {
int actual_seq_len = (sequence_length / (b + 1));
original_offset = b * sequence_length + actual_seq_len;
for (int s = actual_seq_len; s < sequence_length; s++) {
token_offset[token_offset_idx++] = original_offset++;
}
}
assert(token_offset_idx == batch_size * sequence_length);
std::vector<int64_t> input_dims{token_count, hidden_size};
std::vector<float> input_data = random.Gaussian<float>(input_dims, 0.0f, 0.3f);
std::vector<int64_t> weight_dims{hidden_size, 3 * hidden_size};
std::vector<float> weight_data = random.Gaussian<float>(weight_dims, 0.0f, 0.3f);
std::vector<int64_t> bias_dims{3 * hidden_size};
std::vector<float> bias_data = random.Gaussian<float>(bias_dims, 0.0f, 0.1f);
std::vector<int64_t> token_offset_dims{batch_size, sequence_length};
std::vector<int64_t> cum_seq_len_dims{batch_size + 1};
float gpu_threshold = is_float16 ? 0.1f : 0.005f;
bool enable_cuda = HasCudaEnvironment(is_float16 ? 530 : 0);
if (enable_cuda) {
OpTester test("PackedAttention", 1, onnxruntime::kMSDomain);
test.AddAttribute<int64_t>("num_heads", num_heads);
if (is_float16) {
test.AddInput<MLFloat16>("input", input_dims, ToFloat16(input_data));
test.AddInput<MLFloat16>("weight", weight_dims, ToFloat16(weight_data));
test.AddInput<MLFloat16>("bias", bias_dims, ToFloat16(bias_data));
} else {
test.AddInput<float>("input", input_dims, input_data);
test.AddInput<float>("weight", weight_dims, weight_data);
test.AddInput<float>("bias", bias_dims, bias_data);
}
test.AddInput<int32_t>("token_offset", token_offset_dims, token_offset);
test.AddInput<int32_t>("cumulative_sequence_length", cum_seq_len_dims, cum_seq_len);
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
test.AddReferenceOutputs(onnx_model, gpu_threshold, DefaultCudaExecutionProvider());
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
}
TEST(PackedAttentionTest, test_on_random_data) {
std::string onnx_model = "testdata/packed_attention_fp32.onnx";
std::string onnx_model_fp16 = "testdata/packed_attention_fp16.onnx";
for (int batch_size : std::vector<int>({1, 2, 3, 4, 5, 6, 7, 8})) {
for (int sequence_length : std::vector<int>({32, 48, 64, 95, 128})) {
RunModelWithRandomInput(
batch_size,
sequence_length,
onnx_model,
false);
RunModelWithRandomInput(
batch_size,
sequence_length,
onnx_model_fp16,
true);
}
}
}
TEST(PackedAttentionTest, test_on_random_data_large_seq) {
int batch_size = 2;
int sequence_length = 1152; // > 1024
std::string onnx_model = "testdata/packed_attention_fp32.onnx";
std::string onnx_model_fp16 = "testdata/packed_attention_fp16.onnx";
RunModelWithRandomInput(
batch_size,
sequence_length,
onnx_model,
false);
RunModelWithRandomInput(
batch_size,
sequence_length,
onnx_model_fp16,
true);
}
} // namespace test
} // namespace onnxruntime

7
onnxruntime/test/providers/provider_test_utils.cc
View file

@ -413,9 +413,9 @@ struct TensorCheck<BFloat16> {
if (abs_error <= abs_threshold) {
// if the absolute error is small enough, then no need to calculate realative error
EXPECT_NEAR(0, abs_error, abs_threshold) << "provider_type: "
<< provider_type;
<< provider_type;
} else {
//default for existing tests.
// default for existing tests.
const float rel_error = abs_error / max_value;
EXPECT_NEAR(0, rel_error, threshold) << "provider_type: "
<< provider_type;
@ -1406,7 +1406,7 @@ void OpTester::ExecuteModelForEps(
}
};
void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_error) {
void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_error, std::unique_ptr<IExecutionProvider> ep) {
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
@ -1418,6 +1418,7 @@ void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_erro
Status status;
InferenceSession subgraph_session_object{so, GetEnvironment()};
status = subgraph_session_object.RegisterExecutionProvider(std::move(ep));
ASSERT_TRUE((status = subgraph_session_object.Load(model_path)).IsOK()) << status;
ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status;

2
onnxruntime/test/providers/provider_test_utils.h
View file

@ -687,7 +687,7 @@ class OpTester {
}
// Generate the reference outputs with the model file
void AddReferenceOutputs(const std::string& model_path, float abs_error = 0.0f);
void AddReferenceOutputs(const std::string& model_path, float abs_error = 0.0f, std::unique_ptr<IExecutionProvider> ep = nullptr);
void AddCustomOpRegistry(std::shared_ptr<CustomRegistry> registry) {
custom_schema_registries_.push_back(registry->GetOpschemaRegistry());

BIN
onnxruntime/test/testdata/packed_attention_fp16.onnx vendored Normal file
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/packed_attention_fp32.onnx vendored Normal file
View file

Binary file not shown.