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};