saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add mask_filter in Attention related ops' attribute (#14274)

Browse source 
### Description
<!-- Describe your changes. -->


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

https://github.com/microsoft/onnxruntime/issues/12843

Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2023-01-17 12:28:11 -08:00 committed by GitHub
parent caa5900508
commit 2db57a53a3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
30 changed files with 161 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
docs/ContribOperators.md
View file

@ -123,6 +123,8 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Attributes
<dl>
<dt><tt>mask_filter_value</tt> : float</dt>
<dd>The value to be filled in the attention mask. Default value is -10000.0f</dd>
<dt><tt>num_heads</tt> : int (required)</dt>
<dd>Number of attention heads</dd>
<dt><tt>past_present_share_buffer</tt> : int</dt>
@ -967,6 +969,8 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Attributes
<dl>
<dt><tt>mask_filter_value</tt> : float</dt>
<dd>The value to be filled in the attention mask. Default value is negative infinity</dd>
<dt><tt>num_heads</tt> : int (required)</dt>
<dd>Number of attention heads</dd>
</dl>
@ -2120,6 +2124,8 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Attributes
<dl>
<dt><tt>mask_filter_value</tt> : float</dt>
<dd>The value to be filled in the attention mask. Default value is negative infinity</dd>
<dt><tt>num_heads</tt> : int (required)</dt>
<dd>Number of attention heads</dd>
</dl>

1
onnxruntime/contrib_ops/cpu/bert/attention_base.cc
View file

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

2
onnxruntime/contrib_ops/cpu/bert/attention_base.h
View file

@ -37,6 +37,7 @@ class AttentionBase {
num_heads_ = static_cast<int>(num_heads);
is_unidirectional_ = info.GetAttrOrDefault<int64_t>("unidirectional", 0) == 1;
mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
if (!info.GetAttrs<int64_t>("qkv_hidden_sizes", qkv_hidden_sizes_).IsOK()) {
qkv_hidden_sizes_.clear();
@ -68,6 +69,7 @@ class AttentionBase {
std::vector<int64_t> 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

1
onnxruntime/contrib_ops/cpu/bert/attention_common.h
View file

@ -32,6 +32,7 @@ struct AttentionParameters {
int num_heads;
bool is_unidirectional;
bool past_present_share_buffer;
float mask_filter_value;
AttentionMaskType mask_type;
};

2
onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
View file

@ -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<size_t>(batch_size) * num_heads_ * sequence_length * total_sequence_length * sizeof(T);
memset(attention_probs, 0, bytes);

23
onnxruntime/contrib_ops/cpu/bert/attention_helper.h
View file

@ -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<T>(0.0f) : static_cast<T>(-10000.0f);
p_mask[i] = (mask_index[i] > 0) ? static_cast<T>(0.0f) : static_cast<T>(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<T>(-10000.0f);
p_mask[s_i * all_sequence_length + m_i] += static_cast<T>(mask_filter_value);
}
}
p_mask += static_cast<size_t>(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<ptrdiff_t>(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<T>(0.0f) : static_cast<T>(-10000.0f);
p_mask[m_i] = (raw_mask[m_i] > 0) ? static_cast<T>(0.0f) : static_cast<T>(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<T>(-10000.0f);
p_mask[m_i] = static_cast<T>(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<T>(-10000.0f);
p_mask[m_i] = static_cast<T>(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<T>(-10000.0f);
p_mask[s_i * all_sequence_length + m_i] += static_cast<T>(mask_filter_value);
}
}
}

2
onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h
View file

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

11
onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
View file

@ -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<T>(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<T>(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<T>(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<const half*>(gemm_query_buffer),
reinterpret_cast<const half*>(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<const float*>(gemm_query_buffer),
reinterpret_cast<const float*>(gemm_kv_buffer),
key_padding_mask,

1
onnxruntime/contrib_ops/cuda/bert/attention_impl.h
View file

@ -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

55
onnxruntime/contrib_ops/cuda/bert/attention_softmax.h
View file

@ -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<float, TPB>;
__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<T, TPB>(
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 <typename T>
@ -454,21 +456,22 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream,
template <typename T>
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,
<<<grid, blockSize, 0, 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<T, blockSize>
<<<grid, blockSize, 0, 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 <= 128) {
const int blockSize = 128;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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 <= 256) {
const int blockSize = 256;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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 <= 512) {
const int blockSize = 512;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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 <= 1024) {
const int blockSize = 1024;
SoftmaxWithRawMaskSmallKernel<T, blockSize>
<<<grid, blockSize, 0, 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 {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024.");
}

3
onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc
View file

@ -167,6 +167,8 @@ DecoderAttention<T>::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<int>(num_heads);
mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
}
template <typename T>
@ -375,6 +377,7 @@ Status DecoderAttention<T>::ComputeInternal(OpKernelContext* context) const {
use_past_,
has_layer_state_,
has_key_padding_mask_,
mask_filter_value_,
nullptr == gemm_query_buffer_p ? nullptr : reinterpret_cast<const CudaT*>(gemm_query_buffer_p.get()),
nullptr == gemm_kv_buffer_p ? nullptr : reinterpret_cast<const CudaT*>(gemm_kv_buffer_p.get()),
nullptr == key_padding_mask ? nullptr : key_padding_mask->Data<bool>(),

1
onnxruntime/contrib_ops/cuda/bert/decoder_attention.h
View file

@ -19,6 +19,7 @@ class DecoderAttention final : public CudaKernel {
private:
int num_heads_;
float mask_filter_value_;
};
} // namespace cuda

4
onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu
View file

@ -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.

2
onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu
View file

@ -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)

2
onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h
View file

@ -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)

17
onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
View file

@ -35,6 +35,8 @@ MultiHeadAttention<T>::MultiHeadAttention(const OpKernelInfo& info) : CudaKernel
ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0);
num_heads_ = static_cast<int>(num_heads);
mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
disable_fused_runner_ = sizeof(T) != 2 ||
ParseEnvironmentVariableWithDefault<bool>(attention::kDisableFusedAttention, false);
@ -53,13 +55,14 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
auto& device_prop = GetDeviceProp();
AttentionParameters parameters;
ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs<Tensor>(query,
key,
value,
bias,
key_padding_mask,
&parameters,
num_heads_,
device_prop.maxThreadsPerBlock));
key,
value,
bias,
key_padding_mask,
&parameters,
num_heads_,
mask_filter_value_,
device_prop.maxThreadsPerBlock));
int sequence_length = parameters.sequence_length;

1
onnxruntime/contrib_ops/cuda/bert/multihead_attention.h
View file

@ -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<MHARunner> fused_fp16_runner_;

1
onnxruntime/contrib_ops/rocm/bert/attention.cc
View file

@ -127,6 +127,7 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
reinterpret_cast<const void*>(gemm_buffer.get()),
nullptr == mask_index ? nullptr : mask_index->Data<int>(),
nullptr == mask_index ? gsl::span<const int64_t>() : mask_index->Shape().GetDims(),
mask_filter_value_,
nullptr == past ? nullptr : past->Data<T>(),
nullptr == extra_add_qk ? nullptr : extra_add_qk->Data<T>(),
work_space.get(),

12
onnxruntime/contrib_ops/rocm/bert/attention_impl.cu
View file

@ -85,6 +85,7 @@ Status QkvToContext(
T* workspace,
const int* mask_index,
gsl::span<const int64_t> 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<T>(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<const int64_t> 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<const __half*>(past),
@ -233,6 +236,7 @@ Status LaunchAttentionKernel(
reinterpret_cast<float*>(workspace),
mask_index,
mask_index_dims,
mask_filter_value,
is_unidirectional,
past_sequence_length,
reinterpret_cast<const float*>(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<T>(stream, kv_sequence_length, sequence_length, batch_size,
num_heads, nullptr, key_padding_mask, nullptr, scratch1, scratch2,
false, 1, 2, static_cast<int>(0), false, nullptr));
false, 1, 2, static_cast<int>(0), false, nullptr, mask_filter_value));
} else {
ORT_RETURN_IF_ERROR(ComputeSoftmax<T>(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<const half*>(gemm_query_buffer),
reinterpret_cast<const half*>(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<const float*>(gemm_query_buffer),
reinterpret_cast<const float*>(gemm_kv_buffer),
key_padding_mask,

2
onnxruntime/contrib_ops/rocm/bert/attention_impl.h
View file

@ -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<const int64_t> 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

27
onnxruntime/contrib_ops/rocm/bert/attention_softmax.h
View file

@ -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<float, TPB>;
__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<T, TPB>(
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 <typename T>
@ -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<T, blockSize><<<grid, blockSize, 0, 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 <= 128) {
const int blockSize = 128;
SoftmaxWithRawMaskSmallKernel<T, blockSize><<<grid, blockSize, 0, 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 <= 256) {
const int blockSize = 256;
SoftmaxWithRawMaskSmallKernel<T, blockSize><<<grid, blockSize, 0, 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 <= 512) {
const int blockSize = 512;
SoftmaxWithRawMaskSmallKernel<T, blockSize><<<grid, blockSize, 0, 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 <= 1024) {
const int blockSize = 1024;
SoftmaxWithRawMaskSmallKernel<T, blockSize><<<grid, blockSize, 0, 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 {
ORT_THROW("Attention ROCM operator does not support total sequence length > 1024.");
}

8
onnxruntime/core/graph/contrib_ops/bert_defs.cc
View file

@ -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")

9
onnxruntime/python/tools/transformers/fusion_attention.py
View file

@ -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])

14
onnxruntime/python/tools/transformers/fusion_gpt_attention.py
View file

@ -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,

7
onnxruntime/python/tools/transformers/fusion_gpt_attention_megatron.py
View file

@ -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

8
onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py
View file

@ -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])

1
onnxruntime/python/tools/transformers/models/longformer/convert_to_onnx.py
View file

@ -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)

2
onnxruntime/test/contrib_ops/attention_op_test.cc
View file

@ -76,6 +76,7 @@ static void RunAttentionTest(
tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
tester.AddAttribute<int64_t>("unidirectional", static_cast<int64_t>(is_unidirectional ? 1 : 0));
tester.AddAttribute<int64_t>("past_present_share_buffer", static_cast<int64_t>(past_present_share_buffer ? 1 : 0));
tester.AddAttribute<float>("mask_filter_value", static_cast<float>(-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<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
tester.AddAttribute<int64_t>("unidirectional", static_cast<int64_t>(0));
tester.AddAttribute<float>("mask_filter_value", static_cast<float>(-10000.0f));
std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};

1
onnxruntime/test/contrib_ops/cross_attention_op_test.cc
View file

@ -41,6 +41,7 @@ static void RunMultiHeadAttentionTest(
if (enable_cpu || enable_cuda || enable_rocm) {
OpTester tester("MultiHeadAttention", 1, onnxruntime::kMSDomain);
tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
tester.AddAttribute<float>("mask_filter_value", static_cast<float>(-10000.0f));
std::vector<int64_t> query_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> key_dims = {batch_size, kv_sequence_length, hidden_size};

1
onnxruntime/test/contrib_ops/decoder_attention_op_test.cc
View file

@ -42,6 +42,7 @@ static void RunAttentionTest(
if (enable_cpu || enable_cuda || enable_rocm) {
OpTester tester("DecoderAttention", 1, onnxruntime::kMSDomain);
tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(num_heads));
tester.AddAttribute<float>("mask_filter_value", static_cast<float>(-10000.0f));
int head_size = hidden_size / num_heads;
std::vector<int64_t> query_dims = {sequence_length, batch_size, hidden_size};