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

[t5 optimization] kernel changes to t5 (#14928)

Browse source 
### Description
<!-- Describe your changes. -->

1. support optional bias in Attention op (used in T5 encoder)
2. support broadcasting rel_pos_bias in attention_softmax.h
3. add scale in
MHA op's attributes
4. support past_key/past_value and present_key/present_value in MHA
5. UT and parity tests are added
6. fix an issue: https://github.com/microsoft/onnxruntime/issues/14920

note: the fusions will be in another PR since mt5 needs to be tested and
an issue from github will be investigated.

Future works:
1. support shared buffer for past/present
2. enable trt kernels when possible and investigate (trt/cutlass)kernels
with rel_pos_bias)
3. support KV/QKV packing with past/present

### 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. -->

---------

Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2023-03-13 14:29:16 -07:00 committed by GitHub
parent b34e570ad0
commit 538d64891a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 1345 additions and 152 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
docs/ContribOperators.md
View file

@ -144,14 +144,14 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Whether every token can only attend to previous tokens. Default value is 0.</dd>
</dl>
#### Inputs (3 - 7)
#### Inputs (2 - 7)
<dl>
<dt><tt>input</tt> : T</dt>
<dd>Input tensor with shape (batch_size, sequence_length, 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>
<dt><tt>bias</tt> (optional) : T</dt>
<dd>Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) for input projection</dd>
<dt><tt>mask_index</tt> (optional) : M</dt>
<dd>Attention mask with shape (batch_size, 1, max_sequence_length, max_sequence_length), (batch_size, total_sequence_length) or (batch_size, sequence_length, total_sequence_length), or index with shape (batch_size) or (2 * batch_size)</dd>
@ -2381,30 +2381,40 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>The value to be filled in the attention mask. Default value is -10000.0f</dd>
<dt><tt>num_heads</tt> : int (required)</dt>
<dd>Number of attention heads</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 (1 - 6)
#### Inputs (1 - 8)
<dl>
<dt><tt>query</tt> : T</dt>
<dd>Query with shape (batch_size, sequence_length, hidden_size), or packed QKV with shape (batch_size, kv_sequence_length, num_heads, 3, head_size)</dd>
<dt><tt>key</tt> (optional) : T</dt>
<dd>Key with shape (batch_size, kv_sequence_length, hidden_size), or packed KV with shape (batch_size, kv_sequence_length, num_heads, 2, head_size)</dd>
<dd>Key with shape (batch_size, kv_sequence_length, hidden_size), or packed KV with shape (batch_size, kv_sequence_length, num_heads, 2, head_size), or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size)</dd>
<dt><tt>value</tt> (optional) : T</dt>
<dd>Value with shape (batch_size, kv_sequence_length, v_hidden_size)</dd>
<dd>Value with shape (batch_size, kv_sequence_length, v_hidden_size), or past_value with shape (batch_size, num_heads, kv_sequence_length, head_size)</dd>
<dt><tt>bias</tt> (optional) : T</dt>
<dd>Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection</dd>
<dt><tt>key_padding_mask</tt> (optional) : M</dt>
<dd>Key padding mask with shape (batch_size) or (batch_size, kv_sequence_length)</dd>
<dt><tt>relative_position_bias</tt> (optional) : T</dt>
<dd>relative position bias: addition to QxK' with shape (batch_size, num_heads, sequence_length, total_sequence_length) or (1, num_heads, sequence_length, total_sequence_length)</dd>
<dt><tt>past_key</tt> (optional) : T</dt>
<dd>past state for self attention key with shape (batch_size, num_heads, past_sequence_length, head_size)</dd>
<dt><tt>past_value</tt> (optional) : T</dt>
<dd>past state for self attention value with shape (batch_size, num_heads, past_sequence_length, head_size)</dd>
</dl>
#### Outputs
#### Outputs (1 - 3)
<dl>
<dt><tt>output</tt> : T</dt>
<dd>3D output tensor with shape (batch_size, sequence_length, v_hidden_size)</dd>
<dt><tt>present_key</tt> (optional) : T</dt>
<dd>present state for cross attention key with shape (batch_size, num_heads, kv_sequence_length, head_size)or present state for self attention key with shape (batch_size, num_heads, total_sequence_length, head_size)</dd>
<dt><tt>present_value</tt> (optional) : T</dt>
<dd>present state for cross attention value with shape (batch_size, num_heads, kv_sequence_length, head_size)or present state for self attention value with shape (batch_size, num_heads, total_sequence_length, head_size)</dd>
</dl>
#### Type Constraints

2
docs/OperatorKernels.md
View file

@ -813,7 +813,7 @@ Do not modify directly.*
|Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|Irfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|LongformerAttention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask:**T**<br> *in* global_weight:**T**<br> *in* global_bias:**T**<br> *in* global:**G**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|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> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|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)|
|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)|

10
onnxruntime/contrib_ops/cpu/bert/attention_base.cc
View file

@ -193,6 +193,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
}
}
bool broadcast_res_pos_bias = false;
if (relative_position_bias != nullptr) {
const auto& relative_position_bias_dims = relative_position_bias->Shape().GetDims();
@ -202,11 +203,14 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
relative_position_bias_dims.size());
}
if (relative_position_bias_dims[0] != batch_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, got ",
"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 ",
@ -255,6 +259,8 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
output_parameters->mask_filter_value = mask_filter_value_;
output_parameters->scale = scale_;
output_parameters->mask_type = mask_type;
output_parameters->broadcast_res_pos_bias = broadcast_res_pos_bias;
output_parameters->pass_past_in_kv = false;
}
return Status::OK();

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

@ -52,6 +52,8 @@ struct AttentionParameters {
bool is_unidirectional;
bool past_present_share_buffer;
bool do_rotary;
bool broadcast_res_pos_bias;
bool pass_past_in_kv;
float mask_filter_value;
float scale;
AttentionMaskType mask_type;

126
onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h
View file

@ -18,16 +18,21 @@ Status CheckInputs(const T* query,
const T* bias,
const T* key_padding_mask,
const T* relative_position_bias,
const T* past_key,
const T* past_value,
void* parameters,
int num_heads,
float mask_filter_value,
float scale,
int max_threads_per_block) {
// key_padding_mask (K/V) : (B) or (B, L) or None
// relative_position_bias : (B, 1, S, L)
// past_key : (B, N, S*, H)
// past_value : (B, N, S*, H)
// When no packing for q/k/v:
// query (Q) : (B, S, D)
// key (K) : (B, L, D)
// value (V) : (B, L, D_v)
// key (K) : (B, L, D) or (B, N, S*, H)
// value (V) : (B, L, D_v) or (B, N, S*, H)
// bias (Q/K/V) : (D + D + D_v)
// When packed kv is used:
// query (Q) : (B, S, D)
@ -40,6 +45,7 @@ Status CheckInputs(const T* query,
// value (V) : None
// bias (Q/K/V) : None
const auto& query_dims = query->Shape().GetDims();
if (query_dims.size() != 3 && query_dims.size() != 5) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 3 or 5 dimensions, got ",
@ -52,6 +58,63 @@ Status CheckInputs(const T* query,
int head_size = static_cast<int>(hidden_size) / num_heads;
int kv_sequence_length = sequence_length;
int past_sequence_length = 0;
if (past_key != nullptr && past_value != nullptr) {
const auto& past_key_dims = past_key->Shape().GetDims();
const auto& past_value_dims = past_value->Shape().GetDims();
if (past_key_dims.size() != 4) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' is expected to have 4 dimensions, got ",
past_key_dims.size());
}
if (past_value_dims.size() != 4) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_value' is expected to have 4 dimensions, got ",
past_value_dims.size());
}
if (past_key_dims[0] != batch_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' dimension 0 should be batch_size, got ",
past_key_dims[0]);
}
if (past_value_dims[0] != batch_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_value' dimension 0 should be batch_size, got ",
past_value_dims[0]);
}
if (past_key_dims[1] != num_heads) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' dimension 1 should be same as number of heads, got ",
past_key_dims[1]);
}
if (past_value_dims[1] != num_heads) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_value' dimension 1 should be same as number of heads, got ",
past_value_dims[1]);
}
if (past_key_dims[2] != past_value_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' and 'past_value' shall have same dim 2 (past_sequence_length)");
}
if (past_key_dims[3] != head_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' dimension 3 should be same as head_size, got ",
past_key_dims[3]);
}
if (past_value_dims[3] != head_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_value' dimension 3 should be same as head_size, got ",
past_value_dims[3]);
}
past_sequence_length = static_cast<int>(past_key_dims[2]);
} else if (past_key != nullptr || past_value != nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' and 'past_value' shall be both present or both absent");
}
if (key != nullptr) {
if (query_dims.size() != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 3 dimensions when key is given, got ",
@ -59,8 +122,8 @@ Status CheckInputs(const T* query,
}
const auto& key_dims = key->Shape().GetDims();
if (key_dims.size() != 3 && key_dims.size() != 5) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'key' is expected to have 3 or 5 dimensions, got ",
if (key_dims.size() != 3 && key_dims.size() != 4 && key_dims.size() != 5) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'key' is expected to have 3, 4, or 5 dimensions, got ",
key_dims.size());
}
if (query_dims[0] != key_dims[0]) {
@ -73,8 +136,9 @@ Status CheckInputs(const T* query,
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'query' and 'key' shall have same dim 2 (hidden_size)");
}
} else // if (key_dims.size() == 5)
{
kv_sequence_length = static_cast<int>(key_dims[1]);
} else if (key_dims.size() == 5) {
if (static_cast<int>(key_dims[2]) != num_heads || static_cast<int>(key_dims[3]) != 2 || static_cast<int>(key_dims[4]) != head_size) {
return ORT_MAKE_STATUS(
ONNXRUNTIME, INVALID_ARGUMENT,
@ -83,9 +147,17 @@ Status CheckInputs(const T* query,
if (value != nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Expect 'value' be none when 'key' has packed kv format.");
}
}
kv_sequence_length = static_cast<int>(key_dims[1]);
kv_sequence_length = static_cast<int>(key_dims[1]);
} else { // key_dims.size() == 4 (cross-attention with past_key)
if (static_cast<int>(key_dims[1]) != num_heads || static_cast<int>(key_dims[3]) != head_size) {
return ORT_MAKE_STATUS(
ONNXRUNTIME, INVALID_ARGUMENT,
"Expect 'key' shape (batch_size, num_heads, kv_sequence_length, head_size) for past_key");
}
kv_sequence_length = static_cast<int>(key_dims[2]);
}
} else { // packed QKV
if (query_dims.size() != 5) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 5 dimensions when key is empty, got ",
@ -128,11 +200,13 @@ Status CheckInputs(const T* query,
}
}
// NOTE: In Cross-Attention, we pass the past key and value to 'key' and 'value' instead of 'past_key' and 'past_value'.
bool pass_past_in_kv = false;
int v_hidden_size = hidden_size;
if (value != nullptr) {
const auto& value_dims = value->Shape().GetDims();
if (value_dims.size() != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value' is expected to have 3 dimensions, got ",
if (value_dims.size() != 3 && value_dims.size() != 4) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value' is expected to have 3 or 4 dimensions, got ",
value_dims.size());
}
@ -141,13 +215,24 @@ Status CheckInputs(const T* query,
"Input 'query' and 'value' shall have same dim 0 (batch_size)");
}
if (static_cast<int64_t>(kv_sequence_length) != value_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'key' and 'value' shall have same same dim 1 (kv_sequence_length)");
if (value_dims.size() == 3) {
if (static_cast<int64_t>(kv_sequence_length) != value_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'key' and 'value' shall have the same dim 1 (kv_sequence_length)");
}
v_hidden_size = static_cast<int>(value_dims[2]);
} else { // value_dims.size() == 4
if (static_cast<int64_t>(kv_sequence_length) != value_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'past_key' and 'past_value' shall have the same dim 2 (kv_sequence_length)");
}
v_hidden_size = static_cast<int>(value_dims[1]) * static_cast<int>(value_dims[3]);
pass_past_in_kv = true;
}
v_hidden_size = static_cast<int>(value_dims[2]);
}
int total_sequence_length = past_sequence_length + kv_sequence_length;
bool broadcast_res_pos_bias = false;
if (relative_position_bias != nullptr) {
const auto& relative_position_bias_dims = relative_position_bias->Shape().GetDims();
@ -161,6 +246,9 @@ Status CheckInputs(const T* query,
"Input 'relative_position_bias' dimension 0 should be 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 ",
@ -171,7 +259,7 @@ Status CheckInputs(const T* query,
"Input 'relative_position_bias' dimension 2 should be same as sequence_length, got ",
relative_position_bias_dims[2]);
}
if (relative_position_bias_dims[3] != kv_sequence_length) {
if (relative_position_bias_dims[3] != total_sequence_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'relative_position_bias' dimension 3 should be same as total_sequence_length, got ",
relative_position_bias_dims[3]);
@ -182,9 +270,9 @@ Status CheckInputs(const T* query,
AttentionParameters* output_parameters = reinterpret_cast<AttentionParameters*>(parameters);
output_parameters->batch_size = batch_size;
output_parameters->sequence_length = sequence_length;
output_parameters->past_sequence_length = 0;
output_parameters->past_sequence_length = past_sequence_length;
output_parameters->kv_sequence_length = kv_sequence_length;
output_parameters->total_sequence_length = kv_sequence_length;
output_parameters->total_sequence_length = total_sequence_length;
output_parameters->max_sequence_length = 0;
output_parameters->input_hidden_size = 0;
output_parameters->hidden_size = hidden_size;
@ -196,7 +284,9 @@ Status CheckInputs(const T* query,
output_parameters->past_present_share_buffer = false;
output_parameters->mask_filter_value = mask_filter_value;
output_parameters->mask_type = mask_type;
output_parameters->scale = 0.0f;
output_parameters->scale = scale;
output_parameters->broadcast_res_pos_bias = broadcast_res_pos_bias;
output_parameters->pass_past_in_kv = pass_past_in_kv;
}
if (max_threads_per_block > 0 && num_heads > max_threads_per_block) {

11
onnxruntime/contrib_ops/cuda/bert/attention.cc
View file

@ -67,9 +67,12 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
auto& device_prop = GetDeviceProp();
AttentionParameters parameters;
// Use the second dimension from weight for bias to get q_hidden_size when bias is nullptr
std::vector<int64_t> bias_dims{weights->Shape().GetDims()[1]};
const TensorShape bias_shape{bias_dims};
ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(),
weights->Shape(),
bias->Shape(),
bias != nullptr ? bias->Shape() : bias_shape,
mask_index,
past,
relative_position_bias,
@ -200,18 +203,22 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
typedef typename ToCudaType<T>::MappedType CudaT;
AttentionData<CudaT> data;
data.gemm_buffer = reinterpret_cast<CudaT*>(gemm_buffer.get());
data.bias = reinterpret_cast<const CudaT*>(bias->Data<T>());
data.bias = nullptr == bias ? nullptr : reinterpret_cast<const CudaT*>(bias->Data<T>());
data.query = nullptr;
data.key = nullptr;
data.value = nullptr;
data.mask_index = (nullptr == mask_index) ? nullptr : mask_index->Data<int>();
data.mask_index_dims = (nullptr == mask_index) ? gsl::span<const int64_t>() : mask_index->Shape().GetDims();
data.past = (nullptr == past) ? nullptr : reinterpret_cast<const CudaT*>(past->Data<T>());
data.past_key = nullptr;
data.past_value = nullptr;
data.relative_position_bias = (nullptr == relative_position_bias) ? nullptr : reinterpret_cast<const CudaT*>(relative_position_bias->Data<T>());
data.has_qkv_workspace = true;
data.workspace = reinterpret_cast<CudaT*>(work_space.get());
data.output = reinterpret_cast<CudaT*>(output->MutableData<T>());
data.present = (nullptr == present) ? nullptr : reinterpret_cast<CudaT*>(present->MutableData<T>());
data.present_key = nullptr;
data.present_value = nullptr;
data.fused_runner = reinterpret_cast<void*>(fused_runner);
data.fused_cross_attention_kernel = nullptr;
data.use_memory_efficient_attention = use_memory_efficient_attention;

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

@ -333,6 +333,56 @@ Status PrepareQkv(contrib::AttentionParameters& parameters,
data.gemm_buffer, data.bias, qkv, true, v_head_size, qkv_add_bias,
3, parameters.do_rotary, parameters.original_past_sequence_length);
}
}
// cross attention with past/present state
else if (data.past_key != nullptr || data.present_key != nullptr) {
// no bias for T5 cross attention
assert(data.bias == nullptr);
// cross attention with past state
if (data.past_key != nullptr && data.present_key == nullptr) {
assert(data.past_value != nullptr);
assert(data.query != nullptr);
assert(data.key == nullptr);
assert(data.value == nullptr);
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, sequence_length, batch_size, qk_head_size, num_heads,
max_threads_per_block, false, data.query, q));
}
// cross attention with present state or self attention with present state
else if (data.past_key == nullptr && data.present_key != nullptr) {
assert(data.past_value == nullptr);
assert(data.present_value != nullptr);
assert(data.query != nullptr);
assert(data.key != nullptr);
assert(data.value != nullptr);
// TODO: supporting packed qkv for self attention may benefit performance
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, sequence_length, batch_size, qk_head_size, num_heads,
max_threads_per_block, false, data.query, q));
// TODO: supporting packed kv for cross attention may benefit performance
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, kv_sequence_length, batch_size, qk_head_size, num_heads,
max_threads_per_block, false, data.key, data.present_key));
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, kv_sequence_length, batch_size, v_head_size, num_heads,
max_threads_per_block, false, data.value, data.present_value));
}
// self attention with past and present state
else {
assert(data.past_key != nullptr);
assert(data.past_value != nullptr);
assert(data.present_key != nullptr);
assert(data.present_value != nullptr);
assert(data.query != nullptr);
assert(data.key != nullptr);
assert(data.value != nullptr);
// TODO: supporting packed qkv for self attention may benefit performance
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, sequence_length, batch_size, qk_head_size, num_heads,
max_threads_per_block, false, data.query, q));
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, kv_sequence_length, batch_size, qk_head_size, num_heads,
max_threads_per_block, false, data.key, k));
ORT_RETURN_IF_ERROR(LaunchTransQkv(stream, 1, kv_sequence_length, batch_size, v_head_size, num_heads,
max_threads_per_block, false, data.value, v));
}
qkv_format = AttentionQkvFormat::Q_K_V_BNSH;
} else if (data.key == nullptr) { // gemm_buffer == nullptr and packed qkv
assert(data.bias == nullptr);
assert(qk_head_size == v_head_size);
@ -536,6 +586,28 @@ Status QkvToContext(
k = data.present;
v = data.present + batches * present_size_per_batch_k;
}
if (nullptr != data.past_key || nullptr != data.present_key) {
assert(qkv_format == AttentionQkvFormat::Q_K_V_BNSH);
if (nullptr != data.past_key && nullptr == data.present_key) {
k = const_cast<T*>(data.past_key);
v = const_cast<T*>(data.past_value);
} else if (nullptr == data.past_key && nullptr != data.present_key) {
k = data.present_key;
v = data.present_value;
} else {
ORT_RETURN_IF_ERROR(
LaunchConcatTensorToTensor(stream, parameters.total_sequence_length, sequence_length, batch_size, qk_head_size, num_heads,
max_threads_per_block, 1, data.past_key, k, data.present_key));
ORT_RETURN_IF_ERROR(
LaunchConcatTensorToTensor(stream, parameters.total_sequence_length, sequence_length, batch_size, v_head_size, num_heads,
max_threads_per_block, 1, data.past_value, v, data.present_value));
// Update pointers to present_k and present_v.
k = data.present_key;
v = data.present_value;
}
}
} else {
assert(qk_head_size == v_head_size);
assert(data.fused_cross_attention_kernel == nullptr);
@ -544,6 +616,11 @@ Status QkvToContext(
assert(!data.use_memory_efficient_attention);
assert(data.has_qkv_workspace);
if (nullptr != data.past_key || nullptr != data.present_key) {
// TODO: support this case.
ORT_THROW("buffer sharing for no bias case between past and present is not supported yet.");
}
if (data.present != data.past) {
// For easy testing. Production should better avoid this path.
int64_t kv_size = 2LL * (int64_t)batch_size * num_heads * parameters.max_sequence_length * qk_head_size;
@ -719,6 +796,8 @@ Status QkvToContext(
cublasSetStream(cublas, stream);
DUMP_TENSOR_D("q[BNSH]", q, batch_size, num_heads, sequence_length, qk_head_size);
DUMP_TENSOR_D("k[BNSH]", k, batch_size, num_heads, total_sequence_length, qk_head_size);
CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(
cublas, CUBLAS_OP_T, CUBLAS_OP_N,
total_sequence_length, sequence_length, qk_head_size,
@ -743,21 +822,22 @@ Status QkvToContext(
T* persistent_softmax_workspace = scratch1; // replace Q*K' in place with masked score for persistent softmax.
ORT_RETURN_IF_ERROR(
ComputeSoftmaxWithRawMask<T>(stream, total_sequence_length, sequence_length, batch_size, num_heads,
mask_index, nullptr, data.relative_position_bias, scratch1, scratch2,
parameters.is_unidirectional, scale, mask_dimension,
parameters.max_sequence_length, use_persistent_softmax,
persistent_softmax_workspace, mask_filter_value));
mask_index, nullptr, data.relative_position_bias, parameters.broadcast_res_pos_bias,
scratch1, scratch2, parameters.is_unidirectional, scale, mask_dimension,
parameters.max_sequence_length, use_persistent_softmax, persistent_softmax_workspace,
mask_filter_value));
} else if (nullptr != mask_index) { // 1d mask index
assert(mask_index_dims.size() == 1);
// mask_index has 1D shape: either (batch_size) or (2*batch_size). Only the later one has start postions.
const int* mask_start = (mask_index_dims[0] > batch_size) ? mask_index + batch_size : nullptr;
ORT_RETURN_IF_ERROR(ComputeSoftmaxWithMask1D<T>(
stream, total_sequence_length, sequence_length, batch_size, num_heads,
mask_index, mask_start, data.relative_position_bias, scratch1, scratch2, parameters.is_unidirectional));
mask_index, mask_start, data.relative_position_bias, parameters.broadcast_res_pos_bias,
scratch1, scratch2, parameters.is_unidirectional));
} else { // no mask
ORT_RETURN_IF_ERROR(
ComputeSoftmax<T>(stream, total_sequence_length, sequence_length, batch_size, num_heads, data.relative_position_bias,
scratch1, scratch2, parameters.is_unidirectional));
parameters.broadcast_res_pos_bias, scratch1, scratch2, parameters.is_unidirectional));
}
DUMP_TENSOR_D("Softmax", scratch2, batch_size * num_heads, sequence_length, total_sequence_length);
@ -908,13 +988,15 @@ Status DecoderQkvToContext(
if (has_key_padding_mask) {
constexpr int mask_dimension = 2;
constexpr int max_sequence_length = 0;
ORT_RETURN_IF_ERROR(ComputeSoftmaxWithRawMask<T>(stream, kv_sequence_length, sequence_length, batch_size, num_heads,
nullptr, key_padding_mask, add_before_softmax, scratch1, scratch2,
is_unidirectional, 1.0f, mask_dimension, max_sequence_length,
false, nullptr, mask_filter_value));
ORT_RETURN_IF_ERROR(ComputeSoftmaxWithRawMask<T>(stream, kv_sequence_length, sequence_length, batch_size,
num_heads, nullptr, key_padding_mask, add_before_softmax,
false/*broadcast rpb*/, scratch1, scratch2, is_unidirectional,
1.0f, mask_dimension, max_sequence_length, false, nullptr,
mask_filter_value));
} else {
ORT_RETURN_IF_ERROR(ComputeSoftmax<T>(stream, kv_sequence_length, sequence_length, batch_size, num_heads,
add_before_softmax, scratch1, scratch2, is_unidirectional));
add_before_softmax, false/*broadcast rpb*/, scratch1, scratch2,
is_unidirectional));
}
// compute P*V (as V*P), and store in scratch3: BxNxSxH

4
onnxruntime/contrib_ops/cuda/bert/attention_impl.h
View file

@ -57,6 +57,8 @@ struct AttentionData {
const int* mask_index;
gsl::span<const int64_t> mask_index_dims;
const T* past;
const T* past_key;
const T* past_value;
const T* relative_position_bias;
bool has_qkv_workspace;
@ -64,6 +66,8 @@ struct AttentionData {
T* output;
T* present;
T* present_key;
T* present_value;
void* fused_runner;
const void* fused_cross_attention_kernel;

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

@ -39,7 +39,8 @@ __device__ inline void Softmax(const int all_sequence_length,
const int sequence_length,
const int valid_end,
const int valid_start,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output) {
using BlockReduce = cub::BlockReduce<float, TPB>;
@ -50,17 +51,22 @@ __device__ inline void Softmax(const int all_sequence_length,
float thread_data_max(-CUDART_INF_F);
const bool no_add = (add_before_softmax == nullptr);
const bool no_rpb = (rel_pos_bias == nullptr);
// e^x is represented as infinity if x is large enough, like 100.f.
// Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough.
// a math transform as below is leveraged to get a stable softmax:
// e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max))
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
const int size_per_batch = gridDim.x * all_sequence_length;
for (int i = threadIdx.x; i < valid_end; i += TPB) {
if (i >= valid_start) {
const int index = offset + i;
float input_at_idx = no_add ? float(input[index]) : float(input[index] + add_before_softmax[index]);
float input_at_idx = no_rpb
? float(input[index])
: float(input[index] + (broadcast_rel_pos_bias
? rel_pos_bias[index % size_per_batch]
: rel_pos_bias[index]));
if (thread_data_max < input_at_idx) {
thread_data_max = input_at_idx;
}
@ -79,7 +85,7 @@ __device__ inline void Softmax(const int all_sequence_length,
for (int i = threadIdx.x; i < valid_end; i += TPB) {
if (i >= valid_start) {
const int index = offset + i;
float val = no_add ? input[index] : input[index] + add_before_softmax[index];
float val = no_rpb ? input[index] : input[index] + rel_pos_bias[index % size_per_batch];
thread_data_sum += expf(val - max_block);
}
}
@ -92,7 +98,7 @@ __device__ inline void Softmax(const int all_sequence_length,
for (int i = threadIdx.x; i < all_sequence_length; i += TPB) {
const int index = offset + i;
float input_at_idx = no_add ? float(input[index]) : float(input[index] + add_before_softmax[index]);
float input_at_idx = no_rpb ? float(input[index]) : float(input[index] + rel_pos_bias[index % size_per_batch]);
const float val = (i >= valid_start && i < valid_end) ? expf(input_at_idx - max_block) * sum_reverse_block : 0.f;
output[index] = T(val);
}
@ -103,7 +109,8 @@ __device__ inline void SoftmaxSmall(const int all_sequence_length,
const int sequence_length,
const int valid_end,
const int valid_start,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
bool is_unidirectional) {
@ -139,8 +146,13 @@ __device__ inline void SoftmaxSmall(const int all_sequence_length,
// 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_add = (add_before_softmax == nullptr);
float input_data = no_add ? float(input[index]) : float(input[index] + add_before_softmax[index]);
const bool no_rpb = (rel_pos_bias == nullptr);
const int size_per_batch = gridDim.x * all_sequence_length;
float input_data = no_rpb
? float(input[index])
: float(input[index] + (broadcast_rel_pos_bias
? rel_pos_bias[index % size_per_batch]
: rel_pos_bias[index]));
float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F);
const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end);
@ -174,7 +186,8 @@ __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* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
bool is_unidirectional) {
@ -205,6 +218,7 @@ __global__ void SoftmaxLargeKernel(const int all_sequence_length,
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
const int size_per_batch = gridDim.x * all_sequence_length;
float thread_data_max = -CUDART_INF_F;
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
@ -216,8 +230,12 @@ __global__ void SoftmaxLargeKernel(const int all_sequence_length,
// 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);
? (rel_pos_bias
? float(input[index] + (broadcast_rel_pos_bias
? rel_pos_bias[index % size_per_batch]
: rel_pos_bias[index]))
: float(input[index]))
: float(-CUDART_INF_F);
cached_data[seq_idx] = input_data;
thread_data_max = max(thread_data_max, input_data);
}
@ -254,7 +272,8 @@ __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* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
const bool is_unidirectional,
@ -272,16 +291,19 @@ __global__ void SoftmaxWithRawMaskLargeKernel(const int all_sequence_length,
__shared__ float max_block;
float max_thread_data = -CUDART_INF_F;
const int size_per_batch = gridDim.x * all_sequence_length;
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
int base_index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
for (int seq_idx = threadIdx.x; seq_idx < all_sequence_length; seq_idx += TPB) {
float thread_data = -CUDART_INF_F;
int index = base_index + seq_idx;
if (add_before_softmax == nullptr) {
if (rel_pos_bias == nullptr) {
thread_data = float(input[index]) * rsqrt_head_size;
} else {
thread_data = float(input[index] + add_before_softmax[index]) * rsqrt_head_size;
T rel_pos_bias_value = broadcast_rel_pos_bias ?
rel_pos_bias[index % size_per_batch] : rel_pos_bias[index];
thread_data = float(input[index] + rel_pos_bias_value) * rsqrt_head_size;
}
const int sequence_index = blockIdx.x % sequence_length;
@ -357,7 +379,8 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
const int sequence_length,
const int* attention_mask, // 2D, 3D or 4D attention mask
const bool* key_padding_mask,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
const bool is_unidirectional,
@ -374,6 +397,7 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
int index = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length + threadIdx.x;
const int size_per_batch = gridDim.x * all_sequence_length;
float thread_data = -CUDART_INF_F;
if (threadIdx.x < all_sequence_length) {
@ -409,8 +433,10 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
}
}
if (add_before_softmax != nullptr) {
thread_data += float(add_before_softmax[index]);
if (rel_pos_bias != nullptr) {
float rel_pos_bias_value = broadcast_rel_pos_bias ?
float(rel_pos_bias[index % size_per_batch]) : float(rel_pos_bias[index]);
thread_data += rel_pos_bias_value;
}
}
@ -446,62 +472,65 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
template <typename T, unsigned TPB>
__global__ void SoftmaxKernelSmall(const int all_sequence_length,
const int sequence_length,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
bool is_unidirectional) {
SoftmaxSmall<T, TPB>(all_sequence_length, sequence_length, all_sequence_length, 0,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
}
template <typename T, unsigned TPB>
__global__ void SoftmaxKernel(const int all_sequence_length,
const int sequence_length,
const T* add_before_softmax,
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,
add_before_softmax, input, output);
rel_pos_bias, broadcast_rel_pos_bias, input, output);
}
template <typename T>
Status ComputeSoftmax(cudaStream_t stream, const int all_sequence_length, const int sequence_length,
const int batch_size, const int num_heads,
const T* add_before_softmax, const T* input, T* output, bool is_unidirectional) {
const int batch_size, const int num_heads, const T* rel_pos_bias,
const bool broadcast_rel_pos_bias, T* input, T* output, bool is_unidirectional) {
const dim3 grid(sequence_length * num_heads, batch_size, 1);
if (all_sequence_length <= 32) {
const int blockSize = 32;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output, is_unidirectional);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 64) {
const int blockSize = 64;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output, is_unidirectional);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 128) {
const int blockSize = 128;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output, is_unidirectional);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 256) {
const int blockSize = 256;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output, is_unidirectional);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 512) {
const int blockSize = 512;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output, is_unidirectional);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 1024) {
const int blockSize = 1024;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output, is_unidirectional);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (!is_unidirectional) {
const int blockSize = 1024;
SoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output);
all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output);
} else {
const int blockSize = 256;
const int sh_bytes = sizeof(float) * all_sequence_length;
SoftmaxLargeKernel<T, blockSize><<<grid, blockSize, sh_bytes, stream>>>(
all_sequence_length, sequence_length, all_sequence_length, 0, add_before_softmax, input, output, true);
all_sequence_length, sequence_length, all_sequence_length, 0, rel_pos_bias, broadcast_rel_pos_bias,
input, output, true);
}
return CUDA_CALL(cudaGetLastError());
@ -512,7 +541,8 @@ __global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length,
const int sequence_length,
const int* mask_end,
const int* mask_start,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
bool is_unidirectional) {
@ -533,7 +563,7 @@ __global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length,
__syncthreads();
SoftmaxSmall<T, TPB>(all_sequence_length, sequence_length, end_position, start_position,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
}
template <typename T, unsigned TPB>
@ -541,7 +571,8 @@ __global__ void MaskedSoftmaxKernel(const int all_sequence_length,
const int sequence_length,
const int* mask_end,
const int* mask_start,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input, T* output) {
__shared__ int start_position;
__shared__ int end_position;
@ -560,7 +591,7 @@ __global__ void MaskedSoftmaxKernel(const int all_sequence_length,
__syncthreads();
Softmax<T, TPB>(all_sequence_length, sequence_length, end_position, start_position,
add_before_softmax, input, output);
rel_pos_bias, broadcast_rel_pos_bias, input, output);
}
template <typename T, unsigned TPB>
@ -568,7 +599,8 @@ __global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length,
const int sequence_length,
const int* attention_mask,
const bool* key_padding_mask,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
const bool is_unidirectional,
@ -579,7 +611,7 @@ __global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length,
const float mask_filter_value) {
SoftmaxWithRawMaskSmall<T, TPB>(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, output,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input, output,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
skip_softmax, mask_filter_value);
}
@ -592,7 +624,8 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream,
const int num_heads,
const int* mask_index,
const int* mask_start,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
const bool is_unidirectional) {
@ -602,37 +635,37 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream,
const int blockSize = 32;
MaskedSoftmaxKernelSmall<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 64) {
const int blockSize = 64;
MaskedSoftmaxKernelSmall<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 128) {
const int blockSize = 128;
MaskedSoftmaxKernelSmall<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 256) {
const int blockSize = 256;
MaskedSoftmaxKernelSmall<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 512) {
const int blockSize = 512;
MaskedSoftmaxKernelSmall<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (all_sequence_length <= 1024) {
const int blockSize = 1024;
MaskedSoftmaxKernelSmall<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output, is_unidirectional);
rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional);
} else if (!is_unidirectional) {
const int blockSize = 1024;
MaskedSoftmaxKernel<T, blockSize>
<<<grid, blockSize, 0, stream>>>(all_sequence_length, sequence_length, mask_index, mask_start,
add_before_softmax, input, output);
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.");
}
@ -648,7 +681,8 @@ Status ComputeSoftmaxWithRawMask(cudaStream_t stream,
const int num_heads,
const int* attention_mask,
const bool* key_padding_mask,
const T* add_before_softmax,
const T* rel_pos_bias,
const bool broadcast_rel_pos_bias,
const T* input,
T* output,
const bool is_unidirectional,
@ -665,51 +699,51 @@ Status ComputeSoftmaxWithRawMask(cudaStream_t stream,
const int blockSize = 32;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 64) {
const int blockSize = 64;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 128) {
const int blockSize = 128;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 256) {
const int blockSize = 256;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 512) {
const int blockSize = 512;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 1024) {
const int blockSize = 1024;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
} else {
const int blockSize = 256;
const int sh_bytes = sizeof(float) * all_sequence_length;
SoftmaxWithRawMaskLargeKernel<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,
attention_mask, key_padding_mask, rel_pos_bias, broadcast_rel_pos_bias, input,
out, is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
use_persistent_softmax, mask_filter_value);
}

25
onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
View file

@ -42,6 +42,8 @@ MultiHeadAttention<T>::MultiHeadAttention(const OpKernelInfo& info)
mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
scale_ = info.GetAttrOrDefault<float>("scale", 0.0f);
disable_fused_self_attention_ = sizeof(T) != 2 ||
ParseEnvironmentVariableWithDefault<bool>(attention::kDisableFusedSelfAttention, false);
@ -73,6 +75,8 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
const Tensor* bias = context->Input<Tensor>(3);
const Tensor* key_padding_mask = context->Input<Tensor>(4);
const Tensor* relative_position_bias = context->Input<Tensor>(5);
const Tensor* past_key = context->Input<Tensor>(6);
const Tensor* past_value = context->Input<Tensor>(7);
auto& device_prop = GetDeviceProp();
AttentionParameters parameters;
@ -82,9 +86,12 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
bias,
key_padding_mask,
relative_position_bias,
past_key,
past_value,
&parameters,
num_heads_,
mask_filter_value_,
scale_,
device_prop.maxThreadsPerBlock));
int sequence_length = parameters.sequence_length;
@ -95,6 +102,12 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
output_shape[2] = static_cast<int64_t>(parameters.v_hidden_size);
Tensor* output = context->Output(0, output_shape);
std::vector<int64_t> present_dims{
parameters.batch_size, parameters.num_heads, parameters.total_sequence_length, parameters.head_size};
TensorShape present_shape(present_dims);
Tensor* present_key = context->Output(1, present_shape);
Tensor* present_value = context->Output(2, present_shape);
MHARunner* fused_runner = nullptr;
const FusedMultiHeadCrossAttentionKernel* fused_cross_attention_kernel = nullptr;
@ -107,6 +120,7 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
bool use_fused_cross_attention = !disable_fused_cross_attention_ &&
nullptr == key_padding_mask &&
nullptr == relative_position_bias &&
(nullptr == past_key && nullptr == past_value && !parameters.pass_past_in_kv) &&
key != nullptr &&
(value != nullptr || bias == nullptr) && // TODO: new kernel for adding bias to packed KV
parameters.hidden_size == parameters.v_hidden_size &&
@ -128,6 +142,7 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
fused_cross_attention_kernel == nullptr &&
nullptr == relative_position_bias &&
(value != nullptr || key == nullptr) &&
(nullptr == past_key && nullptr == past_value && !parameters.pass_past_in_kv) &&
(nullptr == key_padding_mask || is_mask_1d_seq_len) &&
parameters.hidden_size == parameters.v_hidden_size &&
parameters.sequence_length == parameters.kv_sequence_length &&
@ -195,16 +210,22 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
data.gemm_buffer = nullptr;
data.bias = (nullptr == bias) ? nullptr : reinterpret_cast<const CudaT*>(bias->Data<T>());
data.query = reinterpret_cast<const CudaT*>(query->Data<T>());
data.key = (nullptr == key) ? nullptr : reinterpret_cast<const CudaT*>(key->Data<T>());
data.value = (nullptr == value) ? nullptr : reinterpret_cast<const CudaT*>(value->Data<T>());
data.key = (nullptr == key || parameters.pass_past_in_kv) ? nullptr : reinterpret_cast<const CudaT*>(key->Data<T>());
data.value = (nullptr == value || parameters.pass_past_in_kv) ? nullptr : reinterpret_cast<const CudaT*>(value->Data<T>());
data.mask_index = (nullptr == key_padding_mask) ? nullptr : key_padding_mask->Data<int>();
data.mask_index_dims = (nullptr == key_padding_mask) ? gsl::span<const int64_t>() : key_padding_mask->Shape().GetDims();
data.past = nullptr;
data.past_key = (parameters.pass_past_in_kv) ? reinterpret_cast<const CudaT*>(key->Data<T>())
: (nullptr == past_key) ? nullptr : reinterpret_cast<const CudaT*>(past_key->Data<T>());
data.past_value = (parameters.pass_past_in_kv) ? reinterpret_cast<const CudaT*>(value->Data<T>())
: (nullptr == past_value) ? nullptr : reinterpret_cast<const CudaT*>(past_value->Data<T>());
data.relative_position_bias = (nullptr == relative_position_bias) ? nullptr : reinterpret_cast<const CudaT*>(relative_position_bias->Data<T>());
data.has_qkv_workspace = !no_qkv_workspace;
data.workspace = reinterpret_cast<CudaT*>(work_space.get());
data.output = reinterpret_cast<CudaT*>(output->MutableData<T>());
data.present = nullptr;
data.present_key = (nullptr == present_key) ? nullptr : reinterpret_cast<CudaT*>(present_key->MutableData<T>());
data.present_value = (nullptr == present_value) ? nullptr : reinterpret_cast<CudaT*>(present_value->MutableData<T>());
data.fused_runner = reinterpret_cast<void*>(fused_runner);
data.fused_cross_attention_kernel = fused_cross_attention_kernel;
data.use_memory_efficient_attention = use_memory_efficient_attention;

1
onnxruntime/contrib_ops/cuda/bert/multihead_attention.h
View file

@ -24,6 +24,7 @@ class MultiHeadAttention final : public CudaKernel {
protected:
int num_heads_; // number of attention heads
float mask_filter_value_;
float scale_;
bool disable_fused_self_attention_;
bool enable_trt_flash_attention_;
bool disable_fused_cross_attention_;

4
onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc
View file

@ -200,11 +200,15 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
data.mask_index = (nullptr == mask_index) ? nullptr : mask_index->Data<int>();
data.mask_index_dims = (nullptr == mask_index) ? gsl::span<const int64_t>() : mask_index->Shape().GetDims();
data.past = (nullptr == past_tensor) ? nullptr : reinterpret_cast<const CudaT*>(past_tensor->Data<T>());
data.past_key = nullptr;
data.past_value = nullptr;
data.relative_position_bias = nullptr; // add_qk is not supported in quantized attention
data.has_qkv_workspace = true;
data.workspace = reinterpret_cast<CudaT*>(work_space.get());
data.output = reinterpret_cast<CudaT*>(output->MutableData<T>());
data.present = (nullptr == present) ? nullptr : reinterpret_cast<CudaT*>(present->MutableData<T>());
data.present_key = nullptr;
data.present_value = nullptr;
data.fused_runner = fused_runner;
data.fused_cross_attention_kernel = nullptr;
data.use_memory_efficient_attention = use_memory_efficient_attention;

2
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
View file

@ -997,8 +997,8 @@ Status UpdateDecoderFeeds(
for (int i = 0; i < num_present_tensors; ++i) {
next_inputs[t5_decoder_first_past_input_idx + i] =
last_outputs[t5_decoder_first_present_output_idx + i];
return Status::OK();
}
return Status::OK();
}
return PickT5PastState<T>(last_outputs, next_inputs, num_present_tensors, beam_indices, allocator,

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

@ -133,6 +133,11 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c
// Input 1 (key) has shape (batch_size, kv_sequence_length, hidden_size)
// Input 2 (value) has shape (batch_size, kv_sequence_length, v_hidden_size)
// Q, K and V without packing and past (cross attention):
// Input 0 (query) has shape (batch_size, sequence_length, hidden_size)
// Input 1 (key) has shape (batch_size, num_head, kv_sequence_length, head_size)
// Input 2 (value) has shape (batch_size, num_head, kv_sequence_length, head_size)
// Packed KV:
// Input 0 (query) has shape (batch_size, sequence_length, hidden_size)
// Input 1 (batch_size, kv_sequence_length, num_heads, 2, head_size)
@ -167,14 +172,14 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c
if (hasInputShape(ctx, 2)) {
auto& value_shape = getInputShape(ctx, 2);
auto& value_dims = value_shape.dim();
if (value_dims.size() != 3) {
fail_shape_inference("Inputs 2 (value) shall be 3 dimensions");
if (value_dims.size() != 3 && value_dims.size() != 4) {
fail_shape_inference("Inputs 2 (value) shall be 3 or 4 dimensions");
}
ONNX_NAMESPACE::TensorShapeProto output_shape;
*output_shape.add_dim() = query_dims[0];
*output_shape.add_dim() = query_dims[1];
*output_shape.add_dim() = value_dims[2];
*output_shape.add_dim() = value_dims.size() == 3 ? value_dims[2] : value_dims[1] * value_dims[3];
updateOutputShape(ctx, 0, output_shape);
return;
}
@ -261,7 +266,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
.Input(2,
"bias",
"Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) for input projection",
"T")
"T",
OpSchema::Optional)
.Input(3,
"mask_index",
"Attention mask with shape (batch_size, 1, max_sequence_length, max_sequence_length), "
@ -412,18 +418,23 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
.Attr("num_heads", "Number of attention heads", AttributeProto::INT)
.Attr("mask_filter_value", "The value to be filled in the attention mask. Default value is -10000.0f",
AttributeProto::FLOAT, OPTIONAL_VALUE)
.Attr("scale",
"Custom scale will be used if specified. Default value is 1/sqrt(head_size)",
AttributeProto::FLOAT,
OPTIONAL_VALUE)
.Input(0,
"query",
"Query with shape (batch_size, sequence_length, hidden_size), or packed QKV with shape (batch_size, kv_sequence_length, num_heads, 3, head_size)",
"T")
.Input(1,
"key",
"Key with shape (batch_size, kv_sequence_length, hidden_size), or packed KV with shape (batch_size, kv_sequence_length, num_heads, 2, head_size)",
"Key with shape (batch_size, kv_sequence_length, hidden_size), or packed KV with shape (batch_size, kv_sequence_length, num_heads, 2, head_size), "
"or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size)",
"T",
OpSchema::Optional)
.Input(2,
"value",
"Value with shape (batch_size, kv_sequence_length, v_hidden_size)",
"Value with shape (batch_size, kv_sequence_length, v_hidden_size), or past_value with shape (batch_size, num_heads, kv_sequence_length, head_size)",
"T",
OpSchema::Optional)
.Input(3,
@ -442,10 +453,32 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
" or (1, num_heads, sequence_length, total_sequence_length)",
"T",
OpSchema::Optional)
.Input(6,
"past_key",
"past state for self attention key with shape (batch_size, num_heads, past_sequence_length, head_size)",
"T",
OpSchema::Optional)
.Input(7,
"past_value",
"past state for self attention value with shape (batch_size, num_heads, past_sequence_length, head_size)",
"T",
OpSchema::Optional)
.Output(0,
"output",
"3D output tensor with shape (batch_size, sequence_length, v_hidden_size)",
"T")
.Output(1,
"present_key",
"present state for cross attention key with shape (batch_size, num_heads, kv_sequence_length, head_size)"
"or present state for self attention key with shape (batch_size, num_heads, total_sequence_length, head_size)",
"T",
OpSchema::Optional)
.Output(2,
"present_value",
"present state for cross attention value with shape (batch_size, num_heads, kv_sequence_length, head_size)"
"or present state for self attention value with shape (batch_size, num_heads, total_sequence_length, head_size)",
"T",
OpSchema::Optional)
.TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output to float tensors.")
.TypeConstraint("M", {"tensor(int32)"}, "Constrain mask to integer types")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {

13
onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc
View file

@ -122,9 +122,16 @@ void AttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int p
// When past and present share buffer, they shape is same: (2, batch_size, num_heads, max_sequence_length, head_size)
// Type inference
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 2, 0);
if (ctx.getNumOutputs() > 1) {
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 2, 1);
if (hasInputShape(ctx, 2)) {
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 2, 0);
if (ctx.getNumOutputs() > 1) {
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 2, 1);
}
} else {
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (ctx.getNumOutputs() > 1) {
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1);
}
}
// Shape inference

182
onnxruntime/test/contrib_ops/attention_op_test_helper.cc
View file

@ -2908,6 +2908,188 @@ void GetCrossAttentionData_HeadSize16(AttentionTestData& data) {
}
}
void GetCrossAttentionDataWithPast(AttentionTestData& data) {
data.hidden_size = 8;
data.v_hidden_size = 8;
data.num_heads = 2;
data.batch_size = 1;
data.sequence_length = 2;
data.kv_sequence_length = 3;
data.mask_type = AttentionMaskType::MASK_2D_KEY_PADDING;
data.key_padding_mask_data = {1, 1, 1};
data.skip_kernel_types = {
AttentionKernelType::AttentionKernel_TrtFlashAttention,
AttentionKernelType::AttentionKernel_TrtFusedCrossAttention,
AttentionKernelType::AttentionKernel_TrtFusedAttention,
AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention};
{
data.query_data = {
-0.10939738f, -0.11916742f, -0.23157823f, -0.12894472f,
-0.02661306f, 0.26251313f, 0.30725253f, -0.34759378f,
-0.11695808f, -0.13129434f, -0.17031054f, -0.14986445f,
-0.02826184f, 0.2797631f , 0.27337456f, -0.44312602f
};
}
// The past key and value data will be passed to the kernel as input 'key' and 'value'.
{
data.past_key_data = {
0.5967375f , 0.5966938f , 0.48602432f, 0.5341031f,
0.55797786f, 0.5663399f , 0.57087725f, 0.6240304f,
0.5352563f , 0.5648297f , 0.4972945f , 0.56637144f,
0.44123724f, 0.35872823f, 0.32176313f, 0.4490301f,
0.3643952f , 0.51968557f, 0.50137347f, 0.5743993f,
0.3905106f , 0.4741712f , 0.40881708f, 0.47243845f
};
}
{
data.past_value_data = {
0.40251260f, 0.55487730f, 0.49565578f, 0.42683450f,
0.44379145f, 0.58945787f, 0.54852820f, 0.43376005f,
0.44116694f, 0.44007313f, 0.40293324f, 0.53202707f,
0.35520583f, 0.47293650f, 0.45417705f, 0.33723440f,
0.50175804f, 0.37620395f, 0.24103148f, 0.50958070f,
0.56803876f, 0.37866923f, 0.32273075f, 0.44389135f
};
}
{
data.fp32_output_data = {
0.4291f, 0.5275f, 0.4818f, 0.4645f, 0.4770f, 0.4082f, 0.3372f, 0.4319f,
0.4291f, 0.5276f, 0.4818f, 0.4645f, 0.4768f, 0.4083f, 0.3377f, 0.4315f
};
}
{
data.fp16_output_data = data.fp32_output_data;
}
}
void GetSelfAttentionDataWithPast(AttentionTestData& data) {
data.hidden_size = 8;
data.v_hidden_size = 8;
data.num_heads = 2;
data.batch_size = 1;
data.sequence_length = 2;
data.kv_sequence_length = 3;
data.mask_type = AttentionMaskType::MASK_NONE;
data.skip_kernel_types = {
AttentionKernelType::AttentionKernel_TrtFlashAttention,
AttentionKernelType::AttentionKernel_TrtFusedCrossAttention,
AttentionKernelType::AttentionKernel_TrtFusedAttention,
AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention};
{
data.query_data = {
0.00403503f, 0.08716156f, -0.0358175f , -0.08171791f,
0.48912194f, -0.22679007f, -0.09093101f, -0.5939322f,
0.00878838f, 0.03355761f, -0.08080226f, -0.06677517f,
0.55038965f, -0.2720567f , -0.12977877f, -0.634123f
};
}
{
data.key_data = {
0.2808786f , 0.10041683f, 0.15880886f, 0.45283064f,
0.39884242f, 0.12596075f, 0.4198916f , -0.0651141f,
0.31678027f, 0.11010794f, 0.21594375f, 0.4975329f,
0.436772f , 0.20940652f, 0.44072092f, -0.05601776f
};
}
{
data.value_data = {
0.26421773f, -0.16541699f, -0.0599675f , 0.27200517f,
-0.1074627f , -0.4493224f , -0.03694462f, 0.17997989f,
0.27960598f, -0.16643806f, -0.07019104f, 0.29006317f,
-0.11640988f, -0.47876123f, -0.01979145f, 0.11468418f
};
}
{
data.rel_pos_bias_data = {
0.4781123f , 0.82420444f, 0.654424f , 0.3995186f , 0.5482078f,
0.55570245f, 0.4216576f , 0.46001542f, 0.67183703f, 0.41973996f,
0.28494194f, 0.60367906f, 0.3453173f , 0.44483483f, 0.6770777f,
0.5460559f , 0.31994605f, 0.5470492f , 0.5433419f , 0.60349935f
};
}
{
data.past_key_data = {
0.34734827f, 0.5592256f , 0.5333037f , 0.5122027f,
0.5940516f , 0.44744077f, 0.43128848f, 0.55360645f,
0.57874715f, 0.29512063f, 0.2780432f , 0.4693917f,
0.4450266f , 0.530704f , 0.3124955f , 0.4273598f,
0.44368753f, 0.5890438f , 0.5054336f , 0.46042535f,
0.5352153f , 0.5157861f , 0.39744973f, 0.5441864f
};
}
{
data.past_value_data = {
0.48998538f, 0.5493853f , 0.556647f , 0.7011929f,
0.543909f , 0.5630743f , 0.5087797f , 0.3901024f,
0.53116417f, 0.4086225f , 0.5320247f , 0.5145377f,
0.4086198f , 0.6913348f , 0.50045484f, 0.5338214f,
0.52980417f, 0.5243695f , 0.6046111f , 0.53555113f,
0.44936907f, 0.6010697f , 0.38031512f, 0.427301f
};
}
{
data.fp32_output_data = {
0.4358f, 0.2708f, 0.3201f, 0.4347f, 0.1886f, 0.0845f, 0.2479f, 0.3289f,
0.4157f, 0.2247f, 0.2826f, 0.4321f, 0.1874f, 0.1021f, 0.2427f, 0.3305f
};
}
{
data.fp16_output_data = data.fp32_output_data;
}
{
data.present_key_data = {
0.3473f, 0.5592f, 0.5333f, 0.5122f,
0.5941f, 0.4474f, 0.4313f, 0.5536f,
0.5787f, 0.2951f, 0.2780f, 0.4694f,
0.2809f, 0.1004f, 0.1588f, 0.4528f,
0.3168f, 0.1101f, 0.2159f, 0.4975f,
0.4450f, 0.5307f, 0.3125f, 0.4274f,
0.4437f, 0.5890f, 0.5054f, 0.4604f,
0.5352f, 0.5158f, 0.3974f, 0.5442f,
0.3988f, 0.1260f, 0.4199f, -0.0651f,
0.4368f, 0.2094f, 0.4407f, -0.0560f
};
}
{
data.present_value_data = {
0.4900f, 0.5494f, 0.5566f, 0.7012f,
0.5439f, 0.5631f, 0.5088f, 0.3901f,
0.5312f, 0.4086f, 0.5320f, 0.5145f,
0.2642f, -0.1654f, -0.0600f, 0.2720f,
0.2796f, -0.1664f, -0.0702f, 0.2901f,
0.4086f, 0.6913f, 0.5005f, 0.5338f,
0.5298f, 0.5244f, 0.6046f, 0.5356f,
0.4494f, 0.6011f, 0.3803f, 0.4273f,
-0.1075f, -0.4493f, -0.0369f, 0.1800f,
-0.1164f, -0.4788f, -0.0198f, 0.1147f
};
}
data.is_static_kv = false;
}
bool SkipAttentionKernel(AttentionTestData& data, AttentionKernelType kernel_type) {
return std::find(data.skip_kernel_types.begin(), data.skip_kernel_types.end(), kernel_type) != data.skip_kernel_types.end();
}

13
onnxruntime/test/contrib_ops/attention_op_test_helper.h
View file

@ -11,6 +11,7 @@ using contrib::AttentionMaskType;
namespace test {
struct AttentionTestData {
bool is_static_kv = true;
int hidden_size;
int v_hidden_size;
int num_heads;
@ -27,8 +28,17 @@ struct AttentionTestData {
std::vector<float> qkv_data;
std::vector<float> bias_data;
std::vector<float> rel_pos_bias_data;
std::vector<float> past_key_data;
std::vector<float> past_value_data;
std::vector<float> fp32_output_data;
std::vector<float> fp16_output_data;
std::vector<float> present_key_data;
std::vector<float> present_value_data;
std::vector<AttentionKernelType> skip_kernel_types; // skip some kernels if they do not supported this test case.
};
@ -49,6 +59,9 @@ void GetSelfAttentionData_Batch2_HeadSize32_NoBias_NoMask_PackedQKV(AttentionTes
void GetCrossAttentionData_HeadSize16_8(AttentionTestData& data);
void GetCrossAttentionData_HeadSize16(AttentionTestData& data);
void GetCrossAttentionDataWithPast(AttentionTestData& data);
void GetSelfAttentionDataWithPast(AttentionTestData& data);
bool SkipAttentionKernel(AttentionTestData& data, AttentionKernelType kernel_type);
} // namespace test
} // namespace onnxruntime

197
onnxruntime/test/contrib_ops/multihead_attention_op_test.cc
View file

@ -19,6 +19,11 @@ static void RunMultiHeadAttentionTest(
const std::vector<float>& kv_data, // packed_kv: [batch_size, kv_sequence_length, num_heads, 2, head_size]
const std::vector<float>& qkv_data, // packed_qkv: [batch_size, sequence_length, num_heads, 3, head_size]
const std::vector<float>& bias_data, // bias: [hidden_size + hidden_size + v_hidden_size] or empty
const std::vector<float>& rel_pos_bias_data, // relative_position_bias: [1, num_heads, sequence_length, total_sequence_length]
const std::vector<float>& past_key_data, // past_key: [batch_size, num_heads, kv_sequence_length, head_size]
const std::vector<float>& past_value_data, // past_value: [batch_size, num_heads, kv_sequence_length, head_size]
const std::vector<float>& present_key_data, // present_key: [batch_size, num_heads, total_sequence_length, head_size]
const std::vector<float>& present_value_data, // present_value: [batch_size, num_heads, total_sequence_length, head_size]
const std::vector<int32_t>& key_padding_mask_data, // key_padding_mask: see below
AttentionMaskType mask_type, // 1 for [batch_size], 2 for [batch_size, kv_sequence_length]
const std::vector<float>& output_data, // output: [batch_size, sequence_length, v_hidden_size]
@ -28,6 +33,7 @@ static void RunMultiHeadAttentionTest(
int kv_sequence_length,
int hidden_size,
int v_hidden_size,
bool is_static_kv = true,
bool use_float16 = false,
bool disable_cpu = true, // not supported in cpu right now.
bool disable_cuda = false,
@ -46,10 +52,16 @@ static void RunMultiHeadAttentionTest(
tester.AddAttribute<float>("mask_filter_value", static_cast<float>(-10000.0f));
std::vector<int64_t> query_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> key_dims = {batch_size, kv_sequence_length, hidden_size};
std::vector<int64_t> value_dims = {batch_size, kv_sequence_length, v_hidden_size};
std::vector<int64_t> key_dims = {batch_size, is_static_kv ? kv_sequence_length : sequence_length, hidden_size};
std::vector<int64_t> value_dims = {batch_size, is_static_kv ? kv_sequence_length : sequence_length, v_hidden_size};
std::vector<int64_t> bias_dims = {hidden_size + hidden_size + v_hidden_size};
std::vector<int64_t> rel_pos_bias_dims = {1, num_heads, sequence_length, sequence_length + kv_sequence_length};
std::vector<int64_t> past_key_dims = {batch_size, num_heads, kv_sequence_length, hidden_size / num_heads};
std::vector<int64_t> past_value_dims = past_key_dims;
std::vector<int64_t> output_dims = {batch_size, sequence_length, v_hidden_size};
std::vector<int64_t> present_key_dims =
{batch_size, num_heads, is_static_kv ? kv_sequence_length : sequence_length + kv_sequence_length, hidden_size / num_heads};
std::vector<int64_t> present_value_dims = present_key_dims;
std::vector<float> query = (qkv_data.size() > 0 ? qkv_data : query_data);
std::vector<float> key;
@ -79,12 +91,16 @@ static void RunMultiHeadAttentionTest(
if (key.size()) {
tester.AddInput<MLFloat16>("key", key_dims, ToFloat16(key));
} else if (past_key_data.size() && is_static_kv == true) {
tester.AddInput<MLFloat16>("key", past_key_dims, ToFloat16(past_key_data));
} else {
tester.AddOptionalInputEdge<MLFloat16>();
}
if (value.size()) {
tester.AddInput<MLFloat16>("value", value_dims, ToFloat16(value));
} else if (past_value_data.size() && is_static_kv == true) {
tester.AddInput<MLFloat16>("value", past_value_dims, ToFloat16(past_value_data));
} else {
tester.AddOptionalInputEdge<MLFloat16>();
}
@ -101,22 +117,56 @@ static void RunMultiHeadAttentionTest(
tester.AddOptionalInputEdge<int32_t>();
}
if (rel_pos_bias_data.size()) {
tester.AddInput<MLFloat16>("relative_position_bias", rel_pos_bias_dims, ToFloat16(rel_pos_bias_data));
} else {
tester.AddOptionalInputEdge<MLFloat16>();
}
if (past_key_data.size() && is_static_kv == false) {
tester.AddInput<MLFloat16>("past_key", past_key_dims, ToFloat16(past_key_data));
} else {
tester.AddOptionalInputEdge<MLFloat16>();
}
if (past_value_data.size() && is_static_kv == false) {
tester.AddInput<MLFloat16>("past_value", past_value_dims, ToFloat16(past_value_data));
} else {
tester.AddOptionalInputEdge<MLFloat16>();
}
constexpr float rel_error = 0.0f;
constexpr float abs_error = 0.05f;
tester.AddOutput<MLFloat16>("output", output_dims, ToFloat16(output_data), /*sort*/ false, rel_error, abs_error);
if (present_key_data.size()) {
tester.AddOutput<MLFloat16>("present_key", present_key_dims, ToFloat16(present_key_data), /*sort*/ false, rel_error, abs_error);
} else {
tester.AddOptionalOutputEdge<MLFloat16>();
}
if (present_value_data.size()) {
tester.AddOutput<MLFloat16>("present_value", present_value_dims, ToFloat16(present_value_data), /*sort*/ false, rel_error, abs_error);
} else {
tester.AddOptionalOutputEdge<MLFloat16>();
}
} else {
tester.AddInput<float>("query", query_dims, query);
if (key.size()) {
tester.AddInput<float>("key", key_dims, key);
} else if (past_key_data.size() && is_static_kv == true) {
tester.AddInput<float>("key", past_key_dims, past_key_data);
} else {
tester.AddOptionalInputEdge<MLFloat16>();
tester.AddOptionalInputEdge<float>();
}
if (value.size()) {
tester.AddInput<float>("value", value_dims, value);
} else if (past_value_data.size() && is_static_kv == true) {
tester.AddInput<float>("value", past_value_dims, past_value_data);
} else {
tester.AddOptionalInputEdge<MLFloat16>();
tester.AddOptionalInputEdge<float>();
}
if (bias_data.size()) {
@ -131,9 +181,39 @@ static void RunMultiHeadAttentionTest(
tester.AddOptionalInputEdge<int32_t>();
}
if (rel_pos_bias_data.size()) {
tester.AddInput<float>("relative_position_bias", rel_pos_bias_dims, rel_pos_bias_data);
} else {
tester.AddOptionalInputEdge<float>();
}
if (past_key_data.size() && is_static_kv == false) {
tester.AddInput<float>("past_key", past_key_dims, past_key_data);
} else {
tester.AddOptionalInputEdge<float>();
}
if (past_value_data.size() && is_static_kv == false) {
tester.AddInput<float>("past_value", past_value_dims, past_value_data);
} else {
tester.AddOptionalInputEdge<float>();
}
constexpr float rel_error = 0.0f;
constexpr float abs_error = 0.02f;
tester.AddOutput<float>("output", output_dims, output_data, /*sort*/ false, rel_error, abs_error);
if (present_key_data.size()) {
tester.AddOutput<float>("present_key", present_key_dims, present_key_data, /*sort*/ false, rel_error, abs_error);
} else {
tester.AddOptionalOutputEdge<float>();
}
if (present_value_data.size()) {
tester.AddOutput<float>("present_value", present_value_dims, present_value_data, /*sort*/ false, rel_error, abs_error);
} else {
tester.AddOptionalOutputEdge<float>();
}
}
if (enable_cuda) {
@ -163,6 +243,11 @@ static void RunMultiHeadAttentionKernel(
const std::vector<float>& kv_data, // packed_kv: [batch_size, kv_sequence_length, num_heads, 2, head_size]
const std::vector<float>& qkv_data, // packed_qkv: [batch_size, sequence_length, num_heads, 3, head_size]
const std::vector<float>& bias_data, // bias: [hidden_size + hidden_size + v_hidden_size]
const std::vector<float>& rel_pos_bias_data, // relative_position_bias: [1, num_heads, sequence_length, total_sequence_length]
const std::vector<float>& past_key_data, // past_key: [batch_size, num_heads, kv_sequence_length, head_size]
const std::vector<float>& past_value_data, // past_value: [batch_size, num_heads, kv_sequence_length, head_size]
const std::vector<float>& present_key_data, // present_key: [batch_size, num_heads, total_sequence_length, head_size]
const std::vector<float>& present_value_data, // present_value: [batch_size, num_heads, total_sequence_length, head_size]
const std::vector<int32_t>& key_padding_mask_data, // key_padding_mask: see below
AttentionMaskType mask_type, // 1 for [batch_size], 2 for [batch_size, kv_sequence_length]
const std::vector<float>& output_data, // output: [batch_size, sequence_length, v_hidden_size]
@ -174,6 +259,7 @@ static void RunMultiHeadAttentionKernel(
int v_hidden_size,
AttentionKernelType kernel_type,
bool use_float16 = true,
bool is_static_kv = true,
bool disable_cpu = true, // not supported in cpu right now.
bool disable_cuda = false,
bool disable_rocm = true) {
@ -185,9 +271,10 @@ static void RunMultiHeadAttentionKernel(
{onnxruntime::contrib::attention::kDisableFusedCrossAttention, "0"},
{onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "0"}}};
RunMultiHeadAttentionTest(
query_data, key_data, value_data, kv_data, qkv_data, bias_data, key_padding_mask_data, mask_type, output_data,
num_heads, batch_size, sequence_length, kv_sequence_length, hidden_size, v_hidden_size,
use_float16, disable_cpu, disable_cuda, disable_rocm);
query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data,
past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data,
mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length,
hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm);
return;
}
@ -199,9 +286,10 @@ static void RunMultiHeadAttentionKernel(
{onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"},
{onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}};
RunMultiHeadAttentionTest(
query_data, key_data, value_data, kv_data, qkv_data, bias_data, key_padding_mask_data, mask_type, output_data,
num_heads, batch_size, sequence_length, kv_sequence_length, hidden_size, v_hidden_size,
use_float16, disable_cpu, disable_cuda, disable_rocm);
query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data,
past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data,
mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length,
hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm);
return;
}
@ -213,9 +301,10 @@ static void RunMultiHeadAttentionKernel(
{onnxruntime::contrib::attention::kDisableFusedCrossAttention, "0"},
{onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}};
RunMultiHeadAttentionTest(
query_data, key_data, value_data, kv_data, qkv_data, bias_data, key_padding_mask_data, mask_type, output_data,
num_heads, batch_size, sequence_length, kv_sequence_length, hidden_size, v_hidden_size,
use_float16, disable_cpu, disable_cuda, disable_rocm);
query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data,
past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data,
mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length,
hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm);
return;
}
@ -228,9 +317,10 @@ static void RunMultiHeadAttentionKernel(
{onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"},
{onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "0"}}};
RunMultiHeadAttentionTest(
query_data, key_data, value_data, kv_data, qkv_data, bias_data, key_padding_mask_data, mask_type, output_data,
num_heads, batch_size, sequence_length, kv_sequence_length, hidden_size, v_hidden_size,
use_float16, disable_cpu, disable_cuda, disable_rocm);
query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data,
past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data,
mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length,
hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm);
return;
}
#endif
@ -243,9 +333,10 @@ static void RunMultiHeadAttentionKernel(
{onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"},
{onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}};
RunMultiHeadAttentionTest(
query_data, key_data, value_data, kv_data, qkv_data, bias_data, key_padding_mask_data, mask_type, output_data,
num_heads, batch_size, sequence_length, kv_sequence_length, hidden_size, v_hidden_size,
use_float16, disable_cpu, disable_cuda, disable_rocm);
query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data,
past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data,
mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length,
hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm);
}
}
@ -256,9 +347,11 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data) {
AttentionKernelType kernel_type = AttentionKernelType::AttentionKernel_Unfused;
if (!SkipAttentionKernel(data, kernel_type)) {
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp32_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp32_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
#if USE_FLASH_ATTENTION
@ -267,18 +360,22 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data) {
kernel_type = AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention;
if (!SkipAttentionKernel(data, kernel_type)) {
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp32_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp32_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
}
#endif
kernel_type = AttentionKernelType::AttentionKernel_Default;
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp32_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp32_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
if (data.fp16_output_data.size() > 0) {
@ -286,34 +383,42 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data) {
AttentionKernelType kernel_type = AttentionKernelType::AttentionKernel_TrtFusedCrossAttention;
if (!SkipAttentionKernel(data, kernel_type)) {
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
kernel_type = AttentionKernelType::AttentionKernel_TrtFusedAttention;
if (!SkipAttentionKernel(data, kernel_type)) {
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
#if USE_FLASH_ATTENTION
kernel_type = AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention;
if (!SkipAttentionKernel(data, kernel_type)) {
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
#endif
kernel_type = AttentionKernelType::AttentionKernel_Default;
RunMultiHeadAttentionKernel(
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data, data.key_padding_mask_data, data.mask_type,
data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length,
data.hidden_size, data.v_hidden_size, kernel_type, use_float16);
data.query_data, data.key_data, data.value_data, data.kv_data, data.qkv_data, data.bias_data,
data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data,
data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data,
data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size,
data.v_hidden_size, kernel_type, use_float16, data.is_static_kv);
}
}
@ -370,5 +475,17 @@ TEST(MultiHeadAttentionTest, CrossAttention_Batch1_HeadSize16) {
RunMultiHeadAttentionTests(data);
}
TEST(MultiHeadAttentionTest, CrossAttentionWithPast) {
AttentionTestData data;
GetCrossAttentionDataWithPast(data);
RunMultiHeadAttentionTests(data);
}
TEST(MultiHeadAttentionTest, SelfAttentionWithPast) {
AttentionTestData data;
GetSelfAttentionDataWithPast(data);
RunMultiHeadAttentionTests(data);
}
} // namespace test
} // namespace onnxruntime

580
onnxruntime/test/python/transformers/test_parity_t5_mha.py Normal file
View file

@ -0,0 +1,580 @@
# --------------------------------------------------------------------------
# Copyright 2020 The HuggingFace Inc. team
#
# 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
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
import unittest
import numpy as np
import torch
from torch import nn
torch.manual_seed(0)
def create_t5_mha_graph(
batch_size,
seq_len,
kv_sequence_length,
head_size,
num_heads,
use_past,
is_static_kv,
):
from onnx import TensorProto, helper
use_present = not use_past
if not is_static_kv and use_past:
use_present = True
use_rpb = not is_static_kv
use_mask = not use_rpb
past_sequence_length = kv_sequence_length
total_sequence_length = kv_sequence_length if is_static_kv else seq_len
if not is_static_kv:
kv_sequence_length = seq_len
if not is_static_kv and use_past:
total_sequence_length += past_sequence_length
rpb_length = total_sequence_length if use_past else seq_len
nodes = [
helper.make_node(
"MultiHeadAttention",
[
"query",
"key" if use_present or is_static_kv else "",
"value" if use_present or is_static_kv else "",
"", # bias
"key_padding_mask" if use_mask else "",
"relative_position_bias" if use_rpb else "",
"past_key" if use_past and not is_static_kv else "",
"past_value" if use_past and not is_static_kv else "",
],
[
"output",
"present_key" if use_present else "",
"present_value" if use_present else "",
],
"MHA_0",
num_heads=num_heads,
mask_filter_value=-10000.0,
scale=1.0,
domain="com.microsoft",
),
]
initializers = []
hidden_size = head_size * num_heads
graph_inputs = [
helper.make_tensor_value_info("query", TensorProto.FLOAT, [batch_size, seq_len, hidden_size]),
]
graph_outputs = [
helper.make_tensor_value_info("output", TensorProto.FLOAT, [batch_size, seq_len, hidden_size]),
]
if use_mask:
graph_inputs.append(
helper.make_tensor_value_info("key_padding_mask", TensorProto.INT32, [batch_size, kv_sequence_length])
)
if use_rpb:
graph_inputs.append(
helper.make_tensor_value_info(
"relative_position_bias", TensorProto.FLOAT, [1, num_heads, seq_len, rpb_length]
)
)
if use_past and not is_static_kv:
graph_inputs.append(
helper.make_tensor_value_info(
"past_key", TensorProto.FLOAT, [batch_size, num_heads, past_sequence_length, head_size]
)
)
graph_inputs.append(
helper.make_tensor_value_info(
"past_value", TensorProto.FLOAT, [batch_size, num_heads, past_sequence_length, head_size]
)
)
if use_present:
graph_inputs.append(
helper.make_tensor_value_info("key", TensorProto.FLOAT, [batch_size, kv_sequence_length, hidden_size])
)
graph_inputs.append(
helper.make_tensor_value_info("value", TensorProto.FLOAT, [batch_size, kv_sequence_length, hidden_size])
)
elif is_static_kv and use_past:
graph_inputs.append(
helper.make_tensor_value_info(
"key", TensorProto.FLOAT, [batch_size, num_heads, past_sequence_length, head_size]
)
)
graph_inputs.append(
helper.make_tensor_value_info(
"value", TensorProto.FLOAT, [batch_size, num_heads, past_sequence_length, head_size]
)
)
if use_present:
graph_outputs.append(
helper.make_tensor_value_info(
"present_key", TensorProto.FLOAT, [batch_size, num_heads, total_sequence_length, head_size]
)
)
graph_outputs.append(
helper.make_tensor_value_info(
"present_value", TensorProto.FLOAT, [batch_size, num_heads, total_sequence_length, head_size]
)
)
graph = helper.make_graph(
nodes,
"T5_MHA_Graph",
graph_inputs,
graph_outputs,
initializers,
)
model = helper.make_model(graph)
return model.SerializeToString()
class T5Config:
def __init__(self, is_decoder, batch_size, seq_len, kv_sequence_length, num_heads, head_size, use_past):
self.is_decoder = is_decoder
self.d_model = num_heads * head_size
self.key_value_proj_dim = head_size
self.n_heads = num_heads
self.inner_dim = self.n_heads * self.key_value_proj_dim
# ORT parameters
self.batch_size = batch_size
self.seq_len = seq_len
self.kv_sequence_length = kv_sequence_length
self.head_size = head_size
self.num_heads = num_heads
self.hidden_size = self.d_model
self.use_past = use_past
class T5Attention(nn.Module):
def __init__(self, config: T5Config, is_static_kv):
super().__init__()
self.is_decoder = config.is_decoder
self.is_static_kv = is_static_kv
self.has_relative_attention_bias = not self.is_static_kv
self.d_model = config.d_model
self.key_value_proj_dim = config.head_size
self.n_heads = config.num_heads
self.inner_dim = self.n_heads * self.key_value_proj_dim
# Mesh TensorFlow initialization to avoid scaling before softmax
self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)
self.k = nn.Linear(self.d_model, self.inner_dim, bias=False)
self.v = nn.Linear(self.d_model, self.inner_dim, bias=False)
# self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)
self.pruned_heads = set()
# ORT parameters
self.batch_size = config.batch_size
self.seq_len = config.seq_len
self.kv_sequence_length = config.kv_sequence_length
self.head_size = config.head_size
self.num_heads = config.num_heads
self.hidden_size = self.d_model
self.use_past = config.use_past
# Create onnx graph
self.onnx_graph = create_t5_mha_graph(
self.batch_size,
self.seq_len,
self.kv_sequence_length,
self.head_size,
self.num_heads,
self.use_past,
is_static_kv,
)
def create_inputs(self):
hidden_states = torch.normal(mean=0.5, std=0.1, size=(self.batch_size, self.seq_len, self.hidden_size)).to(
torch.float32
)
key_value_states = torch.normal(
mean=0.5, std=0.1, size=(self.batch_size, self.kv_sequence_length, self.hidden_size)
).to(torch.float32)
past_key = torch.normal(
mean=0.5, std=0.1, size=(self.batch_size, self.num_heads, self.kv_sequence_length, self.head_size)
).to(torch.float32)
past_value = torch.normal(
mean=0.5, std=0.1, size=(self.batch_size, self.num_heads, self.kv_sequence_length, self.head_size)
).to(torch.float32)
past_key_value = (past_key, past_value)
attention_mask = torch.ones((self.batch_size, self.kv_sequence_length)).to(torch.float32)
position_bias_length = self.seq_len if not self.use_past else self.kv_sequence_length + self.seq_len
position_bias = torch.normal(
mean=0.5, std=0.1, size=(1, self.num_heads, position_bias_length, position_bias_length)
).to(torch.float32)
return hidden_states, key_value_states, past_key_value, attention_mask, position_bias
def torch_forward(
self,
hidden_states,
key_value_states=None,
past_key_value=None,
mask=None,
position_bias=None,
use_cache=False,
query_length=None,
):
"""
Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
"""
# Input is (batch_size, seq_length, dim)
# Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length)
# past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head)
batch_size, seq_length = hidden_states.shape[:2]
real_seq_length = seq_length
if past_key_value is not None:
assert (
len(past_key_value) == 2
), f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length
key_length = real_seq_length if key_value_states is None else key_value_states.shape[1]
def shape(states):
"""projection"""
return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
def unshape(states):
"""reshape"""
return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)
def project(hidden_states, proj_layer, key_value_states, past_key_value):
"""projects hidden states correctly to key/query states"""
if key_value_states is None:
# self-attn
# (batch_size, n_heads, seq_length, dim_per_head)
hidden_states = shape(proj_layer(hidden_states))
elif past_key_value is None:
# cross-attn
# (batch_size, n_heads, seq_length, dim_per_head)
hidden_states = shape(proj_layer(key_value_states))
if past_key_value is not None:
if key_value_states is None:
# self-attn
# (batch_size, n_heads, key_length, dim_per_head)
hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
elif past_key_value.shape[2] != key_value_states.shape[1]:
# checking that the `sequence_length` of the `past_key_value` is the same as
# the provided `key_value_states` to support prefix tuning
# cross-attn
# (batch_size, n_heads, seq_length, dim_per_head)
hidden_states = shape(proj_layer(key_value_states))
else:
# cross-attn
hidden_states = past_key_value
return hidden_states
# get query states
query_states = shape(self.q(hidden_states)) # (batch_size, n_heads, seq_length, dim_per_head)
# get key/value states
key_states = project(
hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None
)
value_states = project(
hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None
)
# compute scores
scores = torch.matmul(
query_states, key_states.transpose(3, 2)
) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
if position_bias is None:
if not self.has_relative_attention_bias:
position_bias = torch.zeros(
(1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype
)
else:
assert position_bias is not None
# if key and values are already calculated
# we want only the last query position bias
if past_key_value is not None and position_bias is not None:
position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
if mask is not None:
# Adjust onnx mask shape
mask = (1 - mask.unsqueeze(1).unsqueeze(2)) * torch.finfo(torch.float32).min
position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length)
if self.pruned_heads:
mask = torch.ones(position_bias.shape[1])
mask[list(self.pruned_heads)] = 0
position_bias_masked = position_bias[:, mask.bool()]
else:
position_bias_masked = position_bias
scores += position_bias_masked
attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
scores
) # (batch_size, n_heads, seq_length, key_length)
attn_output = unshape(torch.matmul(attn_weights, value_states)) # (batch_size, seq_length, dim)
# attn_output = self.o(attn_output) # ORT places this matmul outside of MHA op
present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None
outputs = (attn_output,) + (present_key_value_state,)
return outputs
def ort_forward(
self,
hidden_states,
key_value_states=None,
past_key_value=None,
mask=None,
position_bias=None,
use_cache=False,
query_length=None,
):
import onnxruntime
sess_options = onnxruntime.SessionOptions()
cuda_providers = ["CUDAExecutionProvider"]
if cuda_providers[0] not in onnxruntime.get_available_providers():
return None
ort_session = onnxruntime.InferenceSession(self.onnx_graph, sess_options, providers=cuda_providers)
batch_size, seq_length = hidden_states.shape[:2]
real_seq_length = seq_length
if past_key_value is not None:
assert (
len(past_key_value) == 2
), f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length
def project(hidden_states, proj_layer, key_value_states, past_key_value):
if key_value_states is None:
# self-attn
hidden_states = proj_layer(hidden_states)
elif past_key_value is None:
# cross-attn
# (batch_size, n_heads, seq_length, dim_per_head)
hidden_states = proj_layer(key_value_states)
return hidden_states
# get query states
query_states = self.q(hidden_states) # (batch_size, n_heads, seq_length, dim_per_head)
# get key/value states
key_states = project(
hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None
)
value_states = project(
hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None
)
if past_key_value is not None and position_bias is not None:
position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
torch_key_padding_mask = mask.to(torch.int32) if mask is not None else None
torch_position_bias = position_bias if position_bias is not None else None
ort_inputs = None
if past_key_value is None:
ort_inputs = {
"query": np.ascontiguousarray(query_states.detach().numpy()),
"key": np.ascontiguousarray(key_states.detach().numpy()),
"value": np.ascontiguousarray(value_states.detach().numpy()),
}
if torch_key_padding_mask is not None:
ort_inputs["key_padding_mask"] = np.ascontiguousarray(torch_key_padding_mask.detach().numpy())
if torch_position_bias is not None:
ort_inputs["relative_position_bias"] = np.ascontiguousarray(torch_position_bias.detach().numpy())
else:
torch_past_key = past_key_value[0]
torch_past_value = past_key_value[1]
ort_inputs = {
"query": np.ascontiguousarray(query_states.detach().numpy()),
}
if self.is_static_kv:
ort_inputs["key"] = np.ascontiguousarray(torch_past_key.detach().numpy())
ort_inputs["value"] = np.ascontiguousarray(torch_past_value.detach().numpy())
else:
ort_inputs["past_key"] = np.ascontiguousarray(torch_past_key.detach().numpy())
ort_inputs["past_value"] = np.ascontiguousarray(torch_past_value.detach().numpy())
ort_inputs["key"] = np.ascontiguousarray(key_states.detach().numpy())
ort_inputs["value"] = np.ascontiguousarray(value_states.detach().numpy())
if torch_key_padding_mask is not None:
ort_inputs["key_padding_mask"] = np.ascontiguousarray(torch_key_padding_mask.detach().numpy())
if torch_position_bias is not None:
ort_inputs["relative_position_bias"] = np.ascontiguousarray(torch_position_bias.detach().numpy())
ort_output = ort_session.run(None, ort_inputs)
output = None
if past_key_value is not None and self.is_static_kv:
output = torch.tensor(ort_output)
else:
output = (torch.tensor(ort_output[0]),) + ((torch.tensor(ort_output[1]), torch.tensor(ort_output[2])),)
return output
def compare_t5_cross_attention_decoder(batch_size, seq_len, num_heads, head_size, kv_sequence_length):
config = T5Config(
is_decoder=True,
batch_size=batch_size,
seq_len=seq_len,
kv_sequence_length=kv_sequence_length,
num_heads=num_heads,
head_size=head_size,
use_past=True,
)
T5CrossAttention = T5Attention(config, is_static_kv=True)
hidden_states, key_value_states, past_key_value, attention_mask, _ = T5CrossAttention.create_inputs()
torch_output = T5CrossAttention.torch_forward(
hidden_states, key_value_states, past_key_value, attention_mask, position_bias=None, use_cache=False
)
ort_output = T5CrossAttention.ort_forward(
hidden_states, key_value_states, past_key_value, attention_mask, position_bias=None, use_cache=False
)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
def compare_t5_cross_attention_decoder_init(batch_size, seq_len, num_heads, head_size, kv_sequence_length):
config = T5Config(
is_decoder=True,
batch_size=batch_size,
seq_len=seq_len,
kv_sequence_length=kv_sequence_length,
num_heads=num_heads,
head_size=head_size,
use_past=False,
)
T5CrossAttention = T5Attention(config, is_static_kv=True)
hidden_states, key_value_states, _, attention_mask, _ = T5CrossAttention.create_inputs()
torch_output = T5CrossAttention.torch_forward(
hidden_states, key_value_states, None, attention_mask, position_bias=None, use_cache=True
)
ort_output = T5CrossAttention.ort_forward(
hidden_states, key_value_states, None, attention_mask, position_bias=None, use_cache=True
)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
assert torch.allclose(torch_output[1][0], ort_output[1][0], atol=1e-4)
assert torch.allclose(torch_output[1][1], ort_output[1][1], atol=1e-4)
def compare_t5_self_attention_decoder_init(batch_size, seq_len, num_heads, head_size, kv_sequence_length):
config = T5Config(
is_decoder=True,
batch_size=batch_size,
seq_len=seq_len,
kv_sequence_length=kv_sequence_length,
num_heads=num_heads,
head_size=head_size,
use_past=False,
)
T5CrossAttention = T5Attention(config, is_static_kv=False)
hidden_states, _, _, _, position_bias = T5CrossAttention.create_inputs()
torch_output = T5CrossAttention.torch_forward(
hidden_states, None, None, mask=None, position_bias=position_bias, use_cache=True
)
ort_output = T5CrossAttention.ort_forward(
hidden_states, None, None, mask=None, position_bias=position_bias, use_cache=True
)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
assert torch.allclose(torch_output[1][0], ort_output[1][0], atol=1e-4)
assert torch.allclose(torch_output[1][1], ort_output[1][1], atol=1e-4)
def compare_t5_self_attention_decoder(batch_size, seq_len, num_heads, head_size, kv_sequence_length):
config = T5Config(
is_decoder=True,
batch_size=batch_size,
seq_len=seq_len,
kv_sequence_length=kv_sequence_length,
num_heads=num_heads,
head_size=head_size,
use_past=True,
)
T5CrossAttention = T5Attention(config, is_static_kv=False)
hidden_states, _, past_key_value, _, position_bias = T5CrossAttention.create_inputs()
torch_output = T5CrossAttention.torch_forward(
hidden_states, None, past_key_value, mask=None, position_bias=position_bias, use_cache=True
)
ort_output = T5CrossAttention.ort_forward(
hidden_states, None, past_key_value, mask=None, position_bias=position_bias, use_cache=True
)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
assert torch.allclose(torch_output[1][0], ort_output[1][0], atol=1e-4)
assert torch.allclose(torch_output[1][1], ort_output[1][1], atol=1e-4)
class TestT5MHAParity(unittest.TestCase):
def setUp(self):
self.batch_size = 5
self.seq_len = 2
self.num_heads = 2
self.head_size = 4
self.kv_sequence_length = 3
def test_t5_cross_attention_decoder_init(self):
compare_t5_cross_attention_decoder_init(
self.batch_size, self.seq_len, self.num_heads, self.head_size, self.kv_sequence_length
)
def test_t5_self_attention_decoder_init(self):
compare_t5_self_attention_decoder_init(
self.batch_size, self.seq_len, self.num_heads, self.head_size, self.kv_sequence_length
)
def test_t5_cross_attention_decoder(self):
compare_t5_cross_attention_decoder(
self.batch_size, self.seq_len, self.num_heads, self.head_size, self.kv_sequence_length
)
def test_t5_self_attention_decoder(self):
compare_t5_self_attention_decoder(
self.batch_size, self.seq_len, self.num_heads, self.head_size, self.kv_sequence_length
)
if __name__ == "__main__":
unittest.main()