From 1fa219d7d517314ad51bc211702539714ddcbf8a Mon Sep 17 00:00:00 2001
From: mindest <30493312+mindest@users.noreply.github.com>
Date: Sun, 13 Oct 2024 05:43:00 +0900
Subject: [PATCH] DecoderMaskedMultiHeadAttention CPU kernel. (#22292)
### Description
DecoderMaskedMultiHeadAttention CPU kernel.
---
docs/ContribOperators.md | 12 +-
docs/OperatorKernels.md | 1 +
.../contrib_ops/cpu/bert/attention_common.h | 39 ++
.../contrib_ops/cpu/bert/attention_cpu_base.h | 104 ++--
.../decoder_masked_multihead_attention.cc | 471 ++++++++++++++++++
.../bert/decoder_masked_multihead_attention.h | 78 +++
.../contrib_ops/cpu/cpu_contrib_kernels.cc | 2 +
.../decoder_masked_multihead_attention_impl.h | 39 --
.../core/graph/contrib_ops/bert_defs.cc | 16 +-
.../python/transformers/test_parity_t5_mha.py | 177 +++++--
10 files changed, 810 insertions(+), 129 deletions(-)
create mode 100644 onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc
create mode 100644 onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.h
diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index c791c42617..b87532debe 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -1167,7 +1167,7 @@ This version of the operator has been available since version 1 of the 'com.micr
mask_index (optional) : M
Mask values of shape (batch_size, total_sequence_length) or (batch_size, kv_sequence_length)
attention_bias (optional) : T
-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)
past_key (optional) : T
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)`.
past_value (optional) : T
@@ -1175,9 +1175,9 @@ This version of the operator has been available since version 1 of the 'com.micr
past_sequence_length (optional) : M
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.
beam_width (optional) : M
-The beam width that is being used while decoding.If not provided, the beam width will be assumed to be 1.
+The beam width that is being used while decoding. If not provided, the beam width will be assumed to be 1.
cache_indirection (optional) : M
-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
+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
bias (optional) : T
Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection
@@ -1192,7 +1192,7 @@ This version of the operator has been available since version 1 of the 'com.micr
present_value (optional) : T
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).
qk (optional) : V
-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).
#### Type Constraints
@@ -1261,9 +1261,9 @@ This version of the operator has been available since version 1 of the 'com.micr
past_sequence_length : M
When past_present_share_buffer is used, it is required to specify past_sequence_length (could be 0).
beam_width (optional) : M
-The beam width that is being used while decoding.If not provided, the beam width will be assumed to be 1.
+The beam width that is being used while decoding. If not provided, the beam width will be assumed to be 1.
cache_indirection (optional) : M
-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
+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
#### Outputs
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 5963f1a03b..d8de7756ba 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -468,6 +468,7 @@ Do not modify directly.*
|CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)|
|ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|CropAndResize|*in* X:**T1**
*in* rois:**T1**
*in* batch_indices:**T2**
*in* crop_size:**T2**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int32)|
+|DecoderMaskedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* mask_index:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* beam_width:**M**
*in* cache_indirection:**M**
*in* bias:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**V**|1+|**T** = tensor(float)|
|DequantizeLinear|*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**|1+|**T1** = tensor(int16), tensor(int32), tensor(int4), tensor(int8), tensor(uint16), tensor(uint4), tensor(uint8)
**T2** = tensor(float)|
|DynamicQuantizeLSTM|*in* X:**T**
*in* W:**T2**
*in* R:**T2**
*in* B:**T**
*in* sequence_lens:**T1**
*in* initial_h:**T**
*in* initial_c:**T**
*in* P:**T**
*in* W_scale:**T**
*in* W_zero_point:**T2**
*in* R_scale:**T**
*in* R_zero_point:**T2**
*out* Y:**T**
*out* Y_h:**T**
*out* Y_c:**T**|1+|**T** = tensor(float)
**T1** = tensor(int32)
**T2** = tensor(int8), tensor(uint8)|
|DynamicQuantizeMatMul|*in* A:**T1**
*in* B:**T2**
*in* b_scale:**T1**
*in* b_zero_point:**T2**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)|
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_common.h b/onnxruntime/contrib_ops/cpu/bert/attention_common.h
index e0fa581c80..4663855557 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_common.h
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_common.h
@@ -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;
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
index ae2eaf0204..dc9ba80af5 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
@@ -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(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(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(mask_data),
causal, batch_size, sequence_length, past_sequence_length, mask_filter_value_);
DUMP_CPU_TENSOR("Mask3D", static_cast(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() : nullptr;
const T* past_value_data = past_value != nullptr ? past_value->Data() : nullptr;
T* present_value_data = present_value != nullptr ? present_value->MutableData() : nullptr;
+ T* output_qk_data = output_qk != nullptr ? output_qk->MutableData() : nullptr;
const T* attn_bias_data = (attn_bias != nullptr) ? attn_bias->Data() : nullptr;
auto attn_bias_dims = (attn_bias != nullptr) ? attn_bias->Shape().GetDims() : gsl::span{};
+ // 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(past_key->Shape().GetDims()[2]);
+ }
+
// Compute the attention score.
size_t bytes = SafeInt(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(static_cast(attention_probs), Q, K,
static_cast(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(), static_cast(out_tmp_data), static_cast(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
- 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 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 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(past_sequence_length) * head_size; // P x H
const size_t q_input_chunk_length = static_cast(sequence_length) * head_size; // S x H
const size_t kv_input_chunk_length = static_cast(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(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(sizeof(T) * present_chunk_length);
+ double bytes_to_copy_key = (past_present_share_buffer ? kv_input_chunk_length : present_chunk_length) *
+ static_cast(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(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(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(past_sequence_length) * v_head_size; // P x H_v
const ptrdiff_t q_input_chunk_length = SafeInt(sequence_length) * v_head_size; // S x H_v
const ptrdiff_t kv_input_chunk_length = SafeInt(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(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(sequence_length * v_head_size * sizeof(T));
if (present || present_value) {
- double bytes_to_copy_value = static_cast(present_chunk_length * sizeof(T));
+ double bytes_to_copy_value = (past_present_share_buffer ? kv_input_chunk_length : present_chunk_length) *
+ static_cast(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(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(tmp_buffer) + q_input_chunk_length * i;
diff --git a/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc b/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc
new file mode 100644
index 0000000000..b2aaa9cb11
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc
@@ -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()) \
+ .InputMemoryType(OrtMemTypeCPUInput, kPastSequenceLengthInputIndex) \
+ .InputMemoryType(OrtMemTypeCPUInput, kBeamWidthInputIndex), \
+ DecoderMaskedMultiHeadAttention);
+
+REGISTER_KERNEL_TYPED(float)
+
+template
+DecoderMaskedMultiHeadAttention::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(num_heads);
+ mask_filter_value_ = info.GetAttrOrDefault("mask_filter_value", -10000.0f);
+ scale_ = info.GetAttrOrDefault("scale", 0.0f);
+ past_present_share_buffer_ = info.GetAttrOrDefault("past_present_share_buffer", 0LL);
+ output_qk_ = info.GetAttrOrDefault("output_qk", 0LL);
+}
+
+template
+Status DecoderMaskedMultiHeadAttention::Compute(OpKernelContext* context) const {
+ const Tensor* query = context->Input(0);
+ const Tensor* key = context->Input(1);
+ const Tensor* value = context->Input(2);
+ const Tensor* mask_index = context->Input(3);
+ const Tensor* attention_bias = context->Input(4);
+ const Tensor* past_key = context->Input(kPastInputIndex);
+ const Tensor* past_value = context->Input(kPastInputIndex + 1);
+ const Tensor* past_seq_len = context->Input(kPastSequenceLengthInputIndex);
+ const Tensor* beam_width = context->Input(kBeamWidthInputIndex);
+ const Tensor* cache_indir = context->Input(kCacheIndirectionInputIndex);
+ const Tensor* bias = context->Input(kBiasIndex);
+
+ DecoderMaskedMultiHeadAttentionParams parameters;
+
+ bool is_unidirectional = false;
+ ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs(query,
+ key,
+ value,
+ bias,
+ mask_index,
+ attention_bias,
+ past_key,
+ past_value,
+ past_seq_len,
+ ¶meters,
+ 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(batch_size);
+ output_shape[1] = static_cast(sequence_length);
+ output_shape[2] = static_cast(v_hidden_size);
+ Tensor* output = context->Output(0, output_shape);
+
+ std::vector 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();
+ auto* present_value_data = present_value->MutableData();
+ auto* past_key_data = past_key->Data();
+ auto* past_value_data = past_value->Data();
+
+ 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(*beam_width->Data());
+ }
+
+ // 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(
+ 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()->MutableData(),
+ key->Data(),
+ value->Data(),
+ 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(
+ context, allocator, batch_size, num_heads_, 1, head_size, key, bias, hidden_size, K));
+ ORT_RETURN_IF_ERROR(MaybeTransposeToBNSHAndAddBias(
+ 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()->MutableData(),
+ K.GetMutable()->MutableData(),
+ V.GetMutable()->MutableData(),
+ 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()->MutableData(),
+ K.GetMutable()->MutableData(),
+ V.GetMutable()->MutableData(),
+ 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
+Status DecoderMaskedMultiHeadAttention::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(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() : nullptr;
+
+ const int32_t* mask_index_data = mask_index != nullptr ? mask_index->Data() : nullptr;
+ const T* attn_bias_data = attn_bias != nullptr ? attn_bias->Data() : nullptr;
+
+ ComputeAttentionProbsWithBeams(static_cast(attention_probs), Q, K, mask_index_data, batch_size,
+ past_sequence_length, max_sequence_length, head_size, past_key->Data(),
+ present_key->MutableData(), tp, attn_bias_data, broadcast_attn_bias_dim_0,
+ broadcast_attn_bias_dim_1, cache_indir->Data(), 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(batch_size) * num_heads_ * v_head_size * sizeof(T));
+ BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(std::move(allocator)));
+
+ ComputeVxAttentionScoreWithBeams(output->MutableData(), static_cast(out_tmp_data),
+ static_cast(attention_probs), V, batch_size,
+ past_sequence_length, max_sequence_length, v_head_size, past_value->Data(),
+ present_value->MutableData(), cache_indir->Data(), beam_width, tp);
+
+ return Status::OK();
+}
+
+template
+void DecoderMaskedMultiHeadAttention::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(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((SafeInt(2) * head_size - 1) * total_sequence_length);
+ unit_cost.bytes_loaded = static_cast(SafeInt(2) * head_size * total_sequence_length * sizeof(T));
+ unit_cost.bytes_stored = static_cast(SafeInt(head_size) * total_sequence_length * sizeof(T));
+
+ if (attn_bias_data != nullptr) {
+ unit_cost.bytes_loaded += static_cast(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(head_size);
+ unit_cost.bytes_loaded += static_cast(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(attention_probs) + last_offset;
+ math::Dot(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(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(attention_probs) + j + i * probs_matrix_size;
+ math::Dot(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(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
+void DecoderMaskedMultiHeadAttention::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(SafeInt(2) * v_head_size * total_sequence_length);
+ unit_cost.bytes_loaded = static_cast(SafeInt(3) * v_head_size * total_sequence_length * sizeof(T));
+ unit_cost.bytes_stored = static_cast(SafeInt(2) * v_head_size * total_sequence_length * sizeof(T));
+
+ // Cost of appending current value to present value
+ unit_cost.compute_cycles += static_cast(v_head_size);
+ unit_cost.bytes_loaded += static_cast(v_head_size);
+
+ ThreadPool::TryParallelFor(
+ tp, SafeInt(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(v_head_size,
+ static_cast(*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(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(v_head_size,
+ static_cast(*attn_probs_ptr),
+ past_value_vec + j * v_head_size,
+ tmp_buffer + i * v_head_size,
+ nullptr);
+ math::Add(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
diff --git a/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.h b/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.h
new file mode 100644
index 0000000000..68d1b97513
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.h
@@ -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
+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
diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
index 5f59195667..c742cd1e95 100644
--- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
@@ -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,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
// 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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h
index efad338553..0e1c9ce7b1 100644
--- a/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h
+++ b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h
@@ -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,
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index c706c6fc5f..0a261d8f73 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -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.")
diff --git a/onnxruntime/test/python/transformers/test_parity_t5_mha.py b/onnxruntime/test/python/transformers/test_parity_t5_mha.py
index e4f65b07c5..84708ddcf8 100644
--- a/onnxruntime/test/python/transformers/test_parity_t5_mha.py
+++ b/onnxruntime/test/python/transformers/test_parity_t5_mha.py
@@ -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__":