Skip to content

Add support for custom position ids and attention bias to GQA CPU operator #23944

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
derdeljan-msft merged 25 commits into from
Mar 14, 2025
Merged

Add support for custom position ids and attention bias to GQA CPU operator #23944

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
d7f5aa1
Scalar support for custom position ids and mask in GQA
derdeljan-msft Mar 4, 2025
15172c3
Vectorized attention mask application for fp32
derdeljan-msft Mar 6, 2025
d7eae78
Vectorized attention mask application for fp16
derdeljan-msft Mar 6, 2025
9d244dd
Add mask upscale to fp32 if the platform doesn't support fp16
derdeljan-msft Mar 6, 2025
8faee66
Fix typo in fp16 eltwise kernels
derdeljan-msft Mar 6, 2025
147d19b
Add validation for custom attention parameters
derdeljan-msft Mar 7, 2025
4b1262e
Add mlas unit test for eltwise kernels
derdeljan-msft Mar 7, 2025
f7a0788
Refactor python unit GQA tests
derdeljan-msft Mar 7, 2025
9dec056
Cleanup comments
derdeljan-msft Mar 7, 2025
5d23817
Fix CI pipeline errors
derdeljan-msft Mar 7, 2025
42e83d6
Apply suggestions from code review
derdeljan-msft Mar 7, 2025
bc0d69b
Fix docs pipeline build
derdeljan-msft Mar 8, 2025
ab60cbc
Fix docs pipeline build
derdeljan-msft Mar 8, 2025
4e0ca5c
Fix first batch of PR comments
derdeljan-msft Mar 10, 2025
949118f
Fix PR comments
derdeljan-msft Mar 13, 2025
62d39a5
Linter fix
derdeljan-msft Mar 13, 2025
0349678
Update attention_mask input description
derdeljan-msft Mar 13, 2025
0865ddb
Fix build break
derdeljan-msft Mar 13, 2025
55e09c9
Fix docs gen CI pipeline
derdeljan-msft Mar 13, 2025
e3bc338
Apply attention mask after softcap
derdeljan-msft Mar 13, 2025
757af32
Cleanup mlas eltwise module
derdeljan-msft Mar 13, 2025
0c268c9
Fix PR comments
derdeljan-msft Mar 13, 2025
c36a9cf
Fix position_ids handling for the first prompt
derdeljan-msft Mar 13, 2025
86a7737
Fix build break
derdeljan-msft Mar 13, 2025
56fe768
Fix PR comments and fix docs gen CI pipeline
derdeljan-msft Mar 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions cmake/onnxruntime_mlas.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ onnxruntime_add_static_library(onnxruntime_mlas
${MLAS_SRC_DIR}/activate.cpp
${MLAS_SRC_DIR}/logistic.cpp
${MLAS_SRC_DIR}/tanh.cpp
${MLAS_SRC_DIR}/eltwise.h
${MLAS_SRC_DIR}/eltwise.cpp
${MLAS_SRC_DIR}/erf.cpp
${MLAS_SRC_DIR}/compute.cpp
${MLAS_SRC_DIR}/quantize.cpp
Expand Down Expand Up @@ -101,6 +103,9 @@ function(setup_mlas_source_for_windows)
${MLAS_SRC_DIR}/softmax_kernel_neon.h
${MLAS_SRC_DIR}/softmax_kernel_neon.cpp
${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp
${MLAS_SRC_DIR}/eltwise_kernel_neon.h
${MLAS_SRC_DIR}/eltwise_kernel_neon.cpp
${MLAS_SRC_DIR}/eltwise_kernel_neon_fp16.cpp
)

set(mlas_platform_preprocess_srcs
Expand Down Expand Up @@ -387,6 +392,8 @@ else()
${MLAS_SRC_DIR}/hgemm_kernel_neon.cpp
${MLAS_SRC_DIR}/softmax_kernel_neon.h
${MLAS_SRC_DIR}/softmax_kernel_neon.cpp
${MLAS_SRC_DIR}/eltwise_kernel_neon.h
${MLAS_SRC_DIR}/eltwise_kernel_neon.cpp
)
set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp
PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+dotprod")
Expand All @@ -409,6 +416,7 @@ else()
${MLAS_SRC_DIR}/rotary_embedding_kernel_neon_fp16.cpp
${MLAS_SRC_DIR}/halfgemm_kernel_neon_fp16.cpp
${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp
${MLAS_SRC_DIR}/eltwise_kernel_neon_fp16.cpp
)
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/HalfGemmKernelNeon.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/QgemmS8S8KernelSmmla.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ")
Expand All @@ -423,6 +431,7 @@ else()
set_source_files_properties(${MLAS_SRC_DIR}/rotary_embedding_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/halfgemm_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/eltwise_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
endif()

if(ONNXRUNTIME_MLAS_MULTI_ARCH)
Expand Down
6 changes: 5 additions & 1 deletion docs/ContribOperators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2551,7 +2551,7 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Softcap value for attention weights. Default value is 0.</dd>
</dl>

#### Inputs (7 - 9)
#### Inputs (7 - 11)

<dl>
<dt><tt>query</tt> : T</dt>
Expand All @@ -2572,6 +2572,10 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>2D tensor with shape (max_sequence_length, head_size / 2).</dd>
<dt><tt>sin_cache</tt> (optional) : T</dt>
<dd>2D tensor with shape (max_sequence_length, head_size / 2).</dd>
<dt><tt>custom_pos_ids</tt> (optional) : tensor(int64)</dt>
<dd>2D tensor with shape (batch_size, sequence_length).</dd>
<dt><tt>custom_causal_attention_mask</tt> (optional) : T</dt>
<dd>3D tensor with shape (batch_size, sequence_length, total_sequence_length)</dd>
</dl>

#### Outputs
Expand Down
6 changes: 3 additions & 3 deletions docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,7 +520,7 @@ Do not modify directly.*
|Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
|GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *in* attention_mask:**I**<br> *out* sequences:**I**|1+|**T** = tensor(float)|
|GridSample|*in* X:**T1**<br> *in* Grid:**T1**<br> *out* Y:**T2**|1+|**T1** = tensor(float)<br/> **T2** = tensor(float)|
|GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
|GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* custom_pos_ids:**tensor(int64)**<br> *in* custom_causal_attention_mask:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
|Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|MatMulBnb4|*in* A:**T1**<br> *in* B:**T2**<br> *in* absmax:**T1**<br> *out* Y:**T1**|1+|**T1** = tensor(float)<br/> **T2** = tensor(uint8)|
|MatMulFpQ4|*in* A:**T1**<br> *in* B:**T2**<br> *in* B_shape:**T3**<br> *out* Y:**T1**|1+|**T1** = tensor(float)<br/> **T2** = tensor(uint8)<br/> **T3** = tensor(int64)|
Expand Down Expand Up @@ -922,7 +922,7 @@ Do not modify directly.*
|GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *in* attention_mask:**I**<br> *out* sequences:**I**|1+|**T** = tensor(float), tensor(float16)|
|GridSample|*in* X:**T1**<br> *in* Grid:**T1**<br> *out* Y:**T2**|1+|**T1** = tensor(float)<br/> **T2** = tensor(float)|
|GroupNorm|*in* X:**T**<br> *in* gamma:**M**<br> *in* beta:**M**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float16)|
+|GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* custom_pos_ids:**tensor(int64)**<br> *in* custom_causal_attention_mask:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float16)|
|Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|Irfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|LongformerAttention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask:**T**<br> *in* global_weight:**T**<br> *in* global_bias:**T**<br> *in* global:**G**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
Expand Down Expand Up @@ -1399,7 +1399,7 @@ Do not modify directly.*
|FusedMatMulActivation|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|GroupNorm|*in* X:**T**<br> *in* gamma:**M**<br> *in* beta:**M**<br> *out* Y:**T**|1+|**M** = tensor(float), tensor(float16)<br/> **T** = tensor(float), tensor(float16)|
|GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* custom_pos_ids:**tensor(int64)**<br> *in* custom_causal_attention_mask:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
|MatMulIntegerToFloat|*in* A:**T1**<br> *in* B:**T2**<br> *in* a_scale:**T3**<br> *in* b_scale:**T3**<br> *in* a_zero_point:**T1**<br> *in* b_zero_point:**T2**<br> *in* bias:**T3**<br> *out* Y:**T3**|1+|**T1** = tensor(int8), tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(float), tensor(float16)|
|MatMulNBits|*in* A:**T1**<br> *in* B:**T2**<br> *in* scales:**T1**<br> *in* zero_points:**T3**<br> *in* g_idx:**T4**<br> *in* bias:**T1**<br> *out* Y:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)|
|MultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* bias:**T**<br> *in* key_padding_mask:**M**<br> *in* attention_bias:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
Expand Down
5 changes: 5 additions & 0 deletions onnxruntime/contrib_ops/cpu/bert/attention_helper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,11 @@ void ComputeAttentionSoftcapInplace(T* scores, int sequence_length, T softcap) {
MlasComputeSoftcap(scores, scores, sequence_length, softcap);
}

template <typename T>
void ApplyAttentionMask(T* softmax_logits, const T* attention_mask, int N) {
MlasEltwiseAdd(softmax_logits, attention_mask, softmax_logits, N);
}

template <typename T>
void PrepareMask(const int32_t* mask_index,
gsl::span<const int64_t> mask_index_dims,
Expand Down
80 changes: 54 additions & 26 deletions onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ class GQAAttentionBase {
Status ApplyAttention(const T* Q, // Q data with shape BxNxSxH
const T* K, // K data with shape BxN_kvxSxH
const T* V, // V data with shape BxN_kvxSxH
const Tensor* attention_mask, // Causal attention mask to apply before
const Tensor* past_key, // past K input tensor (if not using past state)
const Tensor* past_value, // past V input tensor (if not using past state)
Tensor* output, // output tensor
Expand Down Expand Up @@ -87,13 +88,17 @@ class GQAAttentionBase {
const T* past_value_data = past_value != nullptr ? past_value->Data<T>() : nullptr;
T* present_value_data = present_value != nullptr ? present_value->MutableData<T>() : nullptr;

const T* attention_mask_data = attention_mask != nullptr ? attention_mask->Data<T>() : nullptr;
const size_t attention_mask_total_seqlen =
attention_mask != nullptr ? static_cast<size_t>(attention_mask->Shape()[2]) : static_cast<size_t>(0);

bool past_present_share_buffer = past_key_data == present_key_data && past_value_data == present_value_data;

const T* k = packed_qkv ? Q + num_heads_ * sequence_length * head_size : K;

if (gqa_mlas_supported) {
ComputeAttentionProbs(static_cast<T*>(attention_probs), Q, k, seqlens_k->Data<int32_t>(), batch_size,
sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data,
ComputeAttentionProbs(static_cast<T*>(attention_probs), Q, k, seqlens_k->Data<int32_t>(), attention_mask_data, batch_size,
sequence_length, attention_mask_total_seqlen, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data,
present_key_data, past_present_share_buffer, packed_qkv, is_prompt, tp, allocator);

// Compute the attentionScore * Value: out(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v)
Expand All @@ -104,9 +109,10 @@ class GQAAttentionBase {
hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv,
is_prompt, tp, allocator);
} else {
ComputeAttentionProbs(static_cast<float*>(attention_probs), Q, k, seqlens_k->Data<int32_t>(), batch_size,
sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data,
present_key_data, past_present_share_buffer, packed_qkv, is_prompt, tp, allocator);
ComputeAttentionProbs(static_cast<float*>(attention_probs), Q, k, seqlens_k->Data<int32_t>(), attention_mask_data,
batch_size, sequence_length, attention_mask_total_seqlen, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size,
past_key_data, present_key_data, past_present_share_buffer, packed_qkv, is_prompt, tp,
allocator);

// Compute the attentionScore * Value: out(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v)
const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V;
Expand All @@ -130,8 +136,10 @@ class GQAAttentionBase {
const T* Q, // Q data. Its size is BxNxSxH
const T* K, // k data. Its size is BxNxLxH
const int32_t* seqlens_k, // total - 1 sequence lengths tensor
const T* attention_mask, // optional causal attention mask
const size_t batch_size, // batch size of self-attention
const size_t sequence_length, // sequence length of self-attention (S)
const size_t attention_mask_total_seqlen, // max total seqlen in batch used for attention last dim
const size_t past_buffer_sequence_length, // sequence length of past state
const size_t present_buffer_sequence_length, // sequence length of present state
const size_t head_size, // head size of self-attention
Expand Down Expand Up @@ -189,6 +197,9 @@ class GQAAttentionBase {
const ptrdiff_t output_offset = SafeInt<ptrdiff_t>(i) * sequence_length * present_buffer_sequence_length;
U* output = attention_probs + output_offset;

const ptrdiff_t attention_mask_offset = SafeInt<ptrdiff_t>(batch_index) * sequence_length * attention_mask_total_seqlen;
const T* attention_mask_batch = attention_mask != nullptr ? attention_mask + attention_mask_offset : nullptr;

const T* k;
if (packed_qkv) {
k = K + packed_batch_stride * batch_index + kv_input_chunk_length * (head_index / kv_num_heads_factor);
Expand Down Expand Up @@ -242,35 +253,48 @@ class GQAAttentionBase {
U* output_softmax = output;
for (size_t seq = 0; seq < sequence_length; seq++) {
size_t seq_causal_length = past_seqlen + seq + 1;
if (local_window_size_ > 0 && seq_causal_length > static_cast<size_t>(local_window_size_) + 1) {

const bool should_apply_local_window = local_window_size_ > 0 &&
seq_causal_length > static_cast<size_t>(local_window_size_) + 1;

const size_t start_offset = should_apply_local_window ? seq_causal_length - local_window_size_ - 1 : 0;
const size_t window_size = should_apply_local_window ? local_window_size_ + 1 : seq_causal_length;

// Apply custom attention mask if there is any
if (attention_mask_batch != nullptr) {
if constexpr (std::is_same_v<U, T>) {
ApplyAttentionMask(output_softmax + start_offset, attention_mask_batch + start_offset,
static_cast<int>(window_size));
} else {
size_t bytes = window_size * sizeof(float);
auto attention_mask_batch_fp32 = static_cast<float*>(allocator->Alloc(bytes));
BufferUniquePtr scratch_buffer(attention_mask_batch_fp32, BufferDeleter(allocator));

MlasConvertHalfToFloatBuffer(attention_mask_batch + start_offset, attention_mask_batch_fp32, window_size);
ApplyAttentionMask(output_softmax, attention_mask_batch_fp32, static_cast<int>(window_size));
}
}
Comment thread
derdeljan-msft marked this conversation as resolved.
Outdated

// Mask everything before local window, if local window should be applied
if (should_apply_local_window) {
for (size_t total_seq_id = 0; total_seq_id < seq_causal_length - local_window_size_ - 1; total_seq_id++) {
if constexpr (std::is_same<U, float>::value) {
output_softmax[total_seq_id] = 0.f;
} else {
output_softmax[total_seq_id] = MLFloat16::FromBits(static_cast<uint16_t>(0));
}
}
if (softcap_ > 0.f) {
ComputeAttentionSoftcapInplace(output_softmax + seq_causal_length - local_window_size_ - 1,
local_window_size_ + 1, static_cast<U>(softcap_));
}
if (use_smooth_softmax_) {
ComputeSmoothSoftmaxInplace(output_softmax + seq_causal_length - local_window_size_ - 1, 1,
local_window_size_ + 1, nullptr);
} else {
ComputeAttentionSoftmaxInplace(output_softmax + seq_causal_length - local_window_size_ - 1, 1,
local_window_size_ + 1, nullptr);
}
}

// Calculate softmax
Comment thread
derdeljan-msft marked this conversation as resolved.
Outdated
if (softcap_ > 0.f) {
Comment thread
derdeljan-msft marked this conversation as resolved.
ComputeAttentionSoftcapInplace(output_softmax + start_offset, static_cast<int>(window_size),
static_cast<U>(softcap_));
}
if (use_smooth_softmax_) {
ComputeSmoothSoftmaxInplace(output_softmax + start_offset, 1, static_cast<int>(window_size), nullptr);
} else {
if (softcap_ > 0.f) {
ComputeAttentionSoftcapInplace(output_softmax, static_cast<int>(seq_causal_length),
static_cast<U>(softcap_));
}
if (use_smooth_softmax_) {
ComputeSmoothSoftmaxInplace(output_softmax, 1, static_cast<int>(seq_causal_length), nullptr);
} else {
ComputeAttentionSoftmaxInplace(output_softmax, 1, static_cast<int>(seq_causal_length), nullptr);
}
ComputeAttentionSoftmaxInplace(output_softmax + start_offset, 1, static_cast<int>(window_size), nullptr);
}

// set causal [seq_causal_length, total_seqlen) to 0.f
Expand All @@ -283,6 +307,10 @@ class GQAAttentionBase {
}

output_softmax += present_buffer_sequence_length;

if (attention_mask_batch != nullptr) {
attention_mask_batch += attention_mask_total_seqlen;
}
}
}
});
Expand Down
Loading