diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index e374dfdb0b..22ba16330a 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -123,6 +123,8 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Attributes
+- mask_filter_value : float
+- The value to be filled in the attention mask. Default value is -10000.0f
- num_heads : int (required)
- Number of attention heads
- past_present_share_buffer : int
@@ -967,6 +969,8 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Attributes
+- mask_filter_value : float
+- The value to be filled in the attention mask. Default value is negative infinity
- num_heads : int (required)
- Number of attention heads
@@ -2120,6 +2124,8 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Attributes
+- mask_filter_value : float
+- The value to be filled in the attention mask. Default value is negative infinity
- num_heads : int (required)
- Number of attention heads
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.cc b/onnxruntime/contrib_ops/cpu/bert/attention_base.cc
index 2558c52df3..34a58bd254 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_base.cc
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.cc
@@ -248,6 +248,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
output_parameters->num_heads = num_heads_;
output_parameters->is_unidirectional = is_unidirectional_;
output_parameters->past_present_share_buffer = (past_present_share_buffer_ != 0);
+ output_parameters->mask_filter_value = mask_filter_value_;
output_parameters->mask_type = mask_type;
}
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_base.h
index 71a4ed8c6d..689db55965 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_base.h
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.h
@@ -37,6 +37,7 @@ class AttentionBase {
num_heads_ = static_cast(num_heads);
is_unidirectional_ = info.GetAttrOrDefault("unidirectional", 0) == 1;
+ mask_filter_value_ = info.GetAttrOrDefault("mask_filter_value", -10000.0f);
if (!info.GetAttrs("qkv_hidden_sizes", qkv_hidden_sizes_).IsOK()) {
qkv_hidden_sizes_.clear();
@@ -68,6 +69,7 @@ class AttentionBase {
std::vector qkv_hidden_sizes_; // Q, K, V hidden sizes parsed from the qkv_hidden_sizes attribute.
bool require_same_hidden_size_; // whether the implementation supports different hidden sizes of Q/K/V.
bool past_present_share_buffer_; // whether or not the past (if used) and present tensor share the same buffer
+ float mask_filter_value_; // the value to be used for filtered out positions
};
} // namespace contrib
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_common.h b/onnxruntime/contrib_ops/cpu/bert/attention_common.h
index c313e4d284..a585538b57 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_common.h
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_common.h
@@ -32,6 +32,7 @@ struct AttentionParameters {
int num_heads;
bool is_unidirectional;
bool past_present_share_buffer;
+ float mask_filter_value;
AttentionMaskType mask_type;
};
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
index 2b01d1ce14..7dc3ca9873 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
@@ -123,7 +123,7 @@ class AttentionCPUBase : public AttentionBase {
// mask_data is nullptr when mask_index is nullptr and not unidirectional, otherwise its shape is BxSxT
if (mask_data != nullptr) {
PrepareMask(mask_index, mask_index_dims, mask_data,
- has_unidirectional, batch_size, sequence_length, past_sequence_length);
+ has_unidirectional, batch_size, sequence_length, past_sequence_length, mask_filter_value_);
} else { // no any mask
size_t bytes = static_cast(batch_size) * num_heads_ * sequence_length * total_sequence_length * sizeof(T);
memset(attention_probs, 0, bytes);
diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h
index 317751bbcd..76cf7ed8f8 100644
--- a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h
+++ b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h
@@ -67,7 +67,8 @@ void PrepareMask(const int32_t* mask_index,
bool is_unidirectional,
int batch_size,
int sequence_length,
- int past_sequence_length) {
+ int past_sequence_length,
+ float mask_filter_value) {
const int all_sequence_length = past_sequence_length + sequence_length;
// mask_data has been filled with 0, and its shape is BxSxT
@@ -78,17 +79,17 @@ void PrepareMask(const int32_t* mask_index,
ORT_NOT_IMPLEMENTED("4D mask in attention cpu kernel is not supported");
}
- // For 3D mask, convert values 0 to -10000.0, and 1 to 0.0, then apply unidirectional mask if any.
+ // For 3D mask, convert values 0 to mask_filter_value, and 1 to 0.0, then apply unidirectional mask if any.
if (nullptr != mask_index && mask_index_dims.size() == 3) {
for (int i = 0; i < batch_size * sequence_length * all_sequence_length; i++) {
- p_mask[i] = (mask_index[i] > 0) ? static_cast(0.0f) : static_cast(-10000.0f);
+ p_mask[i] = (mask_index[i] > 0) ? static_cast(0.0f) : static_cast(mask_filter_value);
}
if (is_unidirectional) {
for (int b_i = 0; b_i < batch_size; b_i++) {
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++) {
- p_mask[s_i * all_sequence_length + m_i] += static_cast(-10000.0f);
+ p_mask[s_i * all_sequence_length + m_i] += static_cast(mask_filter_value);
}
}
p_mask += static_cast(sequence_length) * all_sequence_length;
@@ -107,26 +108,26 @@ void PrepareMask(const int32_t* mask_index,
// TODO: mask_index can be used in softmax to save some calculation.
if (nullptr != mask_index) {
if (is_raw_attention_mask) {
- // Raw attention mask has value 0 or 1. Here we convert 0 to -10000.0, and 1 to 0.0.
+ // Raw attention mask has value 0 or 1. Here we convert 0 to mask_filter_value, and 1 to 0.0.
ptrdiff_t off = SafeInt(b_i) * all_sequence_length;
const int32_t* raw_mask = mask_index + off;
for (int m_i = 0; m_i < all_sequence_length; m_i++) {
- p_mask[m_i] = (raw_mask[m_i] > 0) ? static_cast(0.0f) : static_cast(-10000.0f);
+ p_mask[m_i] = (raw_mask[m_i] > 0) ? static_cast(0.0f) : static_cast(mask_filter_value);
}
} else {
// mask_index is 1D: (B) or (2B) => (Bx)T
- // Handle right-side padding: mask value at or after the end position will be -10000.0
+ // Handle right-side padding: mask value at or after the end position will be mask_filter_value
int end_position = mask_index[b_i];
for (int m_i = end_position; m_i < all_sequence_length; m_i++) {
- p_mask[m_i] = static_cast(-10000.0f);
+ p_mask[m_i] = static_cast(mask_filter_value);
}
- // Handle left-side padding: mask value before the start position will be -10000.0
+ // Handle left-side padding: mask value before the start position will be mask_filter_value
if (has_mask_start_position) {
int start_position = std::min(mask_index[b_i + batch_size], all_sequence_length);
for (int m_i = 0; m_i < start_position; m_i++) {
- p_mask[m_i] = static_cast(-10000.0f);
+ p_mask[m_i] = static_cast(mask_filter_value);
}
}
}
@@ -141,7 +142,7 @@ void PrepareMask(const int32_t* mask_index,
if (is_unidirectional) {
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++) {
- p_mask[s_i * all_sequence_length + m_i] += static_cast(-10000.0f);
+ p_mask[s_i * all_sequence_length + m_i] += static_cast(mask_filter_value);
}
}
}
diff --git a/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h
index a1b4156f72..ade1c527f4 100644
--- a/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h
+++ b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h
@@ -19,6 +19,7 @@ Status CheckInputs(const T* query,
const T* key_padding_mask,
void* parameters,
int num_heads,
+ float mask_filter_value,
int max_threads_per_block) {
// query (Q) : (B, S, D)
// key (K) : (B, L, D)
@@ -104,6 +105,7 @@ Status CheckInputs(const T* query,
output_parameters->num_heads = num_heads;
output_parameters->is_unidirectional = false;
output_parameters->past_present_share_buffer = false;
+ output_parameters->mask_filter_value = mask_filter_value;
output_parameters->mask_type = mask_type;
}
diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
index 8206549f38..0818feedbb 100644
--- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
@@ -244,6 +244,7 @@ Status QkvToContext(
const int num_heads = parameters.num_heads;
const int qk_head_size = parameters.head_size;
const int v_head_size = parameters.v_head_size;
+ const float mask_filter_value = parameters.mask_filter_value;
T* qkv = data.workspace;
const int batches = batch_size * num_heads;
@@ -427,8 +428,8 @@ Status QkvToContext(
ComputeSoftmaxWithRawMask(stream, total_sequence_length, sequence_length, batch_size, num_heads,
mask_index, nullptr, data.extra_add_qk, scratch1, scratch2,
parameters.is_unidirectional, rsqrt_head_size, mask_dimension,
- parameters.max_sequence_length,
- use_persistent_softmax, persistent_softmax_workspace));
+ parameters.max_sequence_length, use_persistent_softmax,
+ persistent_softmax_workspace, mask_filter_value));
} else if (nullptr != mask_index) { // 1d mask index
ORT_ENFORCE(mask_index_dims.size() == 1);
// mask_index has 1D shape: either (batch_size) or (2*batch_size). Only the later one has start postions.
@@ -471,6 +472,7 @@ Status DecoderQkvToContext(
const bool use_past,
const bool has_layer_state,
const bool has_key_padding_mask,
+ const float mask_filter_value,
const T* gemm_query_buffer,
const T* gemm_kv_buffer,
const bool* key_padding_mask,
@@ -588,7 +590,7 @@ Status DecoderQkvToContext(
ORT_RETURN_IF_ERROR(ComputeSoftmaxWithRawMask(stream, kv_sequence_length, sequence_length, batch_size, num_heads,
nullptr, key_padding_mask, add_before_softmax, scratch1, scratch2,
is_unidirectional, 1.0f, mask_dimension, max_sequence_length,
- false, nullptr));
+ false, nullptr, mask_filter_value));
} else {
ORT_RETURN_IF_ERROR(ComputeSoftmax(stream, kv_sequence_length, sequence_length, batch_size, num_heads,
add_before_softmax, scratch1, scratch2, is_unidirectional));
@@ -630,6 +632,7 @@ Status LaunchDecoderAttentionKernel(
const bool use_past,
const bool has_layer_state,
const bool has_key_padding_mask,
+ const float mask_filter_value,
const void* gemm_query_buffer,
const void* gemm_kv_buffer,
const bool* key_padding_mask,
@@ -655,6 +658,7 @@ Status LaunchDecoderAttentionKernel(
use_past,
has_layer_state,
has_key_padding_mask,
+ mask_filter_value,
reinterpret_cast(gemm_query_buffer),
reinterpret_cast(gemm_kv_buffer),
key_padding_mask,
@@ -680,6 +684,7 @@ Status LaunchDecoderAttentionKernel(
use_past,
has_layer_state,
has_key_padding_mask,
+ mask_filter_value,
reinterpret_cast(gemm_query_buffer),
reinterpret_cast(gemm_kv_buffer),
key_padding_mask,
diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h
index 79bd38e9e1..d09ad6b031 100644
--- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h
+++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h
@@ -71,6 +71,7 @@ Status LaunchDecoderAttentionKernel(
const bool use_past, // Whether use cache or not
const bool has_layer_state, // Whether output cache or not
const bool has_key_padding_mask, // Whether use key_padding_mask or not
+ const float mask_filter_value, // Mask filter value
const void* gemm_query_buffer, // Query buffer
const void* gemm_kv_buffer, // Key and value buffer
const bool* key_padding_mask, // Key padding mask
diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h
index 1ad908a31f..bf13719030 100644
--- a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h
+++ b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h
@@ -181,7 +181,8 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
const float rsqrt_head_size,
const int mask_dimension,
const int max_sequence_length,
- const bool skip_softmax) {
+ const bool skip_softmax,
+ const float mask_filter_value) {
using BlockReduce = cub::BlockReduce;
__shared__ typename BlockReduce::TempStorage tmp_storage;
@@ -203,7 +204,7 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
if (is_unidirectional) {
int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length.
if (threadIdx.x > from_index) {
- thread_data = -10000.0f;
+ thread_data = mask_filter_value;
}
}
@@ -221,7 +222,7 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
if (nullptr == key_padding_mask) {
const int& mask = attention_mask[mask_offset];
if (mask == 0)
- thread_data += -10000.0f;
+ thread_data += mask_filter_value;
} else {
const bool mask = key_padding_mask[mask_offset];
if (mask) {
@@ -388,12 +389,13 @@ __global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length,
const float rsqrt_head_size,
const int mask_dimension,
const int max_sequence_length,
- const bool skip_softmax) {
+ const bool skip_softmax,
+ const float mask_filter_value) {
SoftmaxWithRawMaskSmall(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, output,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- skip_softmax);
+ skip_softmax, mask_filter_value);
}
template
@@ -454,21 +456,22 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream,
template
Status ComputeSoftmaxWithRawMask(cudaStream_t stream,
- const int all_sequence_length,
- const int sequence_length,
- const int batch_size,
- const int num_heads,
- const int* attention_mask,
- const bool* key_padding_mask,
- const T* add_before_softmax,
- const T* input,
- T* output,
- const bool is_unidirectional,
- const float rsqrt_head_size,
- const int mask_dimension,
- const int max_sequence_length,
- const bool use_persistent_softmax,
- T* persistent_softmax_workspace) {
+ const int all_sequence_length,
+ const int sequence_length,
+ const int batch_size,
+ const int num_heads,
+ const int* attention_mask,
+ const bool* key_padding_mask,
+ const T* add_before_softmax,
+ const T* input,
+ T* output,
+ const bool is_unidirectional,
+ const float rsqrt_head_size,
+ const int mask_dimension,
+ const int max_sequence_length,
+ const bool use_persistent_softmax,
+ T* persistent_softmax_workspace,
+ const float mask_filter_value) {
const dim3 grid(sequence_length * num_heads, batch_size, 1);
T* out = use_persistent_softmax ? persistent_softmax_workspace : output;
@@ -478,42 +481,42 @@ Status ComputeSoftmaxWithRawMask(cudaStream_t stream,
<<>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 64) {
const int blockSize = 64;
SoftmaxWithRawMaskSmallKernel
<<>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 128) {
const int blockSize = 128;
SoftmaxWithRawMaskSmallKernel
<<>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 256) {
const int blockSize = 256;
SoftmaxWithRawMaskSmallKernel
<<>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 512) {
const int blockSize = 512;
SoftmaxWithRawMaskSmallKernel
<<>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 1024) {
const int blockSize = 1024;
SoftmaxWithRawMaskSmallKernel
<<>>(all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024.");
}
diff --git a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc
index b0ac4ad154..299c21df94 100644
--- a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc
@@ -167,6 +167,8 @@ DecoderAttention::DecoderAttention(const OpKernelInfo& info) : CudaKernel(inf
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);
}
template
@@ -375,6 +377,7 @@ Status DecoderAttention::ComputeInternal(OpKernelContext* context) const {
use_past_,
has_layer_state_,
has_key_padding_mask_,
+ mask_filter_value_,
nullptr == gemm_query_buffer_p ? nullptr : reinterpret_cast(gemm_query_buffer_p.get()),
nullptr == gemm_kv_buffer_p ? nullptr : reinterpret_cast(gemm_kv_buffer_p.get()),
nullptr == key_padding_mask ? nullptr : key_padding_mask->Data(),
diff --git a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.h b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.h
index 6f887f43db..5a3a6979b5 100644
--- a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.h
+++ b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.h
@@ -19,6 +19,7 @@ class DecoderAttention final : public CudaKernel {
private:
int num_heads_;
+ float mask_filter_value_;
};
} // namespace cuda
diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu
index 2eedc31acf..de3c3fb6ca 100644
--- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu
@@ -368,7 +368,7 @@ Status LaunchLongformerSoftmaxKernel(
const void* q, // transposed Q with shape (B, N, S, H)
const void* k, // transposed K with shape (B, N, S, H)
const void* v, // transposed V with shape (B, N, S, H)
- const void* attention_mask, // attention mask with shape (B, S), with value 0 not masked and -10000 masked.
+ const void* attention_mask, // attention mask with shape (B, S), with value 0 not masked and value of mask_filter_value.
int max_num_global, // maximum number of global tokens (G)
const bool compact_global_q, // whether global_q has shape (B, N, G, H) instead of (B, N, S, H)
const void* global_q, // Q for global tokens with shape (B, N, S, H).
@@ -842,7 +842,7 @@ Status LongformerQkvToContext(
const size_t element_size,
const T* input, // input for transpose
const T* bias, // bias to add to transposed input
- const T* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked.
+ const T* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 or torch.finfo(dtype).min masked.
const T* global_input, // global input for transpose
const T* global_bias, // bias to add to transposed global input
const int* global_attention, // global attention flags with shape (B, S), with value 0 for local and 1 for global.
diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu
index 21de67e52d..d610051c77 100644
--- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu
@@ -195,7 +195,7 @@ Status LaunchLongformerSoftmaxSimpleKernel(
const void* q, // transposed Q with shape (B, N, S, H)
const void* k, // transposed K with shape (B, N, S, H)
const void* v, // transposed V with shape (B, N, S, H)
- const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked.
+ const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 or torch.finfo(dtype).min masked.
const void* global_q, // Q for global tokens with shape (B, N, S, H)
const void* global_k, // K for global tokens with shape (B, N, S, H)
const void* global_v, // V for global tokens with shape (B, N, S, H)
diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h
index cbf86e6601..d0ff6fe731 100644
--- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h
+++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h
@@ -28,7 +28,7 @@ Status LaunchLongformerSoftmaxSimpleKernel(
const void* q, // transposed Q with shape (B, N, S, H)
const void* k, // transposed K with shape (B, N, S, H)
const void* v, // transposed V with shape (B, N, S, H)
- const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked.
+ const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 or torch.finfo(dtype).min masked.
const void* global_q, // Q for global tokens with shape (B, N, S, H)
const void* global_k, // K for global tokens with shape (B, N, S, H)
const void* global_v, // V for global tokens with shape (B, N, S, H)
diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
index 3d168ca453..0184f5ebe2 100644
--- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
@@ -35,6 +35,8 @@ MultiHeadAttention::MultiHeadAttention(const OpKernelInfo& info) : CudaKernel
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);
+
disable_fused_runner_ = sizeof(T) != 2 ||
ParseEnvironmentVariableWithDefault(attention::kDisableFusedAttention, false);
@@ -53,13 +55,14 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) const {
auto& device_prop = GetDeviceProp();
AttentionParameters parameters;
ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs(query,
- key,
- value,
- bias,
- key_padding_mask,
- ¶meters,
- num_heads_,
- device_prop.maxThreadsPerBlock));
+ key,
+ value,
+ bias,
+ key_padding_mask,
+ ¶meters,
+ num_heads_,
+ mask_filter_value_,
+ device_prop.maxThreadsPerBlock));
int sequence_length = parameters.sequence_length;
diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h
index 6eeecc697c..8c8ad966f5 100644
--- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h
+++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h
@@ -21,6 +21,7 @@ class MultiHeadAttention final : public CudaKernel {
protected:
int num_heads_; // number of attention heads
+ float mask_filter_value_;
bool disable_fused_runner_;
bool enable_flash_attention_;
mutable std::unique_ptr fused_fp16_runner_;
diff --git a/onnxruntime/contrib_ops/rocm/bert/attention.cc b/onnxruntime/contrib_ops/rocm/bert/attention.cc
index d91eec6f1b..756919834a 100644
--- a/onnxruntime/contrib_ops/rocm/bert/attention.cc
+++ b/onnxruntime/contrib_ops/rocm/bert/attention.cc
@@ -127,6 +127,7 @@ Status Attention::ComputeInternal(OpKernelContext* context) const {
reinterpret_cast(gemm_buffer.get()),
nullptr == mask_index ? nullptr : mask_index->Data(),
nullptr == mask_index ? gsl::span() : mask_index->Shape().GetDims(),
+ mask_filter_value_,
nullptr == past ? nullptr : past->Data(),
nullptr == extra_add_qk ? nullptr : extra_add_qk->Data(),
work_space.get(),
diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu b/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu
index bf561689cb..954a129be1 100644
--- a/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu
+++ b/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu
@@ -85,6 +85,7 @@ Status QkvToContext(
T* workspace,
const int* mask_index,
gsl::span mask_index_dims,
+ const float mask_filter_value,
bool is_unidirectional,
int past_sequence_length,
const T* past,
@@ -159,7 +160,7 @@ Status QkvToContext(
ComputeSoftmaxWithRawMask(stream, all_sequence_length, sequence_length, batch_size, num_heads,
mask_index, nullptr, extra_add_qk, scratch1, scratch2,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax, persistent_softmax_workspace));
+ use_persistent_softmax, persistent_softmax_workspace, mask_filter_value));
} else if (nullptr != mask_index) { // 1d mask index
ORT_ENFORCE(mask_index_dims.size() == 1);
// mask_index has 1D shape: either (batch_size) or (2*batch_size). Only the later one has start postions.
@@ -203,6 +204,7 @@ Status LaunchAttentionKernel(
const void* input,
const int* mask_index,
gsl::span mask_index_dims,
+ const float mask_filter_value,
const void* past,
const void* extra_add_qk,
void* workspace,
@@ -219,6 +221,7 @@ Status LaunchAttentionKernel(
reinterpret_cast<__half*>(workspace),
mask_index,
mask_index_dims,
+ mask_filter_value,
is_unidirectional,
past_sequence_length,
reinterpret_cast(past),
@@ -233,6 +236,7 @@ Status LaunchAttentionKernel(
reinterpret_cast(workspace),
mask_index,
mask_index_dims,
+ mask_filter_value,
is_unidirectional,
past_sequence_length,
reinterpret_cast(past),
@@ -258,6 +262,7 @@ Status DecoderQkvToContext(
const bool use_past,
const bool has_layer_state,
const bool has_key_padding_mask,
+ const float mask_filter_value,
const T* gemm_query_buffer,
const T* gemm_kv_buffer,
const bool* key_padding_mask,
@@ -372,7 +377,7 @@ Status DecoderQkvToContext(
if (has_key_padding_mask) {
ORT_RETURN_IF_ERROR(ComputeSoftmaxWithRawMask(stream, kv_sequence_length, sequence_length, batch_size,
num_heads, nullptr, key_padding_mask, nullptr, scratch1, scratch2,
- false, 1, 2, static_cast(0), false, nullptr));
+ false, 1, 2, static_cast(0), false, nullptr, mask_filter_value));
} else {
ORT_RETURN_IF_ERROR(ComputeSoftmax(stream, kv_sequence_length, sequence_length, batch_size,
num_heads, nullptr, scratch1, scratch2, false));
@@ -423,6 +428,7 @@ Status LaunchDecoderAttentionKernel(
const bool use_past,
const bool has_layer_state,
const bool has_key_padding_mask,
+ const float mask_filter_value,
const void* gemm_query_buffer,
const void* gemm_kv_buffer,
const bool* key_padding_mask,
@@ -449,6 +455,7 @@ Status LaunchDecoderAttentionKernel(
use_past,
has_layer_state,
has_key_padding_mask,
+ mask_filter_value,
reinterpret_cast(gemm_query_buffer),
reinterpret_cast(gemm_kv_buffer),
key_padding_mask,
@@ -475,6 +482,7 @@ Status LaunchDecoderAttentionKernel(
use_past,
has_layer_state,
has_key_padding_mask,
+ mask_filter_value,
reinterpret_cast(gemm_query_buffer),
reinterpret_cast(gemm_kv_buffer),
key_padding_mask,
diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_impl.h b/onnxruntime/contrib_ops/rocm/bert/attention_impl.h
index 6747f9135a..7db692083f 100644
--- a/onnxruntime/contrib_ops/rocm/bert/attention_impl.h
+++ b/onnxruntime/contrib_ops/rocm/bert/attention_impl.h
@@ -40,6 +40,7 @@ Status LaunchAttentionKernel(
const void* input, // Input tensor
const int* mask_index, // Attention mask raw data or index. NULL means no mask.
gsl::span mask_index_dims, // Mask index shape
+ const float mask_filter_value, // Mask value for filtered out positions
const void* past, // Past state input
const void* extra_add_qk, // Additional Add
void* workspace, // Temporary buffer
@@ -62,6 +63,7 @@ Status LaunchDecoderAttentionKernel(
const bool use_past, // Whether use cache or not
const bool has_layer_state, // Whether output cache or not
const bool has_key_padding_mask, // Whether use key_padding_mask or not
+ const float mask_filter_value, // Mask filter value
const void* gemm_query_buffer, // Query buffer
const void* gemm_kv_buffer, // Key and value buffer
const bool* key_padding_mask, // Key padding mask
diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h b/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h
index 80f5851164..27ecdf253e 100644
--- a/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h
+++ b/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h
@@ -186,7 +186,8 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
const float rsqrt_head_size,
const int mask_dimension,
const int max_sequence_length,
- const bool skip_softmax) {
+ const bool skip_softmax,
+ const float mask_filter_value) {
using BlockReduce = hipcub::BlockReduce;
__shared__ typename BlockReduce::TempStorage tmp_storage;
@@ -211,7 +212,7 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
if (is_unidirectional) {
int from_index = all_sequence_length - sequence_length + sequence_index; // offset in all sequence length.
if (threadIdx.x > from_index) {
- thread_data = -10000.0f;
+ thread_data = mask_filter_value;
}
}
@@ -229,7 +230,7 @@ __device__ inline void SoftmaxWithRawMaskSmall(const int all_sequence_length,
if (nullptr == key_padding_mask) {
const int& mask = attention_mask[mask_offset];
if (mask == 0)
- thread_data += -10000.0f;
+ thread_data += mask_filter_value;
} else {
const bool mask = key_padding_mask[mask_offset];
if (mask) {
@@ -385,12 +386,13 @@ __global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length,
const float rsqrt_head_size,
const int mask_dimension,
const int max_sequence_length,
- const bool skip_softmax) {
+ const bool skip_softmax,
+ const float mask_filter_value) {
SoftmaxWithRawMaskSmall(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, output,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- skip_softmax);
+ skip_softmax, mask_filter_value);
}
template
@@ -459,7 +461,8 @@ Status ComputeSoftmaxWithRawMask(hipStream_t stream,
const int mask_dimension,
const int max_sequence_length,
const bool use_persistent_softmax,
- T* persistent_softmax_workspace) {
+ T* persistent_softmax_workspace,
+ const float mask_filter_value) {
const dim3 grid(sequence_length * num_heads, batch_size, 1);
T* out = use_persistent_softmax ? persistent_softmax_workspace : output;
@@ -469,42 +472,42 @@ Status ComputeSoftmaxWithRawMask(hipStream_t stream,
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 64) {
const int blockSize = 64;
SoftmaxWithRawMaskSmallKernel<<>>(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 128) {
const int blockSize = 128;
SoftmaxWithRawMaskSmallKernel<<>>(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 256) {
const int blockSize = 256;
SoftmaxWithRawMaskSmallKernel<<>>(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 512) {
const int blockSize = 512;
SoftmaxWithRawMaskSmallKernel<<>>(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else if (all_sequence_length <= 1024) {
const int blockSize = 1024;
SoftmaxWithRawMaskSmallKernel<<>>(
all_sequence_length, sequence_length,
attention_mask, key_padding_mask, add_before_softmax, input, out,
is_unidirectional, rsqrt_head_size, mask_dimension, max_sequence_length,
- use_persistent_softmax);
+ use_persistent_softmax, mask_filter_value);
} else {
ORT_THROW("Attention ROCM operator does not support total sequence length > 1024.");
}
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index 5f5a8d77e1..aa23b1e296 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -199,6 +199,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"(2, batch_size, num_heads, max_sequence_length, head_size)",
AttributeProto::INT,
OPTIONAL_VALUE)
+ .Attr("mask_filter_value",
+ "The value to be filled in the attention mask. Default value is -10000.0f",
+ AttributeProto::FLOAT,
+ OPTIONAL_VALUE)
.Input(0,
"input",
"Input tensor with shape (batch_size, sequence_length, input_hidden_size)",
@@ -271,6 +275,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
OpSchema()
.SetDoc(MultiHeadAttention_ver1_doc)
.Attr("num_heads", "Number of attention heads", AttributeProto::INT)
+ .Attr("mask_filter_value", "The value to be filled in the attention mask. Default value is negative infinity",
+ AttributeProto::FLOAT, OPTIONAL_VALUE)
.Input(0,
"query",
"Query with shape (batch_size, sequence_length, hidden_size) when weights is not available.",
@@ -343,6 +349,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
OpSchema()
.SetDoc(Decoder_Attention_doc)
.Attr("num_heads", "Number of attention heads", AttributeProto::INT)
+ .Attr("mask_filter_value", "The value to be filled in the attention mask. Default value is negative infinity",
+ AttributeProto::FLOAT, OPTIONAL_VALUE)
.Input(0, "query", "3D input tensor with shape (sequence_length, batch_size, hidden_size), hidden_size = num_heads * head_size", "T")
.Input(1, "key", "3D input tensor with shape (total_sequence_length, batch_size, hidden_size)", "T")
.Input(2, "q_weight", "2D input tensor with shape (hidden_size, hidden_size)", "T")
diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py
index 8a403e3318..245ea9322a 100644
--- a/onnxruntime/python/tools/transformers/fusion_attention.py
+++ b/onnxruntime/python/tools/transformers/fusion_attention.py
@@ -101,6 +101,7 @@ class FusionAttention(Fusion):
self.num_heads = num_heads
self.attention_mask = attention_mask
self.use_multi_head_attention = use_multi_head_attention
+ self.mask_filter_value = None
# Flags to show warning only once
self.num_heads_warning = True
@@ -383,6 +384,9 @@ class FusionAttention(Fusion):
[helper.make_attribute("qkv_hidden_sizes", [qw_out_size, kw_out_size, vw_out_size])]
)
+ if self.mask_filter_value is not None:
+ attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
+
return attention_node
def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
@@ -569,6 +573,11 @@ class FusionAttention(Fusion):
logger.debug("fuse_attention: failed to match mask path")
return
+ if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul":
+ _, mul_val = self.model.get_constant_input(mask_nodes[0])
+ if mul_val != -10000:
+ self.mask_filter_value = mul_val
+
if matmul_v.input[0] == root_input and matmul_q.input[0] == root_input and matmul_k.input[0] == root_input:
mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0])
diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py
index 6c1d665d90..7fe3257950 100644
--- a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py
+++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py
@@ -21,6 +21,7 @@ class FusionGptAttentionPastBase(Fusion):
self.num_heads = num_heads
self.utils = FusionUtils(model)
self.casted_attention_mask = {} # map from name of attention mask to the name that casted to int32
+ self.mask_filter_value = None
def match_past_pattern_1(self, concat_k, concat_v, output_name_to_node):
# Pattern 1:
@@ -202,6 +203,9 @@ class FusionGptAttention(FusionGptAttentionPastBase):
]
)
+ if self.mask_filter_value is not None:
+ attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
+
matmul_node = helper.make_node(
"MatMul",
inputs=[attention_node_name + "_output", gemm_qkv.input[1]],
@@ -367,6 +371,12 @@ class FusionGptAttention(FusionGptAttentionPastBase):
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
return
+
+ if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul":
+ _, mul_val = self.model.get_constant_input(mask_nodes[0])
+ if mul_val != -10000:
+ self.mask_filter_value = -mul_val
+
else:
# New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0.
i, qk_nodes, _ = self.model.match_parent_paths(
@@ -408,6 +418,10 @@ class FusionGptAttention(FusionGptAttentionPastBase):
if input_mask_nodes is None:
logger.debug("fuse_attention: failed to match input attention mask path")
return
+ if len(input_mask_nodes) > 1 and input_mask_nodes[0].op_type == "Mul":
+ _, mul_val = self.model.get_constant_input(input_mask_nodes[0])
+ if mul_val != -10000:
+ self.mask_filter_value = mul_val
mask_nodes = self.model.match_parent_path(
where_qk,
diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention_megatron.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention_megatron.py
index a8939a5789..1c0b0b7074 100644
--- a/onnxruntime/python/tools/transformers/fusion_gpt_attention_megatron.py
+++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention_megatron.py
@@ -59,6 +59,8 @@ class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
helper.make_attribute("unidirectional", 0), # unidirectional shall not be ON for 4D attention mask
]
)
+ if self.mask_filter_value is not None:
+ attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
nodes_to_add = [attention_node]
self.nodes_to_add.extend(nodes_to_add)
@@ -80,6 +82,11 @@ class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
return None
(mul_mask, sub_mask, last_slice_mask, slice_mask) = mask_nodes
+ if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul":
+ _, mul_val = self.model.get_constant_input(mask_nodes[0])
+ if mul_val != 10000:
+ self.mask_filter_value = -mul_val
+
if mul_qk.input[1] != last_slice_mask.output[0]:
logger.debug("fuse_attention failed: mul_qk.input[1] != last_slice_mask.output[0]")
return None
diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py
index 183c44c07a..8176be523b 100644
--- a/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py
+++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py
@@ -23,6 +23,7 @@ class FusionGptAttentionNoPast(Fusion):
super().__init__(model, "Attention", ["LayerNormalization", "SkipLayerNormalization"], "without past")
# TODO: detect num_heads from graph like FusionAttention
self.num_heads = num_heads
+ self.mask_filter_value = None
def create_attention_node(self, gemm, gemm_qkv, input, output):
attention_node_name = self.model.create_node_name("Attention")
@@ -39,6 +40,8 @@ class FusionGptAttentionNoPast(Fusion):
helper.make_attribute("unidirectional", 1),
]
)
+ if self.mask_filter_value is not None:
+ attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
matmul_node = helper.make_node(
"MatMul",
@@ -176,6 +179,11 @@ class FusionGptAttentionNoPast(Fusion):
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
return
+ if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul":
+ _, mul_val = self.model.get_constant_input(mask_nodes[0])
+ if mul_val != -10000:
+ self.mask_filter_value = mul_val
+
else:
# New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0.
qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Where", "Div", "MatMul"], [0, 0, 1, 0])
diff --git a/onnxruntime/python/tools/transformers/models/longformer/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/longformer/convert_to_onnx.py
index 8e6196d470..7427b65a2b 100644
--- a/onnxruntime/python/tools/transformers/models/longformer/convert_to_onnx.py
+++ b/onnxruntime/python/tools/transformers/models/longformer/convert_to_onnx.py
@@ -185,6 +185,7 @@ def my_longformer_self_attention_forward_4(
)
input_mask = is_index_masked.float()
+ # TODO: The filtering value may be -10000.0 or -inf. Check the huggingface implementation.
input_mask = input_mask.masked_fill(is_index_masked, -10000.0)
# Yet another way to generate input_mask = torch.masked_fill(attention_mask, is_index_global_attn, 0.0)
diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc
index bfd569e711..0d8cef9686 100644
--- a/onnxruntime/test/contrib_ops/attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/attention_op_test.cc
@@ -76,6 +76,7 @@ static void RunAttentionTest(
tester.AddAttribute("num_heads", static_cast(number_of_heads));
tester.AddAttribute("unidirectional", static_cast(is_unidirectional ? 1 : 0));
tester.AddAttribute("past_present_share_buffer", static_cast(past_present_share_buffer ? 1 : 0));
+ tester.AddAttribute("mask_filter_value", static_cast(-10000.0f));
int32_t qkv_hidden_size_sum;
int32_t v_hidden_size;
@@ -3689,6 +3690,7 @@ TEST(AttentionTest, SharedPrepackedWeights) {
OpTester tester("Attention", 1, onnxruntime::kMSDomain);
tester.AddAttribute("num_heads", static_cast(number_of_heads));
tester.AddAttribute("unidirectional", static_cast(0));
+ tester.AddAttribute("mask_filter_value", static_cast(-10000.0f));
std::vector input_dims = {batch_size, sequence_length, hidden_size};
std::vector weights_dims = {hidden_size, 3 * hidden_size};
diff --git a/onnxruntime/test/contrib_ops/cross_attention_op_test.cc b/onnxruntime/test/contrib_ops/cross_attention_op_test.cc
index 62a150ec92..a8fa83bf99 100644
--- a/onnxruntime/test/contrib_ops/cross_attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/cross_attention_op_test.cc
@@ -41,6 +41,7 @@ static void RunMultiHeadAttentionTest(
if (enable_cpu || enable_cuda || enable_rocm) {
OpTester tester("MultiHeadAttention", 1, onnxruntime::kMSDomain);
tester.AddAttribute("num_heads", static_cast(number_of_heads));
+ tester.AddAttribute("mask_filter_value", static_cast(-10000.0f));
std::vector query_dims = {batch_size, sequence_length, hidden_size};
std::vector key_dims = {batch_size, kv_sequence_length, hidden_size};
diff --git a/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc b/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc
index f8d559216e..9e4c273606 100644
--- a/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc
@@ -42,6 +42,7 @@ static void RunAttentionTest(
if (enable_cpu || enable_cuda || enable_rocm) {
OpTester tester("DecoderAttention", 1, onnxruntime::kMSDomain);
tester.AddAttribute("num_heads", static_cast(num_heads));
+ tester.AddAttribute("mask_filter_value", static_cast(-10000.0f));
int head_size = hidden_size / num_heads;
std::vector query_dims = {sequence_length, batch_size, hidden_size};