diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.cc b/onnxruntime/contrib_ops/cpu/bert/attention.cc index 7698832874..760e86a1a9 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/attention.cc @@ -44,7 +44,7 @@ Status AttentionBase::CheckInputs(const Tensor* input, const auto& dims = input->Shape().GetDims(); if (dims.size() != 3) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 0 is expected to have 3 dimensions, got ", + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'input' is expected to have 3 dimensions, got ", dims.size()); } int batch_size = static_cast(dims[0]); @@ -56,7 +56,7 @@ Status AttentionBase::CheckInputs(const Tensor* input, const auto& weights_dims = weights->Shape().GetDims(); if (weights_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 1 is expected to have 2 dimensions, got ", + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' is expected to have 2 dimensions, got ", weights_dims.size()); } if (weights_dims[0] != dims[2]) { @@ -64,12 +64,12 @@ Status AttentionBase::CheckInputs(const Tensor* input, "Input 1 dimension 0 should have same length as dimension 2 of input 0"); } if (weights_dims[1] != 3 * weights_dims[0]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 1 dimension 1 should be 3 times of dimension 0"); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' dimension 1 should be 3 times of dimension 0"); } const auto& bias_dims = bias->Shape().GetDims(); if (bias_dims.size() != 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 2 is expected to have 1 dimension, got ", + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' is expected to have 1 dimension, got ", bias_dims.size()); } if (bias_dims[0] != weights_dims[1]) { @@ -80,40 +80,40 @@ Status AttentionBase::CheckInputs(const Tensor* input, if (mask_index != nullptr) { // mask_index is optional // unidirectional (like GPT2) does not need mask input. Here we do not allowed the input for unidirectional. if (is_unidirectional_) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 3 (mask_index) is not allowed for unidirectional"); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask_index' is not allowed for unidirectional"); } const auto& mask_dims = mask_index->Shape().GetDims(); if (mask_dims.size() != 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 3 is expected to have 1 dimension, got ", + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask_index' is expected to have 1 dimension, got ", mask_dims.size()); } if (static_cast(mask_dims[0]) != batch_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 3 and 0 shall have same length at dimension 0"); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' and 'input' shall have same length at dimension 0"); } } if (past != nullptr) { // past is optional if (!is_unidirectional_) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 (past) is only allowed for unidirectional"); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past' is only allowed for unidirectional"); } const auto& past_dims = past->Shape().GetDims(); if (past_dims.size() != 5) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 is expected to have 5 dimension, got ", + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past' is expected to have 5 dimension, got ", past_dims.size()); } if (static_cast(past_dims[0]) != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 0 shall have length of 2"); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' dimension 0 shall have length of 2"); } if (static_cast(past_dims[1]) != batch_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 1 shall have same length as dimension 0 of input 0"); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' dimension 1 shall have same length as dimension 0 of input 0"); } if (static_cast(past_dims[2]) != num_heads_) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 2 shall have length of num_heads", num_heads_); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' dimension 2 shall have length of num_heads", num_heads_); } if (static_cast(past_dims[4]) != hidden_size / num_heads_) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 2 shall have length of ", hidden_size / num_heads_); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' dimension 2 shall have length of ", hidden_size / num_heads_); } } @@ -147,7 +147,7 @@ Tensor* AttentionBase::GetPresent(OpKernelContext* context, } template -Attention::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionBase(info) { +Attention::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) { } template @@ -168,19 +168,15 @@ Status Attention::Compute(OpKernelContext* context) const { Tensor* output = context->Output(0, shape); - int past_sequence_length = 0; - Tensor* present = GetPresent(context, past, batch_size, head_size, sequence_length, past_sequence_length); - - // Total sequence length including that of past state: S* = S' + S - const int all_sequence_length = past_sequence_length + sequence_length; - constexpr size_t element_size = sizeof(T); AllocatorPtr allocator; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); auto* tp = context->GetOperatorThreadPool(); - // STEP.1: gemm_data(BS, 3NH) = input(BS, NH) x weights(NH, 3NH) + bias(3NH) + + // Compute Q, K, V + // gemm_data(BS, 3NH) = input(BS, NH) x weights(NH, 3NH) + bias(3NH) auto gemm_data = allocator->Alloc(SafeInt(batch_size) * sequence_length * 3 * hidden_size * element_size); BufferUniquePtr gemm_buffer(gemm_data, BufferDeleter(allocator)); auto Q = reinterpret_cast(gemm_data); @@ -235,51 +231,15 @@ Status Attention::Compute(OpKernelContext* context) const { qkv_dest + qkv_offset, // C head_size, // ldc nullptr // use single-thread - ); + ); } }); } - // STEP.2: compute the attention score. It does 2 things: - // I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) + - // 1 x mask_data(B, N, S, S*) - // II.attention_probs(B, N, S, S*) = Softmax(attention_probs) - size_t attention_probs_bytes = SafeInt(batch_size) * num_heads_ * sequence_length * all_sequence_length * element_size; - auto attention_probs = allocator->Alloc(attention_probs_bytes); - BufferUniquePtr scratch_buffer(attention_probs, BufferDeleter(allocator)); - - size_t mask_data_bytes = 0; - if (mask_index != nullptr) { - mask_data_bytes = SafeInt(batch_size) * sequence_length * all_sequence_length * element_size; - } else if (is_unidirectional_) { - mask_data_bytes = SafeInt(sequence_length) * all_sequence_length * element_size; - } - - void* mask_data = nullptr; - if (mask_data_bytes > 0) { - mask_data = allocator->Alloc(mask_data_bytes); - memset(mask_data, 0, mask_data_bytes); - } - BufferUniquePtr mask_data_buffer(mask_data, BufferDeleter(allocator)); - - const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data() : nullptr; - const T* past_data = past != nullptr ? past->template Data() : nullptr; - T* present_data = present != nullptr ? present->template MutableData() : nullptr; - - ComputeAttentionProbs(static_cast(attention_probs), Q, K, mask_index_data, static_cast(mask_data), - batch_size, sequence_length, past_sequence_length, head_size, num_heads_, is_unidirectional_, - past_data, present_data, tp); - - // STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S*) x V(B, N, S*, H) - auto out_tmp_data = - allocator->Alloc(SafeInt(batch_size) * num_heads_ * sequence_length * head_size * element_size); - BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator)); - - ComputeVxAttentionScore(output->template MutableData(), static_cast(out_tmp_data), static_cast(attention_probs), V, - batch_size, sequence_length, past_sequence_length, head_size, num_heads_, hidden_size, - past_data, present_data, tp); - - return Status::OK(); + // Compute the attention score and apply the score to V + return ApplyAttention(Q, K, V, mask_index, past, output, + batch_size, sequence_length, + head_size, hidden_size, context); } } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.h b/onnxruntime/contrib_ops/cpu/bert/attention.h index fa686284ac..bd82f07f81 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention.h @@ -3,34 +3,15 @@ #pragma once +#include "attention_cpu_base.h" #include "core/common/common.h" #include "core/framework/op_kernel.h" namespace onnxruntime { namespace contrib { -class AttentionBase { - protected: - AttentionBase(const OpKernelInfo& info); - Status CheckInputs(const Tensor* input, - const Tensor* weights, - const Tensor* bias, - const Tensor* mask_index, - const Tensor* past) const; - - Tensor* GetPresent(OpKernelContext* context, - const Tensor* past, - int batch_size, - int head_size, - int sequence_length, - int& past_sequence_length) const; - - int num_heads_; // number of attention heads - bool is_unidirectional_; // whether every token can only attend to previous tokens. -}; - template -class Attention : public OpKernel, public AttentionBase { +class Attention : public OpKernel, public AttentionCPUBase { public: explicit Attention(const OpKernelInfo& info); Status Compute(OpKernelContext* context) const override; diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_base.h new file mode 100644 index 0000000000..b9c1934236 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.h @@ -0,0 +1,33 @@ +// 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 { + +class AttentionBase { + protected: + AttentionBase(const OpKernelInfo& info); + Status CheckInputs(const Tensor* input, + const Tensor* weights, + const Tensor* bias, + const Tensor* mask_index, + const Tensor* past) const; + + Tensor* GetPresent(OpKernelContext* context, + const Tensor* past, + int batch_size, + int head_size, + int sequence_length, + int& past_sequence_length) const; + + int num_heads_; // number of attention heads + bool is_unidirectional_; // whether every token can only attend to previous tokens. +}; + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h new file mode 100644 index 0000000000..1f01504ded --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h @@ -0,0 +1,218 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "attention_base.h" +#include "attention_helper.h" + +#include "core/common/common.h" +#include "core/common/safeint.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { + +class AttentionCPUBase : public AttentionBase { + protected: + AttentionCPUBase(const OpKernelInfo& info) : AttentionBase(info) {} + + template + Status ApplyAttention(const T* Q, // Q data. Its size is BxNxSxH + const T* K, // K data. Its size is BxNxSxH + const T* V, // V value with size BxNxSxH + const Tensor* mask_index, // mask index. nullptr if no mask or its size is B + const Tensor* past, // past state + Tensor* output, // output tensor + int batch_size, // batch size + int sequence_length, // sequence length + int head_size, // head size + int hidden_size, // hidden size + OpKernelContext* context) const { + AllocatorPtr allocator; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); + + auto* tp = context->GetOperatorThreadPool(); + + int past_sequence_length = 0; + Tensor* present = GetPresent(context, past, batch_size, head_size, sequence_length, past_sequence_length); + + // Total sequence length including that of past state: S* = S' + S + const int all_sequence_length = past_sequence_length + sequence_length; + + // Compute the attention score. It does 2 things: + // I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) + + // 1 x mask_data(B, N, S, S*) + // II.attention_probs(B, N, S, S*) = Softmax(attention_probs) + size_t attention_probs_bytes = SafeInt(batch_size) * num_heads_ * sequence_length * all_sequence_length * sizeof(T); + auto attention_probs = allocator->Alloc(attention_probs_bytes); + BufferUniquePtr scratch_buffer(attention_probs, BufferDeleter(allocator)); + + size_t mask_data_bytes = 0; + if (mask_index != nullptr) { + mask_data_bytes = SafeInt(batch_size) * sequence_length * all_sequence_length * sizeof(T); + } else if (is_unidirectional_) { + mask_data_bytes = SafeInt(sequence_length) * all_sequence_length * sizeof(T); + } + + void* mask_data = nullptr; + if (mask_data_bytes > 0) { + mask_data = allocator->Alloc(mask_data_bytes); + memset(mask_data, 0, mask_data_bytes); + } + BufferUniquePtr mask_data_buffer(mask_data, BufferDeleter(allocator)); + + const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data() : nullptr; + const T* past_data = past != nullptr ? past->template Data() : nullptr; + T* present_data = present != nullptr ? present->template MutableData() : nullptr; + + ComputeAttentionProbs(static_cast(attention_probs), Q, K, + mask_index_data, static_cast(mask_data), + batch_size, sequence_length, past_sequence_length, head_size, + past_data, present_data, tp); + + // Compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S*) x V(B, N, S*, H) + auto out_tmp_data = + allocator->Alloc(SafeInt(batch_size) * num_heads_ * sequence_length * head_size * sizeof(T)); + BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator)); + + ComputeVxAttentionScore(output->template MutableData(), static_cast(out_tmp_data), static_cast(attention_probs), V, + batch_size, sequence_length, past_sequence_length, head_size, hidden_size, + past_data, present_data, tp); + + return Status::OK(); + } + + private: + // Helper function to compute the attention probs. It does 2 things: + // I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) + + // 1 x mask_data(B, N, S, S*) + // II.attention_probs(B, N, S, S*) = Softmax(attention_probs) + template + void ComputeAttentionProbs(T* attention_probs, // output buffer for the attention probs. Its size is BxNxSxS + const T* Q, // Q data. Its size is BxNxSxH + const T* K, // k data. Its size is BxNxSxH + const int32_t* mask_index, // mask index. nullptr if no mask or its size is B + T* mask_data, // buffer for mask data. Its size is: SxS* if is_unidirectional_; BxSxS* if mask_index; null otherwise + int batch_size, // batch size of self-attention + int sequence_length, // sequence length of self-attention + int past_sequence_length, // sequence length of past state + int head_size, // head size of self-attention + const T* past, // past state + T* present, // present state + ThreadPool* tp) const { + const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S + const size_t past_chunk_length = static_cast(past_sequence_length * head_size); // S' x H + const size_t input_chunk_length = static_cast(sequence_length * head_size); // S x H + const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H + + { + if (mask_data != nullptr) { + PrepareMask(mask_index, mask_data, is_unidirectional_, batch_size, sequence_length, past_sequence_length); + } else { // no any mask + memset(attention_probs, 0, batch_size * num_heads_ * sequence_length * all_sequence_length * sizeof(T)); + } + + const int loop_len = batch_size * num_heads_; + const float alpha = 1.0f / sqrt(static_cast(head_size)); + + // The cost of Gemm + const double cost = static_cast(head_size * sequence_length * all_sequence_length); + + ThreadPool::TryParallelFor(tp, loop_len, 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_; + + // broadcast mask data: SxS* or (Bx)SxS* -> (BxNx)SxS* + if (mask_data != nullptr) { + const T* broadcast_data_src = is_unidirectional_ ? reinterpret_cast(mask_data) : reinterpret_cast(mask_data) + batch_index * sequence_length * all_sequence_length; + T* broadcast_data_dest = reinterpret_cast(attention_probs) + sequence_length * all_sequence_length * i; + memcpy(broadcast_data_dest, broadcast_data_src, sequence_length * all_sequence_length * sizeof(T)); + } + + const T* k = K + input_chunk_length * i; + if (nullptr != present) { + // concatenate past_K and K : (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH + k = ConcatStateChunk(past, k, present, past_chunk_length, present_chunk_length, i); + } + + // gemm + // original transposed each iteration + // A: Q (B x N x) S x H (B x N x) S x H S x H + // B: K' (B x N x) S* x H (B x N x) H x S* H x S* + // C: attention_probs (B x N x) S x S* (B x N x) S x S* S x S* + math::Gemm(CblasNoTrans, CblasTrans, sequence_length, all_sequence_length, head_size, alpha, + Q + input_chunk_length * i, k, 1.0, + reinterpret_cast(attention_probs) + sequence_length * all_sequence_length * i, nullptr); + } + }); + } + + // attention_probs(B, N, S, S*) = Softmax(attention_probs) + { + const int N = batch_size * num_heads_ * sequence_length; + const int D = all_sequence_length; + ComputeAttentionSoftmaxInplace(attention_probs, N, D, tp); + } + } + + template + void ComputeVxAttentionScore(T* output, // buffer for the result with size BxSxNxH + T* tmp_buffer, // buffer for temp use with size is BxNxSxH + const T* attention_probs, // Attention probs with size BxNxSxS* + const T* V, // V value with size BxNxSxH + int batch_size, // batch size + int sequence_length, // sequence length + int past_sequence_length, // sequence length in past state + int head_size, // head size + int hidden_size, // hidden size + const T* past, // past state + T* present, // present state + ThreadPool* tp) const { + const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S + const size_t past_chunk_length = static_cast(past_sequence_length * head_size); // S' x H + const size_t input_chunk_length = static_cast(sequence_length * head_size); // S x H + const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H + + // Move the pointer of past and present to start of v values. + if (nullptr != past) { + past += batch_size * num_heads_ * past_sequence_length * head_size; + } + if (nullptr != present) { + present += batch_size * num_heads_ * all_sequence_length * head_size; + } + + const double cost = + static_cast(sequence_length) * static_cast(head_size) * static_cast(sequence_length); + + ThreadPool::TryParallelFor(tp, batch_size * num_heads_, cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t i = begin; i != end; ++i) { + const T* v = V + input_chunk_length * i; + if (nullptr != present) { + // concatenate past_V and V: (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH + v = ConcatStateChunk(past, v, present, past_chunk_length, present_chunk_length, i); + } + + T* current_tmp_data = reinterpret_cast(tmp_buffer) + input_chunk_length * i; + math::MatMul(sequence_length, head_size, all_sequence_length, + attention_probs + sequence_length * all_sequence_length * i, + v, current_tmp_data, nullptr); + + // transpose: out(B, S, N, H) = transpose out_tmp(B, N, S, H) + const int batch_index = static_cast(i / num_heads_); + const int head_index = static_cast(i % num_heads_); + T* src = current_tmp_data; + T* dest = output + (batch_index * sequence_length * num_heads_ + head_index) * head_size; + const auto bytes_to_copy = SafeInt(head_size) * sizeof(T); + for (int j = 0; j < sequence_length; j++) { + memcpy(dest, src, bytes_to_copy); + src += head_size; + dest += hidden_size; + } + } + }); + } +}; + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h index 4a94d0b45b..7507cb2ab2 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h @@ -68,7 +68,7 @@ void PrepareMask(const int32_t* mask_index, int past_sequence_length) { const int all_sequence_length = past_sequence_length + sequence_length; T* p_mask = mask_data; - if (is_unidirectional) { + if (is_unidirectional) { // unidirectional mask has shape SxS* for (int s_i = 0; s_i < sequence_length - 1; s_i++) { for (int m_i = past_sequence_length + s_i + 1; m_i < all_sequence_length; m_i++) { @@ -97,7 +97,7 @@ void PrepareMask(const int32_t* mask_index, // Concatenate a past state chunk S'xH with input state chunk SxH into present state chunk S*xH // Returns a pointer to the start of present state chunk. -template +template T* ConcatStateChunk(const T* past, const T* chunk, T* present, size_t past_chunk_length, size_t present_chunk_length, std::ptrdiff_t i) { T* start = present + i * present_chunk_length; @@ -112,138 +112,5 @@ T* ConcatStateChunk(const T* past, const T* chunk, T* present, size_t past_chunk return start; } -// Helper function to compute the attention probs. It does 2 things: -// I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) + -// 1 x mask_data(B, N, S, S*) -// II.attention_probs(B, N, S, S*) = Softmax(attention_probs) -template -void ComputeAttentionProbs(T* attention_probs, // output buffer for the attention probs. Its size is BxNxSxS - const T* Q, // Q data. Its size is BxNxSxH - const T* K, // k data. Its size is BxNxSxH - const int32_t* mask_index, // mask index. nullptr if no mask or its size is B - T* mask_data, // buffer for mask data. Its size is: SxS* if is_unidirectional; BxSxS* if mask_index; null otherwise - int batch_size, // batch size of self-attention - int sequence_length, // sequence length of self-attention - int past_sequence_length, // sequence length of past state - int head_size, // head size of self-attention - int num_heads, // number of heads of self-attention - bool is_unidirectional, // indicate if it is unidrectional. - const T* past, // past state - T* present, // present state - ThreadPool* tp) { - const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S - const size_t past_chunk_length = static_cast(past_sequence_length * head_size); // S' x H - const size_t input_chunk_length = static_cast(sequence_length * head_size); // S x H - const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H - - { - if (mask_data != nullptr) { - PrepareMask(mask_index, mask_data, is_unidirectional, batch_size, sequence_length, past_sequence_length); - } else { // no any mask - memset(attention_probs, 0, batch_size * num_heads * sequence_length * all_sequence_length * sizeof(T)); - } - - const int loop_len = batch_size * num_heads; - const float alpha = 1.0f / sqrt(static_cast(head_size)); - - // The cost of Gemm - const double cost = static_cast(head_size * sequence_length * all_sequence_length); - - ThreadPool::TryParallelFor(tp, loop_len, 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; - - // broadcast mask data: SxS* or (Bx)SxS* -> (BxNx)SxS* - if (mask_data != nullptr) { - const T* broadcast_data_src = is_unidirectional ? reinterpret_cast(mask_data) : reinterpret_cast(mask_data) + batch_index * sequence_length * all_sequence_length; - T* broadcast_data_dest = reinterpret_cast(attention_probs) + sequence_length * all_sequence_length * i; - memcpy(broadcast_data_dest, broadcast_data_src, sequence_length * all_sequence_length * sizeof(T)); - } - - const T* k = K + input_chunk_length * i; - if (nullptr != present) { - // concatenate past_K and K : (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH - k = ConcatStateChunk(past, k, present, past_chunk_length, present_chunk_length, i); - } - - // gemm - // original transposed each iteration - // A: Q (B x N x) S x H (B x N x) S x H S x H - // B: K' (B x N x) S* x H (B x N x) H x S* H x S* - // C: attention_probs (B x N x) S x S* (B x N x) S x S* S x S* - math::Gemm(CblasNoTrans, CblasTrans, sequence_length, all_sequence_length, head_size, alpha, - Q + input_chunk_length * i, k, 1.0, - reinterpret_cast(attention_probs) + sequence_length * all_sequence_length * i, nullptr); - } - }); - } - - // attention_probs(B, N, S, S*) = Softmax(attention_probs) - { - const int N = batch_size * num_heads * sequence_length; - const int D = all_sequence_length; - ComputeAttentionSoftmaxInplace(attention_probs, N, D, tp); - } -} - -template -void ComputeVxAttentionScore(T* output, // buffer for the result with size BxSxNxH - T* tmp_buffer, // buffer for temp use with size is BxNxSxH - const T* attention_probs, // Attention probs with size BxNxSxS* - const T* V, // V value with size BxNxSxH - int batch_size, // batch size - int sequence_length, // sequence length - int past_sequence_length, // sequence length in past state - int head_size, // head size - int num_heads, // number of heads - int hidden_size, // hidden size - const T* past, // past state - T* present, // present state - ThreadPool* tp) { - const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S - const size_t past_chunk_length = static_cast(past_sequence_length * head_size); // S' x H - const size_t input_chunk_length = static_cast(sequence_length * head_size); // S x H - const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H - - // Move the pointer of past and present to start of v values. - if (nullptr != past) { - past += batch_size * num_heads * past_sequence_length * head_size; - } - if (nullptr != present) { - present += batch_size * num_heads * all_sequence_length * head_size; - } - - const double cost = - static_cast(sequence_length) * static_cast(head_size) * static_cast(sequence_length); - - ThreadPool::TryParallelFor(tp, batch_size * num_heads, cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { - for (std::ptrdiff_t i = begin; i != end; ++i) { - - const T* v = V + input_chunk_length * i; - if (nullptr != present) { - // concatenate past_V and V: (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH - v = ConcatStateChunk(past, v, present, past_chunk_length, present_chunk_length, i); - } - - T* current_tmp_data = reinterpret_cast(tmp_buffer) + input_chunk_length * i; - math::MatMul(sequence_length, head_size, all_sequence_length, - attention_probs + sequence_length * all_sequence_length * i, - v, current_tmp_data, nullptr); - - // transpose: out(B, S, N, H) = transpose out_tmp(B, N, S, H) - const int batch_index = static_cast(i / num_heads); - const int head_index = static_cast(i % num_heads); - T* src = current_tmp_data; - T* dest = output + (batch_index * sequence_length * num_heads + head_index) * head_size; - const auto bytes_to_copy = SafeInt(head_size) * sizeof(T); - for (int j = 0; j < sequence_length; j++) { - memcpy(dest, src, bytes_to_copy); - src += head_size; - dest += hidden_size; - } - } - }); -} - } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc index ef344d8db4..b3b6c5775d 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc @@ -34,21 +34,23 @@ REGISTER_KERNEL_TYPED(float, uint8_t, int8_t) REGISTER_KERNEL_TYPED(float, uint8_t, uint8_t) template -QAttention::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionBase(info) { +QAttention::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) { } template Status QAttention::Compute(OpKernelContext* context) const { // Input and output shapes: - // Input 0 - input : (batch_size, sequence_length, hidden_size) - // Input 1 - weights : (hidden_size, 3 * hidden_size) - // Input 2 - bias : (3 * hidden_size) - // Input 3 - input_scale : scalar - // Input 4 - weight_scale : scalar - // Input 5 - mask_index : (batch_size) - // Input 6 - input_zero_point : scalar - // Input 7 - weight_zero_point : scalar - // Output : (batch_size, sequence_length, hidden_size) + // Input 0 - input : (batch_size, sequence_length, hidden_size) + // Input 1 - weights : (hidden_size, 3 * hidden_size) + // Input 2 - bias : (3 * hidden_size) + // Input 3 - input_scale : scalar + // Input 4 - weight_scale : scalar + // Input 5 - mask_index : (batch_size) + // Input 6 - input_zero_point : scalar + // Input 7 - weight_zero_point : scalar + // Input 8 - past : (2, batch_size, num_heads, past_sequence_length, head_size) + // Output 0 : (batch_size, sequence_length, hidden_size) + // Output 1 - present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size) // ORT_RETURN_IF_ERROR(CheckInputs(context)); const Tensor* input = context->Input(0); const Tensor* weights = context->Input(1); @@ -58,8 +60,9 @@ Status QAttention::Compute(OpKernelContext* context) const { const Tensor* mask_index = context->Input(5); const Tensor* i_zp_tensor = context->Input(6); const Tensor* w_zp_tensor = context->Input(7); + const Tensor* past_tensor = context->Input(8); - ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr)); + ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, past_tensor)); ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor), "input scale must be a scalar or 1D tensor of size 1"); @@ -146,51 +149,15 @@ Status QAttention::Compute(OpKernelContext* context) const { &dequant_scale, // output scale bias_data + weights_offset, // bias nullptr // use single-thread - ); + ); } }); } - // STEP.2: compute the attention score. It does 2 things: - // I. attention_probs(B, N, S, S) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S, H -> B, N, H, S) + - // 1 x mask_data(B, N, S, S) - // II.attention_probs(B, N, S, S) = Softmax(attention_probs) - size_t attention_probs_bytes = SafeInt(batch_size) * num_heads_ * sequence_length * sequence_length * element_size; - auto attention_probs = allocator->Alloc(attention_probs_bytes); - BufferUniquePtr scratch_buffer(attention_probs, BufferDeleter(allocator)); - - size_t mask_data_bytes = 0; - if (mask_index != nullptr) { - mask_data_bytes = SafeInt(batch_size) * sequence_length * sequence_length * element_size; - } else if (is_unidirectional_) { - mask_data_bytes = SafeInt(sequence_length) * sequence_length * element_size; - } - - void* mask_data = nullptr; - if (mask_data_bytes > 0) { - mask_data = allocator->Alloc(mask_data_bytes); - memset(mask_data, 0, mask_data_bytes); - } - BufferUniquePtr mask_data_buffer(mask_data, BufferDeleter(allocator)); - - const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data() : nullptr; - - int past_sequence_length = 0; - const T* past_data = nullptr; - T* present_data = nullptr; - ComputeAttentionProbs(static_cast(attention_probs), Q, K, mask_index_data, static_cast(mask_data), - batch_size, sequence_length, past_sequence_length, head_size, num_heads_, is_unidirectional_, - past_data, present_data, tp); - - // STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S) x V(B, N, S, H) - auto out_tmp_data = - allocator->Alloc(SafeInt(batch_size) * num_heads_ * sequence_length * head_size * element_size); - BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator)); - - ComputeVxAttentionScore(output->template MutableData(), static_cast(out_tmp_data), static_cast(attention_probs), V, - batch_size, sequence_length, past_sequence_length, head_size, num_heads_, hidden_size, past_data, present_data, tp); - - return Status::OK(); + // Compute the attention score and apply the score to V + return ApplyAttention(Q, K, V, mask_index, past_tensor, output, + batch_size, sequence_length, + head_size, hidden_size, context); } } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.h b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.h index af5858aca5..961347bf09 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.h +++ b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.h @@ -2,12 +2,12 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" -#include "contrib_ops/cpu/bert/attention.h" +#include "contrib_ops/cpu/bert/attention_cpu_base.h" namespace onnxruntime { namespace contrib { template -class QAttention : public OpKernel, public AttentionBase { +class QAttention : public OpKernel, public AttentionCPUBase { public: QAttention(const OpKernelInfo& info); diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.h b/onnxruntime/contrib_ops/cuda/bert/attention.h index b32ca9ee35..cd9bfe7d10 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention.h @@ -6,7 +6,7 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" #include "core/providers/cuda/cuda_common.h" -#include "contrib_ops/cpu/bert/attention.h" +#include "contrib_ops/cpu/bert/attention_base.h" namespace onnxruntime { namespace contrib { diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc index c1a5a251eb..974fdecf2f 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc @@ -47,7 +47,8 @@ Status QAttention::CheckInputs(const Tensor* input, const Tensor* weight_scale_tensor, const Tensor* mask_index, const Tensor* i_zp_tensor, - const Tensor* w_zp_tensor) const { + const Tensor* w_zp_tensor, + const Tensor* past_tensor) const { // Input and output shapes: // Input 0 - input : (batch_size, sequence_length, hidden_size) // Input 1 - weights : (hidden_size, 3 * hidden_size) @@ -59,7 +60,7 @@ Status QAttention::CheckInputs(const Tensor* input, // Input 7 - weight_zero_point : scalar // Output : (batch_size, sequence_length, hidden_size) - ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr)); + ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, past_tensor)); ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor), "input scale must be a scalar or 1D tensor of size 1"); @@ -88,15 +89,17 @@ Status QAttention::CheckInputs(const Tensor* input, template Status QAttention::ComputeInternal(OpKernelContext* context) const { // Input and output shapes: - // Input 0 - input : (batch_size, sequence_length, hidden_size) - // Input 1 - weights : (hidden_size, 3 * hidden_size) - // Input 2 - bias : (3 * hidden_size) - // Input 3 - input_scale : scalar - // Input 4 - weight_scale : scalar - // Input 5 - mask_index : (batch_size) - // Input 6 - input_zero_point : scalar - // Input 7 - weight_zero_point : scalar - // Output : (batch_size, sequence_length, hidden_size) + // Input 0 - input : (batch_size, sequence_length, hidden_size) + // Input 1 - weights : (hidden_size, 3 * hidden_size) + // Input 2 - bias : (3 * hidden_size) + // Input 3 - input_scale : scalar + // Input 4 - weight_scale : scalar + // Input 5 - mask_index : (batch_size) + // Input 6 - input_zero_point : scalar + // Input 7 - weight_zero_point : scalar + // Input 8 - past : (2, batch_size, num_heads, past_sequence_length, head_size) + // Output 0 - output : (batch_size, sequence_length, hidden_size) + // Output 1 - present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size) // ORT_RETURN_IF_ERROR(CheckInputs(context)); const Tensor* input = context->Input(0); const Tensor* weights = context->Input(1); @@ -106,6 +109,7 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { const Tensor* mask_index = context->Input(5); const Tensor* i_zp_tensor = context->Input(6); const Tensor* w_zp_tensor = context->Input(7); + const Tensor* past_tensor = context->Input(8); ORT_RETURN_IF_ERROR(CheckInputs(input, weights, @@ -114,7 +118,8 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { weight_scale_tensor, mask_index, i_zp_tensor, - w_zp_tensor)); + w_zp_tensor, + past_tensor)); const auto& shape = input->Shape(); int batch_size = static_cast(shape[0]); @@ -160,9 +165,9 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { m, n); - const int past_sequence_length = 0; - const T* past_data = nullptr; - T* present_data = nullptr; + int past_sequence_length = 0; + Tensor* present_tensor = GetPresent(context, past_tensor, batch_size, head_size, sequence_length, past_sequence_length); + size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, past_sequence_length); auto temp_buffer = GetScratchBuffer(workSpaceSize); if (!LaunchAttentionKernel( @@ -178,9 +183,8 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { element_size, is_unidirectional_, past_sequence_length, - past_data, - present_data - )) { + nullptr == past_tensor ? nullptr : past_tensor->template Data(), + nullptr == present_tensor ? nullptr : present_tensor->template MutableData())) { // Get last error to reset it to cudaSuccess. CUDA_CALL(cudaGetLastError()); return Status(common::ONNXRUNTIME, common::FAIL); diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h index ba636d5635..e019c0c1b2 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h @@ -36,7 +36,8 @@ class QAttention final : public CudaKernel, public AttentionBase { const Tensor* weight_scale_tensor, const Tensor* mask_index, const Tensor* i_zp_tensor, - const Tensor* w_zp_tensor) const; + const Tensor* w_zp_tensor, + const Tensor* past_tensor) const; }; } // namespace cuda diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 7b65660efd..eb67c26e52 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -406,11 +406,23 @@ mask_index shall not be provided.)DOC"; "zero point of quantized weight tensor. It's a scalar, which means a per-tensor/layer quantization.", "T2", OpSchema::Optional) + .Input( + 8, + "past", + "past state for key and value with shape (2, batch_size, num_heads, past_sequence_length, head_size).", + "T3", + OpSchema::Optional) .Output( 0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T3") + .Output( + 1, + "present", + "present state for key and value with shape (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)", + "T3", + OpSchema::Optional) .TypeConstraint("T1", {"tensor(int8)", "tensor(uint8)"}, "Constrain input and output types to int8 tensors.") .TypeConstraint("T2", {"tensor(int8)", "tensor(uint8)"}, "Constrain input and output types to int8 tensors.") .TypeConstraint("T3", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.") @@ -2616,7 +2628,9 @@ Example 4: .Attr("axis", "The first normalization dimension: normalization will be performed along dimensions axis : rank(inputs).", AttributeProto::INT, static_cast(-1)) - .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, 1e-5f) + .Attr("epsilon", + "The epsilon value to use to avoid division by zero.", + AttributeProto::FLOAT, 1e-5f) .AllowUncheckedAttributes() .Input(0, "X", "Input data tensor from the previous layer.", "T") .Input(1, "scale", "Scale tensor.", "T") diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index ab8710dc55..1458421a2c 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -1361,8 +1361,9 @@ class ONNXQuantizer: inputs.extend(quantized_input_names) inputs.extend([node.input[2]]) inputs.extend(scale_names) - inputs.extend([node.input[3]]) + inputs.extend([node.input[3] if len(node.input) > 3 else ""]) inputs.extend(zero_point_names) + inputs.extend([node.input[4] if len(node.input) > 4 else ""]) kwargs = {} for attribute in node.attribute: diff --git a/onnxruntime/test/common/tensor_op_test_utils.h b/onnxruntime/test/common/tensor_op_test_utils.h index bb35f89082..ca824e1bea 100644 --- a/onnxruntime/test/common/tensor_op_test_utils.h +++ b/onnxruntime/test/common/tensor_op_test_utils.h @@ -57,6 +57,49 @@ class RandomValueGenerator { return val; } + // Gaussian distribution for float + template + typename std::enable_if< + std::is_floating_point::value, + std::vector>::type + Gaussian(const std::vector& dims, TFloat mean, TFloat stddev) { + std::vector val(detail::SizeFromDims(dims)); + std::normal_distribution distribution(mean, stddev); + for (size_t i = 0; i < val.size(); ++i) { + val[i] = distribution(generator_); + } + return val; + } + + // Gaussian distribution for Integer + template + typename std::enable_if< + std::is_integral::value, + std::vector>::type + Gaussian(const std::vector& dims, TInt mean, TInt stddev) { + std::vector val(detail::SizeFromDims(dims)); + std::normal_distribution distribution(static_cast(mean), static_cast(stddev)); + for (size_t i = 0; i < val.size(); ++i) { + val[i] = static_cast(std::round(distribution(generator_))); + } + return val; + } + + // Gaussian distribution for Integer and Clamp to [min, max] + template + typename std::enable_if< + std::is_integral::value, + std::vector>::type + Gaussian(const std::vector& dims, TInt mean, TInt stddev, TInt min, TInt max) { + std::vector val(detail::SizeFromDims(dims)); + std::normal_distribution distribution(static_cast(mean), static_cast(stddev)); + for (size_t i = 0; i < val.size(); ++i) { + int64_t round_val = static_cast(std::round(distribution(generator_))); + val[i] = static_cast(std::min(std::max(round_val, min), max)); + } + return val; + } + template inline std::vector OneHot(const std::vector& dims, int64_t stride) { std::vector val(detail::SizeFromDims(dims), T(0)); diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 6f7a15ab64..897401c197 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -659,5 +659,34 @@ TEST(AttentionTest, AttentionPastStateBatch2) { use_past_state, past_sequence_length, head_size, &past_data, &present_data); } +TEST(AttentionTest, AttentionPastState_dynamic) { + // create rand inputs + RandomValueGenerator random{}; + + std::vector input_dims{2, 5, 768}; + std::vector input_data = random.Gaussian(input_dims, 0.0f, 0.3f); + + std::vector weight_dims{768, 2304}; + std::vector weight_data = random.Gaussian(weight_dims, 0.0f, 0.3f); + + std::vector bias_dims{2304}; + std::vector bias_data = random.Gaussian(bias_dims, 0.0f, 0.3f); + + std::vector past_dims{2, 2, 12, 15, 64}; + std::vector past_data = random.Gaussian(past_dims, 0.0f, 0.3f); + + OpTester test("Attention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", 12); + test.AddAttribute("unidirectional", 1); + test.AddInput("input", input_dims, input_data); + test.AddInput("weight", weight_dims, weight_data); + test.AddInput("bias", bias_dims, bias_data); + test.AddMissingOptionalInput(); + test.AddInput("past", past_dims, past_data); + + test.AddReferenceOutputs("testdata/attention_past_state.onnx"); + test.Run(); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/dynamic_quantize_matmul_test.cc b/onnxruntime/test/contrib_ops/dynamic_quantize_matmul_test.cc index 839278b8a7..991cbf87bd 100644 --- a/onnxruntime/test/contrib_ops/dynamic_quantize_matmul_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_quantize_matmul_test.cc @@ -21,142 +21,40 @@ namespace onnxruntime { namespace test { template -class DynamicQuantizeMatMulOpTester : public OpTester { - public: - DynamicQuantizeMatMulOpTester(const char* op, - const std::vector& A_dims, - const std::vector& B_dims, - const std::vector& Y_dims, - string&& model_file, - int opset_version = 1, - const char* domain = onnxruntime::kMSDomain) : OpTester(op, opset_version, domain), - A_dims_(A_dims), - B_dims_(B_dims), - Y_dims_(Y_dims), - model_file_(model_file) { - Init(); - } - void Init() { - // create rand inputs - RandomValueGenerator random{}; - A_data_ = random.Uniform(A_dims_, -1.0f, 1.0f); +void TestDynamicQuantizeMatMul(const std::vector& A_dims, + std::vector B_dims, + const std::string& reference_model) { + // create rand inputs + RandomValueGenerator random{}; - std::vector tmp_B_data = random.Uniform(B_dims_, std::numeric_limits::min(), std::numeric_limits::max()); - std::transform(tmp_B_data.begin(), tmp_B_data.end(), std::back_inserter(B_data_), [](int32_t v) -> T { - return static_cast(v); - }); + std::vector A_data = random.Uniform(A_dims, -1.0f, 1.0f); - B_zero_point_ = {static_cast(random.Uniform({1}, std::numeric_limits::min(), std::numeric_limits::max())[0])}; - B_scale_ = random.Uniform({1}, -0.1f, 0.1f); + std::vector B_data; + std::vector tmp_B_data = random.Uniform(B_dims, std::numeric_limits::min(), std::numeric_limits::max()); + std::transform(tmp_B_data.begin(), tmp_B_data.end(), std::back_inserter(B_data), [](int32_t v) -> T { + return static_cast(v); + }); - const int64_t Y_size = std::accumulate(Y_dims_.cbegin(), Y_dims_.cend(), static_cast(1), std::multiplies{}); - Y_data_.resize(Y_size); + std::vector B_scale = random.Uniform({1}, -0.1f, 0.1f); + std::vector B_zero_point = {static_cast(random.Uniform({1}, std::numeric_limits::min(), std::numeric_limits::max())[0])}; - AddInput("A", A_dims_, A_data_); - AddInput("B", B_dims_, B_data_); - AddInput("b_scale", {1}, B_scale_); - AddInput("b_zero_point", {1}, B_zero_point_); + OpTester test("DynamicQuantizeMatMul", 1, onnxruntime::kMSDomain); + test.AddInput("A", A_dims, A_data); + test.AddInput("B", B_dims, B_data); + test.AddInput("b_scale", {1}, B_scale); + test.AddInput("b_zero_point", {1}, B_zero_point); - AddOutput("Y", Y_dims_, Y_data_); - } - void Run() { -#ifndef NDEBUG - // run_called_ to true to avoid a complaining in the destructor of OpTester - run_called_ = true; -#endif - std::vector cpu_fetches; - std::vector subgraph_fetches; - Compute(cpu_fetches); - ComputeOriginalSubgraph(subgraph_fetches); - - // Compare CPU with original subgraph on the output(0) - Y - ASSERT_TRUE(cpu_fetches.size() >= subgraph_fetches.size()); - for (size_t i = 0; i < subgraph_fetches.size(); i++) { - if (cpu_fetches[i].IsTensor() && subgraph_fetches[i].IsTensor()) { - VLOGS_DEFAULT(1) << "Checking tensor " << i; - CheckTensor(subgraph_fetches[i].Get(), cpu_fetches[i].Get(), 1e-3, 1e-3); - } - } - } - - private: - void Compute(std::vector& cpu_fetches) { - auto p_model = BuildGraph(); - auto& graph = p_model->MainGraph(); - - Status status = graph.Resolve(); - ASSERT_TRUE(status.IsOK()) << status; - - // Hookup the inputs and outputs - NameMLValMap feeds; - std::vector output_names; - FillFeedsAndOutputNames(feeds, output_names); - - SessionOptions so; - so.session_logid = op_; - so.session_log_verbosity_level = 1; - - RunOptions run_options; - run_options.run_tag = op_; - run_options.run_log_verbosity_level = 1; - - // run with DynamicQuantizeMatMul - InferenceSession session_object{so, GetEnvironment()}; - std::string s1; - p_model->ToProto().SerializeToString(&s1); - std::istringstream str(s1); - ASSERT_TRUE((status = session_object.Load(str)).IsOK()) << status; - ASSERT_TRUE((status = session_object.Initialize()).IsOK()) << status; - ASSERT_TRUE((status = session_object.Run(run_options, feeds, output_names, &cpu_fetches)).IsOK()); - } - - void ComputeOriginalSubgraph(std::vector& subgraph_fetches) { - NameMLValMap feeds; - OrtValue ml_value; - std::vector output_names; - FillFeedsAndOutputNames(feeds, output_names); - - SessionOptions so; - so.session_logid = op_; - so.session_log_verbosity_level = 1; - so.graph_optimization_level = TransformerLevel::Level1; - - RunOptions run_options; - run_options.run_tag = op_; - run_options.run_log_verbosity_level = 1; - - Status status; - InferenceSession subgraph_session_object{so, GetEnvironment()}; - ASSERT_TRUE((status = subgraph_session_object.Load(model_file_)).IsOK()) << status; - ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status; - ASSERT_TRUE((status = subgraph_session_object.Run(run_options, feeds, output_names, &subgraph_fetches)).IsOK()) << status; - } - - private: - std::vector A_dims_; - std::vector B_dims_; - std::vector Y_dims_; - - std::vector A_data_; - std::vector B_data_; - std::vector B_scale_; - std::vector B_zero_point_; - std::vector Y_data_; - - std::string model_file_; -}; + test.AddReferenceOutputs(reference_model); + test.Run(); +} TEST(DynamicQuantizeMatMul, Int8_test) { #ifdef MLAS_SUPPORTS_GEMM_U8X8 std::vector A_dims{4, 128}; std::vector B_dims{128, 128}; std::vector Y_dims{4, 128}; - DynamicQuantizeMatMulOpTester test("DynamicQuantizeMatMul", - A_dims, - B_dims, - Y_dims, - "testdata/dynamic_quantize_matmul_int8.onnx"); - test.Run(); + + TestDynamicQuantizeMatMul(A_dims, B_dims, "testdata/dynamic_quantize_matmul_int8.onnx"); #endif } @@ -164,12 +62,8 @@ TEST(DynamicQuantizeMatMul, UInt8_test) { std::vector A_dims{4, 128}; std::vector B_dims{128, 128}; std::vector Y_dims{4, 128}; - DynamicQuantizeMatMulOpTester test("DynamicQuantizeMatMul", - A_dims, - B_dims, - Y_dims, - "testdata/dynamic_quantize_matmul_uint8.onnx"); - test.Run(); + + TestDynamicQuantizeMatMul(A_dims, B_dims, "testdata/dynamic_quantize_matmul_uint8.onnx"); } } // namespace test diff --git a/onnxruntime/test/contrib_ops/layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/layer_norm_op_test.cc index 63c4499e2a..c8c62434d2 100644 --- a/onnxruntime/test/contrib_ops/layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/layer_norm_op_test.cc @@ -18,206 +18,29 @@ using namespace std; namespace onnxruntime { namespace test { -class LayerNormOpTester : public OpTester { - public: - LayerNormOpTester(const char* op, - const std::vector& X_dims, - const std::vector& scale_dims, - const std::vector& B_dims, - const std::vector& Y_dims, - float epsilon, - int64_t axis = -1, - int64_t keep_dims = 1, - int opset_version = 1, - const char* domain = onnxruntime::kOnnxDomain) : OpTester(op, opset_version, domain), - X_dims_(X_dims), - scale_dims_(scale_dims), - B_dims_(B_dims), - Y_dims_(Y_dims), - epsilon_(epsilon), - axis_(axis), - keep_dims_(keep_dims) { - Init(); - } - void Init() { - AddAttribute("axis", axis_); - AddAttribute("keep_dims", keep_dims_); - AddAttribute("epsilon", epsilon_); - - // create rand inputs - RandomValueGenerator random{}; - X_data_ = random.Uniform(X_dims_, 0.0f, 1.0f); - scale_data_ = random.Uniform(scale_dims_, 0.0f, 1.0f); - B_data_ = random.Uniform(B_dims_, 0.0f, 1.0f); - - const int64_t Y_size = std::accumulate(Y_dims_.cbegin(), Y_dims_.cend(), static_cast(1), std::multiplies{}); - Y_data_.resize(Y_size); - - AddInput("X", X_dims_, X_data_); - AddInput("scale", scale_dims_, scale_data_, true); - AddInput("B", B_dims_, B_data_, true); - - AddOutput("output", Y_dims_, Y_data_); - AddOutput("mean", mean_); - AddOutput("inv_std_var", inv_std_var_); - } - void Run() { -#ifndef NDEBUG - run_called_ = true; -#endif - std::vector cpu_fetches; - std::vector cuda_fetches; - std::vector subgraph_fetches; - ComputeWithCPU(cpu_fetches); - ComputeWithCUDA(cuda_fetches); - ComputeOriSubgraphWithCPU(subgraph_fetches); - - // Compare CPU with original subgraph on the output(0) - Y - ASSERT_TRUE(cpu_fetches.size() >= subgraph_fetches.size()); - for (size_t i = 0; i < subgraph_fetches.size(); i++) { - if (cpu_fetches[i].IsTensor() && subgraph_fetches[i].IsTensor()) { - VLOGS_DEFAULT(1) << "Checking tensor " << i; - CheckTensor(subgraph_fetches[i].Get(), cpu_fetches[i].Get(), 1e-3, 1e-3); - } - } - - // Compare GPU with original subgraph on the output(0) - Y - if (DefaultCudaExecutionProvider()) { - ASSERT_TRUE(cuda_fetches.size() >= subgraph_fetches.size()); - for (size_t i = 0; i < subgraph_fetches.size(); i++) { - if (cuda_fetches[i].IsTensor() && subgraph_fetches[i].IsTensor()) { - VLOGS_DEFAULT(1) << "Checking tensor " << i; - CheckTensor(subgraph_fetches[i].Get(), cuda_fetches[i].Get(), 1e-3, 1e-3); - } - } - } - } - - private: - void ComputeWithCPU(std::vector& cpu_fetches); - void ComputeWithCUDA(std::vector& cuda_fetches); - void ComputeOriSubgraphWithCPU(std::vector& subgraph_fetches); - - private: - std::vector X_dims_; - std::vector scale_dims_; - std::vector B_dims_; - std::vector Y_dims_; - - std::vector X_data_; - std::vector scale_data_; - std::vector B_data_; - std::vector Y_data_; - float mean_; - float inv_std_var_; - - float epsilon_; - int64_t axis_; - int64_t keep_dims_; -}; - -void LayerNormOpTester::ComputeWithCPU(std::vector& cpu_fetches) { - auto p_model = BuildGraph(); - auto& graph = p_model->MainGraph(); - - Status status = graph.Resolve(); - ASSERT_TRUE(status.IsOK()) << status; - - // Hookup the inputs and outputs - std::unordered_map feeds; - std::vector output_names; - FillFeedsAndOutputNames(feeds, output_names); - - SessionOptions so; - so.session_logid = op_; - so.session_log_verbosity_level = 1; - - RunOptions run_options; - run_options.run_tag = op_; - run_options.run_log_verbosity_level = 1; - - // run with LayerNormalization - InferenceSession layernorm_session_object{so, GetEnvironment()}; - std::string s1; - p_model->ToProto().SerializeToString(&s1); - std::istringstream str(s1); - ASSERT_TRUE((status = layernorm_session_object.Load(str)).IsOK()) << status; - ASSERT_TRUE((status = layernorm_session_object.Initialize()).IsOK()) << status; - ASSERT_TRUE((status = layernorm_session_object.Run(run_options, feeds, output_names, &cpu_fetches)).IsOK()); -} - -void LayerNormOpTester::ComputeWithCUDA(std::vector& cuda_fetches) { - if (DefaultCudaExecutionProvider() == nullptr) { - return; - } - - auto p_model = BuildGraph(); - auto& graph = p_model->MainGraph(); - - Status status = graph.Resolve(); - ASSERT_TRUE(status.IsOK()) << status; - - // Hookup the inputs and outputs - std::unordered_map feeds; - std::vector output_names; - FillFeedsAndOutputNames(feeds, output_names); - - SessionOptions so; - so.session_logid = op_; - so.session_log_verbosity_level = 1; - - RunOptions run_options; - run_options.run_tag = op_; - run_options.run_log_verbosity_level = 1; - - auto cuda_execution_provider = DefaultCudaExecutionProvider(); - InferenceSession cuda_session_object{so, GetEnvironment()}; - EXPECT_TRUE(cuda_session_object.RegisterExecutionProvider(std::move(cuda_execution_provider)).IsOK()); - - std::string s; - p_model->ToProto().SerializeToString(&s); - std::istringstream str2(s); - - EXPECT_TRUE((status = cuda_session_object.Load(str2)).IsOK()) << status; - EXPECT_TRUE((status = cuda_session_object.Initialize()).IsOK()) << status; - EXPECT_TRUE((status = cuda_session_object.Run(run_options, feeds, output_names, &cuda_fetches)).IsOK()) << status; -} - -void LayerNormOpTester::ComputeOriSubgraphWithCPU(std::vector& subgraph_fetches) { - NameMLValMap feeds; - OrtValue ml_value; - std::vector output_names{"Y"}; - - CreateMLValue(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), X_dims_, X_data_, &ml_value); - feeds.insert(std::make_pair("X", ml_value)); - CreateMLValue(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), scale_dims_, scale_data_, &ml_value); - feeds.insert(std::make_pair("Scale", ml_value)); - CreateMLValue(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), B_dims_, B_data_, &ml_value); - feeds.insert(std::make_pair("B", ml_value)); - - SessionOptions so; - so.session_logid = op_; - so.session_log_verbosity_level = 1; - - RunOptions run_options; - run_options.run_tag = op_; - run_options.run_log_verbosity_level = 1; - - Status status; - InferenceSession subgraph_session_object{so, GetEnvironment()}; - ASSERT_TRUE((status = subgraph_session_object.Load("testdata/layernorm.onnx")).IsOK()) << status; - ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status; - ASSERT_TRUE((status = subgraph_session_object.Run(run_options, feeds, output_names, &subgraph_fetches)).IsOK()) << status; -} - TEST(LayerNormTest, BERTLayerNorm) { - float epsilon = 1e-12f; + OpTester tester("LayerNormalization", 1 /*opset_version*/); + tester.AddAttribute("axis", -1); + tester.AddAttribute("epsilon", 1e-12f); + + // create rand inputs + RandomValueGenerator random{}; + std::vector X_dims{4, 128}; + std::vector X_data = random.Uniform(X_dims, 0.0f, 1.0f); + tester.AddInput("X", X_dims, X_data); + std::vector scale_dims{128}; + std::vector scale_data = random.Uniform(scale_dims, 0.0f, 1.0f); + tester.AddInput("Scale", scale_dims, scale_data); + std::vector B_dims{128}; - std::vector Y_dims{4, 128}; - LayerNormOpTester test("LayerNormalization", X_dims, scale_dims, B_dims, Y_dims, epsilon, -1, 1, 9); - test.Run(); + std::vector B_data = random.Uniform(B_dims, 0.0f, 1.0f); + tester.AddInput("B", B_dims, B_data); + + tester.AddReferenceOutputs("testdata/layernorm.onnx"); + + tester.Run(); } } // namespace test diff --git a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc index a77fa49ba2..4b4dfb2b79 100644 --- a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc @@ -705,5 +705,70 @@ TEST(QAttentionTest, QAttentionUnidirectional_CUDA) { true /*is_unidirectional*/); } +template +void TestQuantizedAttentionPastState(int64_t batch, + int64_t seq_len, + int64_t past_seq_len, + int64_t hidden_size, + int64_t head_number, + int64_t head_size, + const std::string& reference_model) { + // create rand inputs + RandomValueGenerator random{}; + + constexpr InputT input_min = std::numeric_limits::min(); + constexpr InputT input_max = std::numeric_limits::max(); + constexpr int32_t input_range = input_max - input_min; + + InputT input_mean = (input_min + input_max) / 2 + 1; + std::vector input_zero_point{input_mean}; + + std::vector input_dims{batch, seq_len, hidden_size}; + std::vector input_data = random.Gaussian(input_dims, input_mean, static_cast(input_range / 6), input_min, input_max); + + constexpr WeightT weight_min = std::numeric_limits::min(); + constexpr WeightT weight_max = std::numeric_limits::max(); + constexpr int32_t weight_range = weight_max - weight_min; + + WeightT weight_mean = (weight_min + weight_max) / 2 + 1; + std::vector weight_zero_point{weight_mean}; + + std::vector weight_dims{hidden_size, 3 * hidden_size}; + std::vector weight_data = random.Gaussian(weight_dims, weight_mean, static_cast(weight_range / 6), weight_min, weight_max); + + std::vector bias_dims{3 * hidden_size}; + std::vector bias_data = random.Gaussian(bias_dims, 0.0f, 0.3f); + + std::vector input_scale{0.005f}; + std::vector weight_scale{0.005f}; + + std::vector past_dims{2, batch, head_number, past_seq_len, head_size}; + std::vector past_data = random.Gaussian(past_dims, 0.0f, 0.3f); + + OpTester test("QAttention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", head_number); + test.AddAttribute("unidirectional", 1); + test.AddInput("input", input_dims, input_data); + test.AddInput("weight", weight_dims, weight_data); + test.AddInput("bias", bias_dims, bias_data); + test.AddInput("input_scale", {1}, input_scale); + test.AddInput("weight_scale", {1}, weight_scale); + test.AddMissingOptionalInput(); + test.AddInput("input_zero_point", {1}, input_zero_point); + test.AddInput("weight_zero_point", {1}, weight_zero_point); + test.AddInput("past", past_dims, past_data); + + test.AddReferenceOutputs(reference_model); + test.Run(); +} + +TEST(QAttentionTest, QAttentionPastState_u8u8) { + TestQuantizedAttentionPastState(2, 5, 15, 768, 12, 64, "testdata/attention_past_state.u8u8.onnx"); +} + +TEST(QAttentionTest, QAttentionPastState_u8s8) { + TestQuantizedAttentionPastState(2, 5, 15, 768, 12, 64, "testdata/attention_past_state.u8s8.onnx"); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index ab826ca0f5..889571ec27 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -393,6 +393,16 @@ void OpTester::FillFeedsAndOutputNames( } } +void OpTester::FillFeeds(std::unordered_map& feeds) { + for (size_t i = 0; i < input_data_.size(); ++i) { + if (std::find(initializer_index_.begin(), initializer_index_.end(), i) == + initializer_index_.end() && + input_data_[i].def_.Exists()) { + feeds[input_data_[i].def_.Name()] = input_data_[i].data_; + } + } +} + void OpTester::SetOutputAbsErr(const char* name, float v) { auto it = std::find_if(output_data_.begin(), output_data_.end(), @@ -524,7 +534,7 @@ std::vector OpTester::ExecuteModel( // Disable expected_failure_string checks for OpenVINO EP if (provider_type != kOpenVINOExecutionProvider) { EXPECT_THAT(status.ErrorMessage(), - testing::HasSubstr(expected_failure_string)); + testing::HasSubstr(expected_failure_string)); } } else { LOGS_DEFAULT(ERROR) << "Initialize failed with status: " @@ -804,7 +814,7 @@ void OpTester::Run( valid = false; for (auto& custom_session_registry : custom_session_registries_) { if (KernelRegistry::HasImplementationOf(*custom_session_registry->GetKernelRegistry(), - node, execution_provider->Type())) { + node, execution_provider->Type())) { valid = true; break; } @@ -843,6 +853,56 @@ void OpTester::Run( } } +void OpTester::AddReferenceOutputs(const std::string& model_path) { + SessionOptions so; + so.session_logid = op_; + so.session_log_verbosity_level = 1; + + RunOptions run_options; + run_options.run_tag = op_; + run_options.run_log_verbosity_level = 1; + + Status status; + InferenceSession subgraph_session_object{so, GetEnvironment()}; + ASSERT_TRUE((status = subgraph_session_object.Load(model_path)).IsOK()) << status; + ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status; + + // Retrieve output names + auto model_outputs = subgraph_session_object.GetModelOutputs(); + ASSERT_TRUE(model_outputs.first.IsOK()); + std::vector output_names; + std::transform(model_outputs.second->begin(), + model_outputs.second->end(), + std::back_inserter(output_names), + [](const onnxruntime::NodeArg* node_arg) -> std::string { return node_arg->Name(); }); + + NameMLValMap feeds; + FillFeeds(feeds); + + std::vector subgraph_fetches; + ASSERT_TRUE((status = subgraph_session_object.Run(run_options, feeds, output_names, &subgraph_fetches)).IsOK()) << status; + + for (size_t out_idx = 0; out_idx < subgraph_fetches.size(); out_idx++) { + // Retrieve TypeProto + ASSERT_TRUE(subgraph_fetches[out_idx].Type()->IsTensorType()) << status; + const Tensor& t = subgraph_fetches[out_idx].Get(); + const TensorTypeBase* tensor_type = DataTypeImpl::TensorTypeFromONNXEnum(t.GetElementType()); + + // Construct a temp TypeProto with shape information + ONNX_NAMESPACE::TypeProto tmp_type_proto(*(tensor_type->GetTypeProto())); + auto mutable_shape = tmp_type_proto.mutable_tensor_type()->mutable_shape(); + for (auto i : t.Shape().GetDims()) { + auto* mutable_dim = mutable_shape->add_dim(); + mutable_dim->set_dim_value(i); + } + + output_data_.push_back(Data(NodeArg(output_names[out_idx], &tmp_type_proto), + std::move(subgraph_fetches[out_idx]), + optional(), + optional())); + } +} + #ifdef ENABLE_TRAINING template std::vector OpTester::ExecuteModel( Model& model, training::TrainingSession& session_object, diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index 162556586f..6a8c7830e9 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -210,7 +210,8 @@ const SequenceTensorTypeProto SequenceTensorType::s_sequence // 1. Create one with the op name // 2. Call AddAttribute with any attributes // 3. Call AddInput for all the inputs -// 4. Call AddOutput with all expected outputs +// 4. Call AddOutput with all expected outputs, +// Or call AddReferenceOutputs to compute reference outputs with the model // 5. Call Run // Not all tensor types and output types are added, if a new input type is used, add it to the TypeToDataType list // above for new output types, add a new specialization for Check<> See current usage for an example, should be self @@ -382,6 +383,9 @@ class OpTester { optional(), optional())); } + // Generate the reference outputs with the model file + void AddReferenceOutputs(const std::string& model_path); + void AddCustomOpRegistry(std::shared_ptr registry) { custom_schema_registries_.push_back(registry->GetOpschemaRegistry()); custom_session_registries_.push_back(registry); @@ -483,6 +487,8 @@ class OpTester { void FillFeedsAndOutputNames(std::unordered_map& feeds, std::vector& output_names); + void FillFeeds(std::unordered_map& feeds); + template std::vector ExecuteModel(Model& model, SessionType& session_object, @@ -545,7 +551,7 @@ class OpTester { const static std::unordered_set reserved_symbolic{"batch", "seq"}; for (size_t i = 0; i < dim_params_data.size(); ++i) { - if (reserved_symbolic.find(dim_params_data[i])!= reserved_symbolic.end()) { + if (reserved_symbolic.find(dim_params_data[i]) != reserved_symbolic.end()) { new_shape.add_dim()->set_dim_param(dim_params_data[i]); } else { ASSERT_TRUE(std::stoi(dim_params_data[i]) == dims[i]); diff --git a/onnxruntime/test/testdata/attention_past_state.onnx b/onnxruntime/test/testdata/attention_past_state.onnx new file mode 100644 index 0000000000..4bcb19d3b2 Binary files /dev/null and b/onnxruntime/test/testdata/attention_past_state.onnx differ diff --git a/onnxruntime/test/testdata/attention_past_state.u8s8.onnx b/onnxruntime/test/testdata/attention_past_state.u8s8.onnx new file mode 100644 index 0000000000..b6fe39ff3c Binary files /dev/null and b/onnxruntime/test/testdata/attention_past_state.u8s8.onnx differ diff --git a/onnxruntime/test/testdata/attention_past_state.u8u8.onnx b/onnxruntime/test/testdata/attention_past_state.u8u8.onnx new file mode 100644 index 0000000000..5ce855c0d8 Binary files /dev/null and b/onnxruntime/test/testdata/attention_past_state.u8u8.onnx differ