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

DecoderMaskedMultiHeadAttention CPU kernel. (#22292)

Browse source 
### Description
DecoderMaskedMultiHeadAttention CPU kernel.
This commit is contained in:
mindest 2024-10-13 05:43:00 +09:00 committed by GitHub
parent 332173509d
commit 1fa219d7d5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 810 additions and 129 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
docs/ContribOperators.md
View file

@ -1167,7 +1167,7 @@ This version of the operator has been available since version 1 of the 'com.micr
<dt><tt>mask_index</tt> (optional) : M</dt>
<dd>Mask values of shape (batch_size, total_sequence_length) or (batch_size, kv_sequence_length)</dd>
<dt><tt>attention_bias</tt> (optional) : T</dt>
<dd>additional add to QxK' with shape (batch_size, num_heads, sequence_length, total_sequence_length)</dd>
<dd>additional add to QxK' with shape (batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)</dd>
<dt><tt>past_key</tt> (optional) : T</dt>
<dd>past state for key with shape (batch_size, num_heads, past_sequence_length, head_size) for self attentionWhen past_present_share_buffer is set, its shape is (batch_size, num_heads, max_sequence_length, head_size). The keys buffer is re-ordered in such a way that its virtual sub-tensor of shape (batch_size, num_heads, max_sequence_length, head_size) which may be perceived as being of shape (batch_size, num_heads, max_sequence_length, head_size / x, x) is reordered to become (batch_size, num_heads, head_size / x, max_sequence_length, x) where `x = 16 / sizeof(T)`.</dd>
<dt><tt>past_value</tt> (optional) : T</dt>
@ -1175,9 +1175,9 @@ This version of the operator has been available since version 1 of the 'com.micr
<dt><tt>past_sequence_length</tt> (optional) : M</dt>
<dd>When past_present_share_buffer is used, it is required to specify past_sequence_length (could be 0).Cross Attention doesn't need this input.</dd>
<dt><tt>beam_width</tt> (optional) : M</dt>
<dd>The beam width that is being used while decoding.If not provided, the beam width will be assumed to be 1.</dd>
<dd>The beam width that is being used while decoding. If not provided, the beam width will be assumed to be 1.</dd>
<dt><tt>cache_indirection</tt> (optional) : M</dt>
<dd>A buffer of shape [batch_size, beam_width, max_output_length] where an [i, j, k] entry specifieswhich beam the 'k' th token came from for the 'j' th beam for batch 'i' in the current iteration</dd>
<dd>A buffer of shape [batch_size, beam_width, max_output_length] where an `[i, j, k]` entry specifies which beam the `k`-th token came from for the `j`-th beam for batch `i` in the current iteration</dd>
<dt><tt>bias</tt> (optional) : T</dt>
<dd>Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection</dd>
</dl>
@ -1192,7 +1192,7 @@ This version of the operator has been available since version 1 of the 'com.micr
<dt><tt>present_value</tt> (optional) : T</dt>
<dd>present state for value with shape (batch_size, num_heads, total_sequence_length, head_size). If past_present_share_buffer is set, its shape is (batch_size, num_heads, max_sequence_length, head_size), while effective_seq_length = (past_sequence_length + kv_sequence_length).</dd>
<dt><tt>qk</tt> (optional) : V</dt>
<dd>normalized Q * K, of shape (batch_size, num_heads, 1, head_size). </dd>
<dd>normalized Q * K, of shape (batch_size, num_heads, 1, total_sequence_length). </dd>
</dl>
#### Type Constraints
@ -1261,9 +1261,9 @@ This version of the operator has been available since version 1 of the 'com.micr
<dt><tt>past_sequence_length</tt> : M</dt>
<dd>When past_present_share_buffer is used, it is required to specify past_sequence_length (could be 0).</dd>
<dt><tt>beam_width</tt> (optional) : M</dt>
<dd>The beam width that is being used while decoding.If not provided, the beam width will be assumed to be 1.</dd>
<dd>The beam width that is being used while decoding. If not provided, the beam width will be assumed to be 1.</dd>
<dt><tt>cache_indirection</tt> (optional) : M</dt>
<dd>A buffer of shape [batch_size, beam_width, max_output_length] where an [i, j, k] entry specifieswhich beam the 'k' th token came from for the 'j' th beam for batch 'i' in the current iteration</dd>
<dd>A buffer of shape [batch_size, beam_width, max_output_length] where an `[i, j, k]` entry specifies which beam the `k`-th token came from for the `j`-th beam for batch `i` in the current iteration</dd>
</dl>
#### Outputs

1
docs/OperatorKernels.md
View file

@ -468,6 +468,7 @@ Do not modify directly.*
|CDist|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(double), tensor(float)|
|ConvTransposeWithDynamicPads|*in* X:**T**<br> *in* W:**T**<br> *in* Pads:**tensor(int64)**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
|CropAndResize|*in* X:**T1**<br> *in* rois:**T1**<br> *in* batch_indices:**T2**<br> *in* crop_size:**T2**<br> *out* Y:**T1**|1+|**T1** = tensor(float)<br/> **T2** = tensor(int32)|
|DecoderMaskedMultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* mask_index:**M**<br> *in* attention_bias:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* past_sequence_length:**M**<br> *in* beam_width:**M**<br> *in* cache_indirection:**M**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**<br> *out* qk:**V**|1+|**T** = tensor(float)|
|DequantizeLinear|*in* x:**T1**<br> *in* x_scale:**T2**<br> *in* x_zero_point:**T1**<br> *out* y:**T2**|1+|**T1** = tensor(int16), tensor(int32), tensor(int4), tensor(int8), tensor(uint16), tensor(uint4), tensor(uint8)<br/> **T2** = tensor(float)|
|DynamicQuantizeLSTM|*in* X:**T**<br> *in* W:**T2**<br> *in* R:**T2**<br> *in* B:**T**<br> *in* sequence_lens:**T1**<br> *in* initial_h:**T**<br> *in* initial_c:**T**<br> *in* P:**T**<br> *in* W_scale:**T**<br> *in* W_zero_point:**T2**<br> *in* R_scale:**T**<br> *in* R_zero_point:**T2**<br> *out* Y:**T**<br> *out* Y_h:**T**<br> *out* Y_c:**T**|1+|**T** = tensor(float)<br/> **T1** = tensor(int32)<br/> **T2** = tensor(int8), tensor(uint8)|
|DynamicQuantizeMatMul|*in* A:**T1**<br> *in* B:**T2**<br> *in* b_scale:**T1**<br> *in* b_zero_point:**T2**<br> *in* bias:**T1**<br> *out* Y:**T1**|1+|**T1** = tensor(float)<br/> **T2** = tensor(int8), tensor(uint8)|

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

@ -79,6 +79,45 @@ struct AttentionParameters {
AttentionQkvFormat qkv_format;
};
struct DecoderMaskedMultiHeadAttentionParams : AttentionParameters {
int beam_width = 1;
// Only NeoX style rotary embedding is supported
int rotary_embedding_dim = 0;
int t_step = 0;
// Whether to use multihead attention(excludes matmul and bias)
bool is_mha = false;
bool is_cross_attention = false;
bool is_packed_qkv = false;
// Useful to better use global memory bandwidth on certain CUDA architectures.
// Turned off by default for now until we fully understand performance implications
// for all types of workloads.
// Can be turned on by appropriate environment variable (see attention_common.h).
bool kv_data_in_flight = false;
void* q = nullptr;
void* q_bias = nullptr;
void* k = nullptr;
void* k_bias = nullptr;
void* v = nullptr;
void* v_bias = nullptr;
void* attention_bias = nullptr;
void* k_cache = nullptr;
void* v_cache = nullptr;
void* out = nullptr;
void* out_qk = nullptr;
const int32_t* cache_indir = nullptr;
const int32_t* mask = nullptr; // [B, total_sequence_length]
};
// Parameters deduced from node attributes and inputs/outputs.
struct PackedAttentionParameters {
int batch_size;

104
onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
View file

@ -36,18 +36,22 @@ class AttentionCPUBase : public AttentionBase {
int v_head_size, // head size of V (H_v)
int v_hidden_size, // hidden size of V (D_v)
const Tensor* attn_bias, // additive bias applied on scaled QK.
OpKernelContext* context) const {
OpKernelContext* context,
Tensor* output_qk = nullptr, // output buffer for QK (if needed)
int past_sequence_length = 0, // sequence length of past state
bool past_present_share_buffer = false) const {
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));
auto* tp = context->GetOperatorThreadPool();
int past_sequence_length = 0;
Tensor* present = nullptr;
if (present_key == nullptr && present_value == nullptr) {
present = GetPresent(context, past, batch_size, v_head_size, kv_sequence_length, past_sequence_length);
} else if (past_key != nullptr && past_value != nullptr) {
past_sequence_length = static_cast<int>(past_key->Shape().GetDims()[2]);
if (past_sequence_length == 0) {
if (present_key == nullptr && present_value == nullptr) {
present = GetPresent(context, past, batch_size, v_head_size, kv_sequence_length, past_sequence_length);
} else if (past_key != nullptr && past_value != nullptr) {
past_sequence_length = static_cast<int>(past_key->Shape().GetDims()[2]);
}
}
// Total sequence length including that of past state: T = P + L
@ -71,7 +75,7 @@ class AttentionCPUBase : public AttentionBase {
if (mask_data != nullptr) {
// Convert mask from boolean (0/1) to float (mask_filter_value/0.0f).
// Merge padding mask with causual mask, and broadcast to 3D (BxSxT).
// Merge padding mask with causal mask, and broadcast to 3D (BxSxT).
PrepareMask(mask_index_data, mask_index_dims, static_cast<T*>(mask_data),
causal, batch_size, sequence_length, past_sequence_length, mask_filter_value_);
DUMP_CPU_TENSOR("Mask3D", static_cast<T*>(mask_data), batch_size, sequence_length, total_sequence_length);
@ -85,10 +89,18 @@ class AttentionCPUBase : public AttentionBase {
T* present_key_data = present_key != nullptr ? present_key->MutableData<T>() : nullptr;
const T* past_value_data = past_value != nullptr ? past_value->Data<T>() : nullptr;
T* present_value_data = present_value != nullptr ? present_value->MutableData<T>() : nullptr;
T* output_qk_data = output_qk != nullptr ? output_qk->MutableData<T>() : nullptr;
const T* attn_bias_data = (attn_bias != nullptr) ? attn_bias->Data<T>() : nullptr;
auto attn_bias_dims = (attn_bias != nullptr) ? attn_bias->Shape().GetDims() : gsl::span<const int64_t>{};
// Used for DecoderMaskedMultiHeadAttention
int max_sequence_length = 0;
if (past_present_share_buffer) {
ORT_ENFORCE(past_key != nullptr && past_value != nullptr);
max_sequence_length = static_cast<int>(past_key->Shape().GetDims()[2]);
}
// Compute the attention score.
size_t bytes = SafeInt<size_t>(batch_size) * num_heads_ * sequence_length * total_sequence_length * sizeof(T);
auto attention_probs = allocator->Alloc(bytes);
@ -96,8 +108,9 @@ class AttentionCPUBase : public AttentionBase {
ComputeAttentionProbs<T>(static_cast<T*>(attention_probs), Q, K,
static_cast<T*>(mask_data),
batch_size, sequence_length, kv_sequence_length, past_sequence_length,
qk_head_size == 0 ? v_head_size : qk_head_size, past_data, past_key_data,
present_data, present_key_data, tp, scale, attn_bias_data, attn_bias_dims);
qk_head_size == 0 ? v_head_size : qk_head_size, past_data, past_key_data, present_data,
present_key_data, tp, scale, attn_bias_data, attn_bias_dims, output_qk_data,
past_present_share_buffer, max_sequence_length);
// Compute the attentionScore * Value: out_tmp(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v)
auto out_tmp_data =
@ -106,7 +119,8 @@ class AttentionCPUBase : public AttentionBase {
ComputeVxAttentionScore(output->MutableData<T>(), static_cast<T*>(out_tmp_data), static_cast<T*>(attention_probs),
V, batch_size, sequence_length, kv_sequence_length, past_sequence_length, v_head_size,
v_hidden_size, past_data, past_value_data, present_data, present_value_data, tp);
v_hidden_size, past_data, past_value_data, present_data, present_value_data, tp,
past_present_share_buffer, max_sequence_length);
return Status::OK();
}
@ -117,29 +131,32 @@ class AttentionCPUBase : public AttentionBase {
// 1 x mask_data(B, N, S, T)
// attention_probs(B, N, S, T) = Softmax(attention_probs)
template <typename T>
void ComputeAttentionProbs(T* attention_probs, // output buffer with size BxNxSxT
const T* Q, // Q data. Its size is BxNxSxH
const T* K, // k data. Its size is BxNxLxH
T* mask_data, // buffer for mask data.
int batch_size, // batch size of self-attention
int sequence_length, // sequence length of self-attention (S)
int kv_sequence_length, // sequence length of cross-attention (L)
int past_sequence_length, // sequence length of past state
int head_size, // head size of self-attention
const T* past, // past state
const T* past_key, // past key only (if not using past state)
T* present, // present state
T* present_key, // present key only (if not using present state)
ThreadPool* tp, // thread pool
float scale, // scale factor
const T* attn_bias_data, // attention bias
gsl::span<const int64_t> attn_bias_dims // attention bias shape
) const {
void ComputeAttentionProbs(T* attention_probs, // output buffer with size BxNxSxT
const T* Q, // Q data. Its size is BxNxSxH
const T* K, // k data. Its size is BxNxLxH
T* mask_data, // buffer for mask data.
int batch_size, // batch size of self-attention
int sequence_length, // sequence length of self-attention (S)
int kv_sequence_length, // sequence length of cross-attention (L)
int past_sequence_length, // sequence length of past state
int head_size, // head size of self-attention
const T* past, // past state
const T* past_key, // past key only (if not using past state)
T* present, // present state
T* present_key, // present key only (if not using present state)
ThreadPool* tp, // thread pool
float scale, // scale factor
const T* attn_bias_data, // attention bias
gsl::span<const int64_t> attn_bias_dims, // attention bias shape
T* output_qk_data = nullptr, // scaled output QK buffer
bool past_present_share_buffer = false,
int max_sequence_length = 0) const {
const int total_sequence_length = past_sequence_length + kv_sequence_length; // T = P + L
const size_t past_chunk_length = static_cast<size_t>(past_sequence_length) * head_size; // P x H
const size_t q_input_chunk_length = static_cast<size_t>(sequence_length) * head_size; // S x H
const size_t kv_input_chunk_length = static_cast<size_t>(kv_sequence_length) * head_size; // L x H
const size_t present_chunk_length = past_chunk_length + kv_input_chunk_length; // T x H
const size_t cache_chunk_length = static_cast<size_t>(max_sequence_length) * head_size; // M x H
DUMP_CPU_TENSOR_INIT();
DUMP_CPU_TENSOR("Q", Q, batch_size, num_heads_, sequence_length, head_size);
@ -164,7 +181,8 @@ class AttentionCPUBase : public AttentionBase {
}
if (present || present_key) {
double bytes_to_copy_key = static_cast<double>(sizeof(T) * present_chunk_length);
double bytes_to_copy_key = (past_present_share_buffer ? kv_input_chunk_length : present_chunk_length) *
static_cast<double>(sizeof(T));
unit_cost.bytes_loaded += bytes_to_copy_key;
unit_cost.bytes_stored += bytes_to_copy_key;
}
@ -214,7 +232,12 @@ class AttentionCPUBase : public AttentionBase {
// Concatenate past_K and K : (BxNx)PxH, (BxNx)LxH -> (BxNx)TxH
k = ConcatStateChunk(past, k, present, past_chunk_length, present_chunk_length, i);
} else if (nullptr != present_key) {
k = ConcatStateChunk(past_key, k, present_key, past_chunk_length, present_chunk_length, i);
if (past_present_share_buffer) {
k = present_key + cache_chunk_length * i;
memcpy(const_cast<T*>(k) + past_chunk_length, K + head_size * i, head_size * sizeof(T));
} else {
k = ConcatStateChunk(past_key, k, present_key, past_chunk_length, present_chunk_length, i);
}
}
// Compute Q*K' + AttentionMask
@ -230,6 +253,12 @@ class AttentionCPUBase : public AttentionBase {
});
}
if (output_qk_data != nullptr) {
// Output the scaled Q*K^T if needed.
memcpy(output_qk_data, attention_probs,
SafeInt<size_t>(batch_size) * num_heads_ * sequence_length * total_sequence_length * sizeof(T));
}
DUMP_CPU_TENSOR("QK (scaled)", attention_probs, batch_size, num_heads_, sequence_length, total_sequence_length);
// attention_probs(B, N, S, T) = Softmax(attention_probs)
@ -257,12 +286,15 @@ class AttentionCPUBase : public AttentionBase {
const T* past_value, // past value only (if not using past state)
T* present, // present state
T* present_value, // present value only (if not using present state)
ThreadPool* tp) const {
ThreadPool* tp,
bool past_present_share_buffer = false,
int max_sequence_length = 0) const {
const int total_sequence_length = past_sequence_length + kv_sequence_length; // T = P + L
const ptrdiff_t past_chunk_length = SafeInt<ptrdiff_t>(past_sequence_length) * v_head_size; // P x H_v
const ptrdiff_t q_input_chunk_length = SafeInt<ptrdiff_t>(sequence_length) * v_head_size; // S x H_v
const ptrdiff_t kv_input_chunk_length = SafeInt<ptrdiff_t>(kv_sequence_length) * v_head_size; // L x H_v
const ptrdiff_t present_chunk_length = past_chunk_length + kv_input_chunk_length; // T x H_v
const ptrdiff_t cache_chunk_length = SafeInt<ptrdiff_t>(max_sequence_length) * v_head_size; // M x H_v
// Move the pointer of past and present to start of v values.
if (nullptr != past) {
@ -281,7 +313,8 @@ class AttentionCPUBase : public AttentionBase {
unit_cost.bytes_stored = static_cast<double>(sequence_length * v_head_size * sizeof(T));
if (present || present_value) {
double bytes_to_copy_value = static_cast<double>(present_chunk_length * sizeof(T));
double bytes_to_copy_value = (past_present_share_buffer ? kv_input_chunk_length : present_chunk_length) *
static_cast<double>(sizeof(T));
unit_cost.bytes_loaded += bytes_to_copy_value;
unit_cost.bytes_stored += bytes_to_copy_value;
}
@ -299,7 +332,12 @@ class AttentionCPUBase : public AttentionBase {
// Concatenate past_V and V: (BxNx)PxH_v, (BxNx)LxH_v -> (BxNx)TxH_v
v = ConcatStateChunk(past, v, present, past_chunk_length, present_chunk_length, i);
} else if (nullptr != present_value) {
v = ConcatStateChunk(past_value, v, present_value, past_chunk_length, present_chunk_length, i);
if (past_present_share_buffer) {
v = present_value + cache_chunk_length * i;
memcpy(const_cast<T*>(v) + past_chunk_length, V + v_head_size * i, v_head_size * sizeof(T));
} else {
v = ConcatStateChunk(past_value, v, present_value, past_chunk_length, present_chunk_length, i);
}
}
T* current_tmp_data = reinterpret_cast<T*>(tmp_buffer) + q_input_chunk_length * i;

471
onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc Normal file
View file

@ -0,0 +1,471 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "attention_cpu_base.h"
#include "attention_utils.h"
#include "core/platform/env_var_utils.h"
#include "contrib_ops/cpu/bert/multihead_attention_helper.h"
#include "contrib_ops/cpu/bert/decoder_masked_multihead_attention.h"
using namespace ::onnxruntime::common;
using namespace ONNX_NAMESPACE;
namespace onnxruntime {
namespace contrib {
// TODO: refactor
static constexpr int kPastSequenceLengthInputIndex = 7;
static constexpr int kBeamWidthInputIndex = 8;
static constexpr int kCacheIndirectionInputIndex = 9;
static constexpr int kPastInputIndex = 5;
static constexpr int kPresentOutputIndex = 1;
static constexpr int kQKOutputIndex = 3;
static constexpr int kBiasIndex = 10;
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
DecoderMaskedMultiHeadAttention, \
kMSDomain, \
1, \
T, \
kCpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.MayInplace(kPastInputIndex, kPresentOutputIndex) \
.MayInplace(kPastInputIndex + 1, kPresentOutputIndex + 1) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.InputMemoryType(OrtMemTypeCPUInput, kPastSequenceLengthInputIndex) \
.InputMemoryType(OrtMemTypeCPUInput, kBeamWidthInputIndex), \
DecoderMaskedMultiHeadAttention<T>);
REGISTER_KERNEL_TYPED(float)
template <typename T>
DecoderMaskedMultiHeadAttention<T>::DecoderMaskedMultiHeadAttention(const OpKernelInfo& info)
: OpKernel(info), AttentionCPUBase(info, false) {
int64_t num_heads = 0;
ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0);
num_heads_ = static_cast<int>(num_heads);
mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
scale_ = info.GetAttrOrDefault<float>("scale", 0.0f);
past_present_share_buffer_ = info.GetAttrOrDefault<int64_t>("past_present_share_buffer", 0LL);
output_qk_ = info.GetAttrOrDefault<int64_t>("output_qk", 0LL);
}
template <typename T>
Status DecoderMaskedMultiHeadAttention<T>::Compute(OpKernelContext* context) const {
const Tensor* query = context->Input<Tensor>(0);
const Tensor* key = context->Input<Tensor>(1);
const Tensor* value = context->Input<Tensor>(2);
const Tensor* mask_index = context->Input<Tensor>(3);
const Tensor* attention_bias = context->Input<Tensor>(4);
const Tensor* past_key = context->Input<Tensor>(kPastInputIndex);
const Tensor* past_value = context->Input<Tensor>(kPastInputIndex + 1);
const Tensor* past_seq_len = context->Input<Tensor>(kPastSequenceLengthInputIndex);
const Tensor* beam_width = context->Input<Tensor>(kBeamWidthInputIndex);
const Tensor* cache_indir = context->Input<Tensor>(kCacheIndirectionInputIndex);
const Tensor* bias = context->Input<Tensor>(kBiasIndex);
DecoderMaskedMultiHeadAttentionParams parameters;
bool is_unidirectional = false;
ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs<Tensor>(query,
key,
value,
bias,
mask_index,
attention_bias,
past_key,
past_value,
past_seq_len,
&parameters,
num_heads_,
mask_filter_value_,
scale_,
is_unidirectional,
past_present_share_buffer_,
kDecoderMaskedMultiHeadAttention));
int batch_size = parameters.batch_size;
int sequence_length = parameters.sequence_length;
int head_size = parameters.head_size;
int v_head_size = parameters.v_head_size;
int hidden_size = parameters.hidden_size;
int v_hidden_size = parameters.v_hidden_size;
// This kernel is for decoding only (i.e.) sequence length has to be 1
if (sequence_length != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input sequence length should be 1 to use DecoderMaskedMultiHeadAttention. "
"Actual length is ",
sequence_length);
}
if (head_size != v_head_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
"QK head size should be same as V head size to use DecoderMaskedMultiHeadAttention");
}
if (parameters.mask_type != AttentionMaskType::MASK_2D_KEY_PADDING &&
parameters.mask_type != AttentionMaskType::MASK_NONE) {
return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
"DecoderMaskedMultiHeadAttention only supports no mask or 2D key "
"padding mask of shape [batch, total_seq_length] currently");
}
TensorShapeVector output_shape(3);
output_shape[0] = static_cast<int64_t>(batch_size);
output_shape[1] = static_cast<int64_t>(sequence_length);
output_shape[2] = static_cast<int64_t>(v_hidden_size);
Tensor* output = context->Output(0, output_shape);
std::vector<int64_t> present_dims{
parameters.batch_size, parameters.num_heads,
past_present_share_buffer_ ? parameters.max_sequence_length : parameters.total_sequence_length,
head_size};
TensorShape present_shape(present_dims);
Tensor* present_key = context->Output(kPresentOutputIndex, present_shape);
Tensor* present_value = context->Output(kPresentOutputIndex + 1, present_shape);
Tensor* output_qk = nullptr;
// Decoder cross-attention
if (past_key == nullptr && present_key == nullptr) {
if (attention_bias != nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
"DecoderMaskedMultiHeadAttention does not support attention bias for cross-attention");
}
parameters.is_cross_attention = true;
parameters.total_sequence_length = parameters.kv_sequence_length;
parameters.max_sequence_length = parameters.kv_sequence_length;
} else {
// Sanity check
ORT_ENFORCE(past_present_share_buffer_);
ORT_ENFORCE(past_key != nullptr && past_value != nullptr);
auto* present_key_data = present_key->MutableData<T>();
auto* present_value_data = present_value->MutableData<T>();
auto* past_key_data = past_key->Data<T>();
auto* past_value_data = past_value->Data<T>();
if (present_key_data != past_key_data) {
std::memcpy(present_key_data, past_key_data, past_key->SizeInBytes());
}
if (present_value_data != past_value_data) {
std::memcpy(present_value_data, past_value_data, past_value->SizeInBytes());
}
parameters.is_cross_attention = false;
}
if (output_qk_) {
int64_t qk_dims[] = {parameters.batch_size, parameters.num_heads, 1, parameters.total_sequence_length};
TensorShape qk_shape(&qk_dims[0], sizeof(qk_dims) / sizeof(qk_dims[0]));
output_qk = context->Output(kQKOutputIndex, qk_shape);
}
// Beam width (in case we are using this op inside BeamSearch)
int beam_width_value = 1;
if (beam_width != nullptr) {
beam_width_value = static_cast<int>(*beam_width->Data<int32_t>());
}
// Cache indirection (in case we are using this op inside BeamSearch)
if (beam_width_value > 1 && cache_indir == nullptr) {
// If beam width > 1, then cache indirection buffer MUST be present
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"If beam width is greater than 1, then cache indirection buffer MUST be present");
}
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));
OrtValue Q;
ORT_RETURN_IF_ERROR(MaybeTransposeToBNSHAndAddBias<T>(
context, allocator, batch_size, num_heads_, 1, head_size, query, bias, 0, Q));
// Cross-attention case
if (parameters.is_cross_attention) {
return ApplyAttention(Q.GetMutable<Tensor>()->MutableData<T>(),
key->Data<T>(),
value->Data<T>(),
mask_index, nullptr /* past */, past_key, past_value, output, present_key, present_value,
batch_size, 1 /* sequence_length */, parameters.kv_sequence_length,
head_size, v_head_size, v_hidden_size, attention_bias, context, output_qk);
}
OrtValue K, V;
ORT_RETURN_IF_ERROR(MaybeTransposeToBNSHAndAddBias<T>(
context, allocator, batch_size, num_heads_, 1, head_size, key, bias, hidden_size, K));
ORT_RETURN_IF_ERROR(MaybeTransposeToBNSHAndAddBias<T>(
context, allocator, batch_size, num_heads_, 1, v_head_size, value, bias, 2 * hidden_size, V));
// Self-attention, !has_beams
if (cache_indir == nullptr) {
return ApplyAttention(Q.GetMutable<Tensor>()->MutableData<T>(),
K.GetMutable<Tensor>()->MutableData<T>(),
V.GetMutable<Tensor>()->MutableData<T>(),
mask_index, nullptr /* past */, past_key, past_value, output, present_key, present_value,
batch_size, 1 /* sequence_length */, parameters.kv_sequence_length,
head_size, v_head_size, v_hidden_size, attention_bias, context, output_qk,
parameters.past_sequence_length, true /* past_present_share_buffer */);
}
// Self-attention, has_beams
return ApplyAttentionWithBeams(Q.GetMutable<Tensor>()->MutableData<T>(),
K.GetMutable<Tensor>()->MutableData<T>(),
V.GetMutable<Tensor>()->MutableData<T>(),
mask_index, past_key, past_value, output, present_key, present_value,
batch_size, parameters.past_sequence_length, parameters.max_sequence_length,
head_size, v_head_size, attention_bias, parameters.broadcast_attn_bias_dim_0,
parameters.broadcast_attn_bias_dim_1, cache_indir, context,
beam_width_value, output_qk);
}
template <typename T>
Status DecoderMaskedMultiHeadAttention<T>::ApplyAttentionWithBeams(
const T* Q,
const T* K,
const T* V,
const Tensor* mask_index,
const Tensor* past_key,
const Tensor* past_value,
Tensor* output,
Tensor* present_key,
Tensor* present_value,
int batch_size,
int past_sequence_length,
int max_sequence_length,
int head_size,
int v_head_size,
const Tensor* attn_bias,
bool broadcast_attn_bias_dim_0,
bool broadcast_attn_bias_dim_1,
const Tensor* cache_indir,
OpKernelContext* context,
int beam_width,
Tensor* output_qk) const {
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));
auto* tp = context->GetOperatorThreadPool();
int total_sequence_length = past_sequence_length + 1;
size_t bytes = SafeInt<size_t>(batch_size) * num_heads_ * total_sequence_length * sizeof(T);
auto attention_probs = allocator->Alloc(bytes);
BufferUniquePtr scratch_buffer(attention_probs, BufferDeleter(allocator));
T* output_qk_data = (output_qk != nullptr) ? output_qk->MutableData<T>() : nullptr;
const int32_t* mask_index_data = mask_index != nullptr ? mask_index->Data<int32_t>() : nullptr;
const T* attn_bias_data = attn_bias != nullptr ? attn_bias->Data<T>() : nullptr;
ComputeAttentionProbsWithBeams(static_cast<T*>(attention_probs), Q, K, mask_index_data, batch_size,
past_sequence_length, max_sequence_length, head_size, past_key->Data<T>(),
present_key->MutableData<T>(), tp, attn_bias_data, broadcast_attn_bias_dim_0,
broadcast_attn_bias_dim_1, cache_indir->Data<int32_t>(), beam_width, output_qk_data);
// Compute the attentionScore * Value: out_tmp(B, N, 1, H_v) = attention_probs(B, N, 1, T) x V(B, N, T, H_v)
auto out_tmp_data = allocator->Alloc(SafeInt<size_t>(batch_size) * num_heads_ * v_head_size * sizeof(T));
BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(std::move(allocator)));
ComputeVxAttentionScoreWithBeams(output->MutableData<T>(), static_cast<T*>(out_tmp_data),
static_cast<const T*>(attention_probs), V, batch_size,
past_sequence_length, max_sequence_length, v_head_size, past_value->Data<T>(),
present_value->MutableData<T>(), cache_indir->Data<int32_t>(), beam_width, tp);
return Status::OK();
}
template <typename T>
void DecoderMaskedMultiHeadAttention<T>::ComputeAttentionProbsWithBeams(
T* attention_probs,
const T* Q,
const T* K,
const int32_t* mask_index_data,
int batch_size,
int past_sequence_length,
int max_sequence_length,
int head_size,
const T* past_key_data,
T* present_key_data,
ThreadPool* tp,
const T* attn_bias_data,
bool broadcast_attn_bias_dim_0,
bool broadcast_attn_bias_dim_1,
const int32_t* cache_indir_data,
int beam_width,
T* output_qk_data) const {
float scale = scale_ == 0.0f ? 1.0f / sqrt(static_cast<float>(head_size)) : scale_;
TensorOpCost unit_cost;
auto total_sequence_length = past_sequence_length + 1;
const ptrdiff_t probs_matrix_size = total_sequence_length;
const ptrdiff_t probs_matrix_bytes = probs_matrix_size * sizeof(T);
unit_cost.compute_cycles = static_cast<double>((SafeInt<ptrdiff_t>(2) * head_size - 1) * total_sequence_length);
unit_cost.bytes_loaded = static_cast<double>(SafeInt<ptrdiff_t>(2) * head_size * total_sequence_length * sizeof(T));
unit_cost.bytes_stored = static_cast<double>(SafeInt<ptrdiff_t>(head_size) * total_sequence_length * sizeof(T));
if (attn_bias_data != nullptr) {
unit_cost.bytes_loaded += static_cast<double>(probs_matrix_bytes) * 2;
unit_cost.bytes_stored += probs_matrix_bytes;
}
if (mask_index_data != nullptr) {
unit_cost.bytes_stored += probs_matrix_bytes;
}
// Cost of appending current key to present key
unit_cost.compute_cycles += static_cast<double>(head_size);
unit_cost.bytes_loaded += static_cast<double>(head_size);
// Parallel for loop
const int loop_len = batch_size * num_heads_;
ThreadPool::TryParallelFor(tp, loop_len, unit_cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) {
for (std::ptrdiff_t i = begin; i != end; ++i) {
const std::ptrdiff_t batch_index = i / num_heads_;
const std::ptrdiff_t head_index = i % num_heads_;
const std::ptrdiff_t beam_batch_index = batch_index / beam_width;
const T* q_vec = Q + i * head_size;
const std::ptrdiff_t attn_bias_base_offset = ((broadcast_attn_bias_dim_0 ? 0 : (beam_batch_index * num_heads_)) +
(broadcast_attn_bias_dim_1 ? 0 : head_index)) *
probs_matrix_size;
{
// Calculate the latest position of the attention_probs
// (1, H) x (T, H)^T -> (1, T)
// Decompose into T (1, H) x (1, H)^T -> (1, 1) operations
auto last_offset = past_sequence_length + i * probs_matrix_size;
T* attention_probs_ptr = reinterpret_cast<T*>(attention_probs) + last_offset;
math::Dot<float, CPUMathUtil>(head_size, q_vec, K + i * head_size, attention_probs_ptr, nullptr);
// Apply the attention bias and mask
if (attn_bias_data != nullptr) {
*attention_probs_ptr += attn_bias_data[attn_bias_base_offset + past_sequence_length];
}
bool is_masked = (mask_index_data != nullptr) &&
(mask_index_data[(batch_index + 1) * total_sequence_length - 1] == 0);
if (is_masked) {
*attention_probs_ptr += mask_filter_value_;
}
*attention_probs_ptr *= scale;
}
{
// Calculate the rest of the attention_probs
for (std::ptrdiff_t j = 0; j < past_sequence_length; ++j) {
const int* beam_indices = &cache_indir_data[batch_index * max_sequence_length];
const std::ptrdiff_t beam_offset = static_cast<std::ptrdiff_t>(beam_indices[j]) * num_heads_ *
max_sequence_length * head_size;
const std::ptrdiff_t beam_batch_offset = (beam_batch_index * beam_width * num_heads_ + head_index) *
max_sequence_length * head_size;
const T* past_k_vec = past_key_data + beam_batch_offset + beam_offset + j * head_size;
T* output = reinterpret_cast<T*>(attention_probs) + j + i * probs_matrix_size;
math::Dot<float, CPUMathUtil>(head_size, q_vec, past_k_vec, output, nullptr);
// Apply the attention bias and mask
if (attn_bias_data != nullptr) {
*output += attn_bias_data[attn_bias_base_offset + j];
}
bool is_masked = (mask_index_data != nullptr) &&
(mask_index_data[batch_index * total_sequence_length + j] == 0);
if (is_masked) {
*output += mask_filter_value_;
}
*output *= scale;
}
}
// Append current key to present key (past_present_share_buffer_ is true)
memcpy(present_key_data + i * max_sequence_length * head_size, K + i * head_size, head_size * sizeof(T));
}
});
if (output_qk_data != nullptr) {
// Output the scaled Q*K^T if needed.
memcpy(output_qk_data, attention_probs,
SafeInt<size_t>(batch_size) * num_heads_ * total_sequence_length * sizeof(T));
}
// attention_probs(B, N, 1, T) = Softmax(attention_probs)
{
const int N = batch_size * num_heads_;
const int D = total_sequence_length;
ComputeAttentionSoftmaxInplace(attention_probs, N, D, tp);
}
}
template <typename T>
void DecoderMaskedMultiHeadAttention<T>::ComputeVxAttentionScoreWithBeams(
T* output,
T* tmp_buffer,
const T* attention_probs,
const T* V,
int batch_size,
int past_sequence_length,
int max_sequence_length,
int v_head_size,
const T* past_value_data,
T* present_value_data,
const int32_t* cache_indir_data,
int beam_width,
ThreadPool* tp) const {
const int total_sequence_length = past_sequence_length + 1;
TensorOpCost unit_cost;
unit_cost.compute_cycles = static_cast<double>(SafeInt<ptrdiff_t>(2) * v_head_size * total_sequence_length);
unit_cost.bytes_loaded = static_cast<double>(SafeInt<ptrdiff_t>(3) * v_head_size * total_sequence_length * sizeof(T));
unit_cost.bytes_stored = static_cast<double>(SafeInt<ptrdiff_t>(2) * v_head_size * total_sequence_length * sizeof(T));
// Cost of appending current value to present value
unit_cost.compute_cycles += static_cast<double>(v_head_size);
unit_cost.bytes_loaded += static_cast<double>(v_head_size);
ThreadPool::TryParallelFor(
tp, SafeInt<ptrdiff_t>(batch_size) * num_heads_, unit_cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) {
for (std::ptrdiff_t i = begin; i != end; ++i) {
const std::ptrdiff_t batch_index = i / num_heads_;
const std::ptrdiff_t head_index = i % num_heads_;
const std::ptrdiff_t beam_batch_index = batch_index / beam_width;
// Compute the attention score
// (1, T) x (T, H_v) -> (1, H_v)
// Decompose into T (1, 1) x (1, H_v) -> (1, H_v) operations and accumulate.
{
const T* attn_probs_ptr = attention_probs + (i + 1) * total_sequence_length - 1;
math::Scale<T, CPUMathUtil>(v_head_size,
static_cast<float>(*attn_probs_ptr),
V + i * v_head_size,
output + i * v_head_size,
nullptr);
}
{
for (std::ptrdiff_t j = 0; j < past_sequence_length; ++j) {
const int* beam_indices = &cache_indir_data[batch_index * max_sequence_length];
const std::ptrdiff_t beam_offset = static_cast<std::ptrdiff_t>(beam_indices[j]) * num_heads_ *
max_sequence_length * v_head_size;
const std::ptrdiff_t beam_batch_offset = (beam_batch_index * beam_width * num_heads_ + head_index) *
max_sequence_length * v_head_size;
const T* past_value_vec = past_value_data + beam_offset + beam_batch_offset;
const T* attn_probs_ptr = attention_probs + j + i * total_sequence_length;
math::Scale<T, CPUMathUtil>(v_head_size,
static_cast<float>(*attn_probs_ptr),
past_value_vec + j * v_head_size,
tmp_buffer + i * v_head_size,
nullptr);
math::Add<T, CPUMathUtil>(v_head_size,
output + i * v_head_size,
tmp_buffer + i * v_head_size,
output + i * v_head_size,
nullptr);
}
}
// Append current value to present value (past_present_share_buffer_ is true)
memcpy(present_value_data + i * max_sequence_length * v_head_size,
V + i * v_head_size,
v_head_size * sizeof(T));
}
});
}
} // namespace contrib
} // namespace onnxruntime

78
onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.h Normal file
View file

@ -0,0 +1,78 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
template <typename T>
class DecoderMaskedMultiHeadAttention final : public OpKernel, public AttentionCPUBase {
public:
DecoderMaskedMultiHeadAttention(const OpKernelInfo& info);
Status ApplyAttentionWithBeams(const T* Q,
const T* K,
const T* V,
const Tensor* mask_index,
const Tensor* past_key,
const Tensor* past_value,
Tensor* output,
Tensor* present_key,
Tensor* present_value,
int batch_size,
int past_sequence_length,
int max_sequence_length,
int head_size,
int v_head_size,
const Tensor* attn_bias,
bool broadcast_attn_bias_dim_0,
bool broadcast_attn_bias_dim_1,
const Tensor* cache_indir,
OpKernelContext* context,
int beam_width,
Tensor* scaled_qk = nullptr) const;
void ComputeAttentionProbsWithBeams(T* attention_probs,
const T* Q,
const T* K,
const int32_t* mask_index_data,
int batch_size,
int past_sequence_length,
int max_sequence_length,
int head_size,
const T* past_key,
T* present_key,
ThreadPool* tp,
const T* attn_bias_data,
bool broadcast_attn_bias_dim_0,
bool broadcast_attn_bias_dim_1,
const int32_t* cache_indir_data,
int beam_width,
T* scaled_qk_data = nullptr) const;
void ComputeVxAttentionScoreWithBeams(T* output,
T* tmp_buffer,
const T* attention_probs,
const T* V,
int batch_size,
int past_sequence_length,
int max_sequence_length,
int v_head_size,
const T* past_value,
T* present_value,
const int32_t* cache_indir_data,
int beam_width,
ThreadPool* tp) const;
Status Compute(OpKernelContext* context) const override;
protected:
int num_heads_; // number of attention heads
float mask_filter_value_;
float scale_;
bool past_present_share_buffer_;
bool output_qk_;
};
} // namespace contrib
} // namespace onnxruntime

2
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
View file

@ -61,6 +61,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastG
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, NGramRepeatBlock);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BifurcationDetector);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QuickGelu);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, DecoderMaskedMultiHeadAttention);
// ******** Start: Quantization ******************* //
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulInteger16);
@ -335,6 +336,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, NGramRepeatBlock)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BifurcationDetector)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QuickGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, DecoderMaskedMultiHeadAttention)>,
// These ops were experimental ops in onnx domain which have been removed now. We add them here as
// contrib ops to main backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, Affine)>,

39
onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h
View file

@ -10,45 +10,6 @@ namespace onnxruntime {
namespace contrib {
namespace cuda {
struct DecoderMaskedMultiHeadAttentionParams : AttentionParameters {
int beam_width = 1;
// Only NeoX style rotary embedding is supported
int rotary_embedding_dim = 0;
int t_step = 0;
// Weather to use multihead attention(excludes matmul and bias)
bool is_mha = false;
bool is_cross_attention = false;
bool is_packed_qkv = false;
// Useful to better use global memory bandwidth on certain CUDA architectures.
// Turned off by default for now until we fully understand performance implications
// for all types of workloads.
// Can be turned on by appropriate environment variable (see attention_common.h).
bool kv_data_in_flight = false;
void* q = nullptr;
void* q_bias = nullptr;
void* k = nullptr;
void* k_bias = nullptr;
void* v = nullptr;
void* v_bias = nullptr;
void* attention_bias = nullptr;
void* k_cache = nullptr;
void* v_cache = nullptr;
void* out = nullptr;
void* out_qk = nullptr;
const int32_t* cache_indir = nullptr;
const int32_t* mask = nullptr; // [B, total_sequence_length]
};
template <
// The type of the inputs. Supported types: float and half.
typename T,

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

@ -787,14 +787,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"M")
.Input(7,
"beam_width",
"The beam width that is being used while decoding."
"The beam width that is being used while decoding. "
"If not provided, the beam width will be assumed to be 1.",
"M",
OpSchema::Optional)
.Input(8,
"cache_indirection",
"A buffer of shape [batch_size, beam_width, max_output_length] where an [i, j, k] entry specifies"
"which beam the 'k' th token came from for the 'j' th beam for batch 'i' in the current iteration",
"A buffer of shape [batch_size, beam_width, max_output_length] where an `[i, j, k]` entry specifies "
"which beam the `k`-th token came from for the `j`-th beam for batch `i` in the current iteration",
"M",
OpSchema::Optional)
.Output(0,
@ -871,7 +871,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
OpSchema::Optional)
.Input(4,
"attention_bias",
"additional add to QxK' with shape (batch_size, num_heads, sequence_length, total_sequence_length)",
"additional add to QxK' with shape (batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)",
"T",
OpSchema::Optional)
.Input(5,
@ -902,15 +902,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
OpSchema::Optional)
.Input(8,
"beam_width",
"The beam width that is being used while decoding."
"The beam width that is being used while decoding. "
"If not provided, the beam width will be assumed to be 1.",
"M",
OpSchema::Optional)
.Input(9,
"cache_indirection",
// This input is useful for CUDA EP only.
"A buffer of shape [batch_size, beam_width, max_output_length] where an [i, j, k] entry specifies"
"which beam the 'k' th token came from for the 'j' th beam for batch 'i' in the current iteration",
"A buffer of shape [batch_size, beam_width, max_output_length] where an `[i, j, k]` entry specifies "
"which beam the `k`-th token came from for the `j`-th beam for batch `i` in the current iteration",
"M",
OpSchema::Optional)
.Input(10,
@ -940,7 +940,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
OpSchema::Optional)
.Output(3,
"qk",
"normalized Q * K, of shape (batch_size, num_heads, 1, head_size). ",
"normalized Q * K, of shape (batch_size, num_heads, 1, total_sequence_length). ",
"V",
OpSchema::Optional)
.TypeConstraint("V", {"tensor(float)"}, "Constrain qk output types to float32 tensors.")

177
onnxruntime/test/python/transformers/test_parity_t5_mha.py
View file

@ -159,6 +159,7 @@ def create_t5_decoder_masked_mha_graph(
head_size,
num_heads,
is_cross_attention,
beam_width=1,
):
nodes = [
helper.make_node(
@ -172,6 +173,8 @@ def create_t5_decoder_masked_mha_graph(
"past_key" if not is_cross_attention else "",
"past_value" if not is_cross_attention else "",
"past_sequence_length" if not is_cross_attention else "",
"beam_width" if beam_width > 1 else "",
"cache_indirection" if beam_width > 1 else "",
],
[
"output",
@ -233,6 +236,15 @@ def create_t5_decoder_masked_mha_graph(
)
)
graph_inputs.append(helper.make_tensor_value_info("past_sequence_length", TensorProto.INT32, [1]))
if beam_width > 1:
graph_inputs.append(helper.make_tensor_value_info("beam_width", TensorProto.INT32, [1]))
graph_inputs.append(
helper.make_tensor_value_info(
"cache_indirection", TensorProto.INT32, [batch_size, beam_width, past_sequence_length + 1]
)
)
graph_outputs.append(
helper.make_tensor_value_info(
"present_key", TensorProto.FLOAT, [batch_size, num_heads, past_sequence_length + 1, head_size]
@ -275,7 +287,7 @@ class T5Config:
class T5Attention(nn.Module):
def __init__(self, config: T5Config, is_static_kv, use_decoder_masked_kernel: bool = False):
def __init__(self, config: T5Config, is_static_kv, use_decoder_masked_kernel: bool = False, beam_width=1):
super().__init__()
self.is_decoder = config.is_decoder
self.is_static_kv = is_static_kv
@ -284,6 +296,7 @@ class T5Attention(nn.Module):
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
self.beam_width = beam_width
# Mesh TensorFlow initialization to avoid scaling before softmax
self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)
@ -312,6 +325,7 @@ class T5Attention(nn.Module):
self.head_size,
self.num_heads,
is_static_kv,
beam_width=self.beam_width,
)
else:
self.onnx_graph = create_t5_mha_graph(
@ -371,7 +385,17 @@ class T5Attention(nn.Module):
position_bias = torch.normal(mean=5, std=0.1, size=(1, self.num_heads, 1, position_bias_length)).to(
torch.float32
)
return hidden_states, key_value_states, past_key_value, attention_mask, position_bias
inputs = [hidden_states, key_value_states, past_key_value, attention_mask, position_bias]
if self.beam_width > 1:
# Treat total_sequence_length as max_sequence_length here.
max_sequence_length = self.kv_sequence_length + self.seq_len
# Use random generated values here, which may not be valid in real case.
cache_indirection = torch.randint(
0, self.beam_width, (self.batch_size, self.beam_width, max_sequence_length)
).to(torch.int32)
inputs.append(cache_indirection)
return inputs
def torch_forward(
self,
@ -497,16 +521,17 @@ class T5Attention(nn.Module):
past_key_value=None,
mask=None,
position_bias=None,
use_cache=False,
use_cuda=True,
query_length=None,
cache_indirection=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)
execution_providers = ["CUDAExecutionProvider"] if use_cuda else ["CPUExecutionProvider"]
if execution_providers[0] not in onnxruntime.get_available_providers():
return
ort_session = onnxruntime.InferenceSession(self.onnx_graph, sess_options, providers=execution_providers)
batch_size, seq_length = hidden_states.shape[:2]
@ -578,13 +603,17 @@ class T5Attention(nn.Module):
torch_past_value_padded[:, :, : torch_past_value.shape[2], :] = torch_past_value
if self.is_static_kv:
if self.use_decoder_masked_kernel:
reordered_past_key = self.reorder_key_cache(
torch_past_key.flatten(),
batch_size=batch_size,
num_heads=self.num_heads,
sequence_length=self.kv_sequence_length,
head_size=self.head_size,
max_sequence_length=self.kv_sequence_length,
reordered_past_key = (
self.reorder_key_cache(
torch_past_key.flatten(),
batch_size=batch_size,
num_heads=self.num_heads,
sequence_length=self.kv_sequence_length,
head_size=self.head_size,
max_sequence_length=self.kv_sequence_length,
)
if use_cuda
else torch_past_key
)
ort_inputs["key"] = reordered_past_key.reshape(torch_past_key.shape)
ort_inputs["value"] = torch_past_value
@ -595,13 +624,17 @@ class T5Attention(nn.Module):
ort_inputs["key"] = np.ascontiguousarray(key_states.detach().numpy())
ort_inputs["value"] = np.ascontiguousarray(value_states.detach().numpy())
if self.use_decoder_masked_kernel:
reordered_past_key = self.reorder_key_cache(
torch_past_key_padded.flatten(),
batch_size=batch_size,
num_heads=self.num_heads,
sequence_length=self.kv_sequence_length,
head_size=self.head_size,
max_sequence_length=max_seq_len,
reordered_past_key = (
self.reorder_key_cache(
torch_past_key_padded.flatten(),
batch_size=batch_size,
num_heads=self.num_heads,
sequence_length=self.kv_sequence_length,
head_size=self.head_size,
max_sequence_length=max_seq_len,
)
if use_cuda
else torch_past_key_padded
)
ort_inputs["past_key"] = reordered_past_key.reshape(torch_past_value_padded.shape)
ort_inputs["past_value"] = torch_past_value_padded
@ -617,6 +650,10 @@ class T5Attention(nn.Module):
if torch_position_bias is not None:
ort_inputs["attention_bias"] = np.ascontiguousarray(torch_position_bias.detach().numpy())
if self.beam_width > 1:
ort_inputs["beam_width"] = np.ascontiguousarray(np.array([self.beam_width], dtype=np.int32))
ort_inputs["cache_indirection"] = np.ascontiguousarray(cache_indirection.detach().numpy())
ort_output = ort_session.run(None, ort_inputs)
output = None
@ -628,7 +665,9 @@ class T5Attention(nn.Module):
return output
def compare_t5_cross_attention_decoder(batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=False):
def compare_t5_cross_attention_decoder(
batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=False, use_cuda=True
):
config = T5Config(
is_decoder=True,
batch_size=batch_size,
@ -646,7 +685,7 @@ def compare_t5_cross_attention_decoder(batch_size, seq_len, num_heads, head_size
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
hidden_states, key_value_states, past_key_value, attention_mask, position_bias=None, use_cuda=use_cuda
)
if ort_output is not None:
@ -669,9 +708,7 @@ def compare_t5_cross_attention_decoder_init(batch_size, seq_len, num_heads, head
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
)
ort_output = T5CrossAttention.ort_forward(hidden_states, key_value_states, None, attention_mask, position_bias=None)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
@ -695,9 +732,7 @@ def compare_t5_self_attention_decoder_init(batch_size, seq_len, num_heads, head_
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
)
ort_output = T5CrossAttention.ort_forward(hidden_states, None, None, mask=None, position_bias=position_bias)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
@ -705,7 +740,9 @@ def compare_t5_self_attention_decoder_init(batch_size, seq_len, num_heads, head_
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, use_dmmha=False):
def compare_t5_self_attention_decoder(
batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=False, use_cuda=True, beam_width=1
):
config = T5Config(
is_decoder=True,
batch_size=batch_size,
@ -716,21 +753,45 @@ def compare_t5_self_attention_decoder(batch_size, seq_len, num_heads, head_size,
use_past=True,
)
T5CrossAttention = T5Attention(config, is_static_kv=False, use_decoder_masked_kernel=use_dmmha) # noqa: N806
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
T5CrossAttention = T5Attention( # noqa: N806
config, is_static_kv=False, use_decoder_masked_kernel=use_dmmha, beam_width=beam_width
)
hidden_states, _, past_key_value, _, position_bias, *maybe_cache_indirection = T5CrossAttention.create_inputs()
cache_indirection = maybe_cache_indirection[0] if beam_width > 1 else None
if beam_width > 1:
# When beam_width > 1, use ORT CUDA result as reference
ref_output = T5CrossAttention.ort_forward(
hidden_states,
None,
past_key_value,
mask=None,
position_bias=position_bias,
cache_indirection=cache_indirection,
use_cuda=True,
)
if ref_output is None:
# Return directly if CUDA EP is not available
return
else:
ref_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
hidden_states,
None,
past_key_value,
mask=None,
position_bias=position_bias,
cache_indirection=cache_indirection,
use_cuda=use_cuda,
)
if ort_output is not None:
assert torch.allclose(torch_output[0], ort_output[0], atol=1e-4)
assert torch.allclose(ref_output[0], ort_output[0], atol=1e-4)
if not use_dmmha:
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)
assert torch.allclose(ref_output[1][0], ort_output[1][0], atol=1e-4)
assert torch.allclose(ref_output[1][1], ort_output[1][1], atol=1e-4)
class TestT5MHAParity(unittest.TestCase):
@ -761,23 +822,53 @@ class TestT5MHAParity(unittest.TestCase):
self.batch_size, self.seq_len, self.num_heads, self.head_size, self.kv_sequence_length
)
def test_t5_cross_attention_decoder_masked_mha(self):
def test_t5_cross_attention_decoder_masked_mha(self, use_cuda=True):
batch_size = 2
seq_len = 1
num_heads = 2
head_size = 32
kv_sequence_length = 2
compare_t5_cross_attention_decoder(
batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=True
batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=True, use_cuda=use_cuda
)
def test_t5_self_attention_decoder_masked_mha(self):
def test_t5_self_attention_decoder_masked_mha(self, use_cuda=True):
batch_size = 2
seq_len = 1
num_heads = 2
head_size = 32
kv_sequence_length = 2
compare_t5_self_attention_decoder(batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=True)
compare_t5_self_attention_decoder(
batch_size, seq_len, num_heads, head_size, kv_sequence_length, use_dmmha=True, use_cuda=use_cuda
)
def test_t5_cross_attention_decoder_masked_mha_cpu(self):
return self.test_t5_cross_attention_decoder_masked_mha(use_cuda=False)
def test_t5_self_attention_decoder_masked_mha_cpu(self):
return self.test_t5_self_attention_decoder_masked_mha(use_cuda=False)
def test_t5_self_attention_decoder_masked_mha_with_beams(self):
"""
Test DecoderMaskedMultiHeadAttention self-attention case with beam_width > 1.
Compare the results on CUDA and CPU EPs.
"""
batch_size = 4
seq_len = 1
num_heads = 2
head_size = 32
kv_sequence_length = 2
beam_width = 2
compare_t5_self_attention_decoder(
batch_size,
seq_len,
num_heads,
head_size,
kv_sequence_length,
use_dmmha=True,
use_cuda=False,
beam_width=beam_width,
)
if __name__ == "__main__":