diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.cc b/onnxruntime/contrib_ops/cpu/bert/attention.cc index 08bb7e275e..1da8aa411d 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/attention.cc @@ -33,13 +33,13 @@ AttentionBase::AttentionBase(const OpKernelInfo& info) { Status AttentionBase::CheckInputs(const TensorShape& input_shape, const TensorShape& weights_shape, const TensorShape& bias_shape, - const Tensor* mask_index, + const Tensor*& mask_index, const Tensor* past) const { // Input shapes: // input : (batch_size, sequence_length, hidden_size) // weights : (hidden_size, 3 * hidden_size) // bias : (3 * hidden_size) - // mask_index : nullptr, (batch_size), (2 * batch_size), or (batch_size, past_sequence_length + sequence_length) + // mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length) // past : (2, batch_size, num_heads, past_sequence_length, head_size) const auto& dims = input_shape.GetDims(); @@ -112,7 +112,14 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape, } } else if (mask_dims.size() == 2) { if (static_cast(mask_dims[0]) != batch_size || static_cast(mask_dims[1]) != past_sequence_length + sequence_length) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' with raw attention mask shall have shape batch_size x (past_sequence_length + sequence_length)"); + // Add operator supports broadcasting. Here we handle a case with only one element in the 2nd dimension. + if ((static_cast(mask_dims[0]) == batch_size || static_cast(mask_dims[0]) == 1) && static_cast(mask_dims[1]) == 1) { + // Mask will have same value after propogation, which has same effect as no mask. + mask_index = nullptr; + } + else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' with raw attention mask shall have shape batch_size x (past_sequence_length + sequence_length)"); + } } } else { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask_index' is expected to have 1 or 2 dimensions, got ", diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_base.h index d856f8f5b3..d7fe44f077 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.h @@ -12,10 +12,11 @@ namespace contrib { class AttentionBase { protected: AttentionBase(const OpKernelInfo& info); + Status CheckInputs(const TensorShape& input_shape, const TensorShape& weights_shape, const TensorShape& bias_shape, - const Tensor* mask_index, + const Tensor*& mask_index, // For dummy mask with shape (1, 1) or (batch_size, 1), it will be updated to nullptr. const Tensor* past) const; Tensor* GetPresent(OpKernelContext* context, diff --git a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc index 43da9c3a5e..4dd0c95ad8 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc @@ -107,7 +107,7 @@ Status QAttention::Compute(OpKernelContext* context) const { // Input 2 - bias : (3 * hidden_size) // Input 3 - input_scale : scalar // Input 4 - weight_scale : scalar - // Input 5 - mask_index : (batch_size) + // Input 5 - mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length) // 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) diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc index 6e76b343ca..67d51b53d5 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc @@ -45,7 +45,7 @@ Status QAttention::CheckInputs(const Tensor* input, const Tensor* bias, const Tensor* input_scale_tensor, const Tensor* weight_scale_tensor, - const Tensor* mask_index, + const Tensor*& mask_index, const Tensor* i_zp_tensor, const Tensor* w_zp_tensor, const Tensor* past_tensor) const { @@ -55,7 +55,7 @@ Status QAttention::CheckInputs(const Tensor* input, // Input 2 - bias : (3 * hidden_size) // Input 3 - input_scale : scalar // Input 4 - weight_scale : scalar - // Input 5 - mask_index : (batch_size) + // Input 5 - mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length) // Input 6 - input_zero_point : scalar // Input 7 - weight_zero_point : scalar // Output : (batch_size, sequence_length, hidden_size) diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h index e019c0c1b2..403ec37a75 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h @@ -34,7 +34,7 @@ class QAttention final : public CudaKernel, public AttentionBase { const Tensor* bias, const Tensor* input_scale_tensor, const Tensor* weight_scale_tensor, - const Tensor* mask_index, + const Tensor*& mask_index, const Tensor* i_zp_tensor, const Tensor* w_zp_tensor, const Tensor* past_tensor) const; diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 831505ad0e..299e55ade4 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -12,7 +12,8 @@ namespace test { enum MaskIndexType { kMaskIndexEnd = 0, kMaskIndexEndAndStart, - kMaskRaw + kMaskRaw, + kMaskDummy // Dummy mask with shape [1, 1] or [batch_size, 1] }; static void RunAttentionTest( @@ -49,7 +50,25 @@ static void RunAttentionTest( std::vector mask_index_dims_1 = {batch_size}; std::vector mask_index_dims_2 = {2 * batch_size}; std::vector mask_index_dims_3 = {batch_size, past_sequence_length + sequence_length}; - std::vector mask_index_dims = (mask_index_type == kMaskIndexEnd ? mask_index_dims_1 : (mask_index_type == kMaskIndexEndAndStart ? mask_index_dims_2 : mask_index_dims_3)); + std::vector mask_index_dims_4 = {batch_size, 1}; + std::vector mask_index_dims; + switch (mask_index_type) { + case kMaskIndexEnd: + mask_index_dims = mask_index_dims_1; + break; + case kMaskIndexEndAndStart: + mask_index_dims = mask_index_dims_2; + break; + case kMaskRaw: + mask_index_dims = mask_index_dims_3; + break; + case kMaskDummy: + mask_index_dims = mask_index_dims_4; + break; + default: + assert(0); // shall not reach here. + break; + } std::vector past_dims = {2, batch_size, number_of_heads, past_sequence_length, head_size}; std::vector present_dims = {2, batch_size, number_of_heads, past_sequence_length + sequence_length, head_size}; @@ -1134,6 +1153,48 @@ TEST(AttentionTest, AttentionMask2DNoWord) { use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskRaw); } + +TEST(AttentionTest, AttentionDummyMask2D) { + int batch_size = 2; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.5f, 0.2f, 0.3f, -0.6f, + 0.8f, -0.5f, 0.0f, 1.f, + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + std::vector mask_index_data = {1, 1}; + + std::vector output_data = { + 3.9696791172027588f, 0.073143675923347473f, 4.25f, 5.65f, + 3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.65f, + 3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.65f, + 3.9696791172027588f, 0.073143675923347473f, 4.25f, 5.65f}; + + bool use_float16 = false; + bool is_unidirectional = false; + bool use_past_state = false; + int past_sequence_length = 0; + const std::vector* past_data = nullptr; + const std::vector* present_data = nullptr; + + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, + use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskDummy); +} + TEST(AttentionTest, AttentionMaskIndexOutOfRange) { int batch_size = 2; int sequence_length = 2;