Skip to content

Support attention bias with windowed CPU GQA #32302

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
Tianlei Wu (tianleiwu) merged 2 commits into from
Aug 31, 2026
Merged

Support attention bias with windowed CPU GQA #32302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fa6a199
feat(cpu): support bias with windowed GQA
tianleiwu Aug 28, 2026
e705d16
fix(cpu): harden windowed GQA bias handling
tianleiwu Aug 28, 2026
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
107 changes: 73 additions & 34 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 @@ -137,6 +137,7 @@ class GQAAttentionBase {
const T* V, // V data with shape BxN_kvxSxH
const T* head_sink, // Head sink for smooth softmax, nullptr if not used
const Tensor* attention_bias, // Attention bias to add to QxK'
const int32_t* attention_bias_offsets, // Per-batch absolute KV origin, or nullptr
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 @@ -196,6 +197,7 @@ class GQAAttentionBase {

if (gqa_mlas_supported) {
ComputeAttentionProbs(static_cast<T*>(attention_probs), Q, k, head_sink, seqlens_k->Data<int32_t>(), attention_bias_data,
attention_bias_offsets,
batch_size, sequence_length, kv_sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache,
seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer,
past_present_share_buffer, packed_qkv, is_prompt, tp, allocator);
Expand All @@ -209,6 +211,7 @@ class GQAAttentionBase {
is_prompt, tp, allocator);
} else {
ComputeAttentionProbs(static_cast<float*>(attention_probs), Q, k, head_sink, seqlens_k->Data<int32_t>(), attention_bias_data,
attention_bias_offsets,
batch_size, sequence_length, kv_sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache,
seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer,
past_present_share_buffer, packed_qkv, is_prompt, tp, allocator);
Expand All @@ -229,16 +232,17 @@ class GQAAttentionBase {
// quantized present K/V (uint8_t storage).
template <typename T>
Status ApplyAttentionQuantized(
const T* Q, // Q data [B, N, S, H] BNSH
const T* K, // K data [B, N_kv, L, H] or nullptr for packed_qkv
const T* V, // V data [B, N_kv, L, H] or nullptr for packed_qkv
const T* head_sink, // smooth softmax sink per head, or nullptr
const Tensor* attention_bias, // additive bias or nullptr
const Tensor* past_key, // past K (uint8_t)
const Tensor* past_value, // past V (uint8_t)
Tensor* output, // output [B, S, N*H] T
Tensor* present_key, // present K (uint8_t)
Tensor* present_value, // present V (uint8_t)
const T* Q, // Q data [B, N, S, H] BNSH
const T* K, // K data [B, N_kv, L, H] or nullptr for packed_qkv
const T* V, // V data [B, N_kv, L, H] or nullptr for packed_qkv
const T* head_sink, // smooth softmax sink per head, or nullptr
const Tensor* attention_bias, // additive bias or nullptr
const int32_t* attention_bias_offsets, // per-batch absolute KV origin, or nullptr
const Tensor* past_key, // past K (uint8_t)
const Tensor* past_value, // past V (uint8_t)
Tensor* output, // output [B, S, N*H] T
Tensor* present_key, // present K (uint8_t)
Tensor* present_value, // present V (uint8_t)
Tensor* output_qk,
const Tensor* seqlens_k,
const float* k_scale,
Expand Down Expand Up @@ -440,6 +444,9 @@ class GQAAttentionBase {
if (attention_bias_shape[1] != 1) {
bias_offset += SafeInt<ptrdiff_t>(head_index) * bias_matrix_size;
}
if (attention_bias_offsets != nullptr) {
bias_offset += attention_bias_offsets[batch_index];
}
attn_bias = attention_bias_data + bias_offset;
}

Expand Down Expand Up @@ -596,15 +603,16 @@ class GQAAttentionBase {
// Uses online softmax with KV block tiling for reduced memory usage.
template <typename T>
Status ApplyAttentionQuantizedFlash(
const T* Q, // Q data [B, N, S, H] BNSH
const T* K, // K data [B, N_kv, L, H] or nullptr for packed_qkv
const T* V, // V data [B, N_kv, L, H] or nullptr for packed_qkv
const Tensor* attention_bias, // additive bias [B|1, N|1, S, T] or nullptr
const Tensor* past_key, // past K (uint8_t)
const Tensor* past_value, // past V (uint8_t)
Tensor* output, // output [B, S, N*H] T
Tensor* present_key, // present K (uint8_t)
Tensor* present_value, // present V (uint8_t)
const T* Q, // Q data [B, N, S, H] BNSH
const T* K, // K data [B, N_kv, L, H] or nullptr for packed_qkv
const T* V, // V data [B, N_kv, L, H] or nullptr for packed_qkv
const Tensor* attention_bias, // additive bias [B|1, N|1, S, T] or nullptr
const int32_t* attention_bias_offsets, // per-batch absolute KV origin, or nullptr
const Tensor* past_key, // past K (uint8_t)
const Tensor* past_value, // past V (uint8_t)
Tensor* output, // output [B, S, N*H] T
Tensor* present_key, // present K (uint8_t)
Tensor* present_value, // present V (uint8_t)
const Tensor* seqlens_k,
const float* k_scale,
const float* v_scale,
Expand Down Expand Up @@ -795,8 +803,14 @@ class GQAAttentionBase {
min_total_seqlen = std::min(min_total_seqlen, total_sl);
}
const bool ragged_seqlens = (max_total_seqlen != min_total_seqlen);
bool ragged_bias_offsets = false;
if (attention_bias_data != nullptr && attention_bias_offsets != nullptr) {
for (int b = 1; b < batch_size; ++b) {
ragged_bias_offsets = ragged_bias_offsets || attention_bias_offsets[b] != attention_bias_offsets[0];
}
}

if (ragged_seqlens) {
if (ragged_seqlens || ragged_bias_offsets) {
Comment thread
tianleiwu marked this conversation as resolved.
// Ragged seqlens: each batch item has its own total_seqlen (and therefore
// past_seqlen). Must use per-batch invocation regardless of past_key/prompt state.
common_past_seqlen = -1; // sentinel: per-batch
Expand All @@ -818,9 +832,11 @@ class GQAAttentionBase {
thread_count = std::max(thread_count, 1);

// Flash decoding: for decode (sequence_length==1), partition KV across threads
// to improve parallelism when batch*heads < thread_count.
// to improve parallelism when batch*heads < thread_count. The per-batch fallback
// uses the regular tiled kernel and therefore needs its larger scratch layout.
const int kv_chunk_count = (max_total_seqlen + kv_block_size - 1) / kv_block_size;
const bool use_flash_decoding = (sequence_length == 1 &&
common_past_seqlen >= 0 &&
batch_size * num_heads_ < thread_count &&
kv_chunk_count > 1);

Expand Down Expand Up @@ -885,10 +901,14 @@ class GQAAttentionBase {
args.v_scale = v_scale;
if constexpr (std::is_same_v<T, float>) {
args.output = output->MutableData<float>();
args.attention_bias = attention_bias_data;
args.attention_bias = attention_bias_data == nullptr
? nullptr
: attention_bias_data + (attention_bias_offsets == nullptr ? 0 : attention_bias_offsets[0]);
} else {
args.output_fp16 = output->MutableData<T>();
args.attention_bias_fp16 = attention_bias_data;
args.attention_bias_fp16 = attention_bias_data == nullptr
? nullptr
: attention_bias_data + (attention_bias_offsets == nullptr ? 0 : attention_bias_offsets[0]);
}
args.attention_bias_seqlen_stride = attention_bias_seqlen_stride;
args.attention_bias_broadcast_batch = attention_bias_broadcast_batch;
Expand Down Expand Up @@ -958,6 +978,9 @@ class GQAAttentionBase {
const size_t bias_head_extent = attention_bias_broadcast_head ? 1 : static_cast<size_t>(num_heads_);
batch_bias += static_cast<size_t>(SafeInt<size_t>(b) * bias_head_extent * sequence_length * attention_bias_seqlen_stride);
}
if (batch_bias != nullptr && attention_bias_offsets != nullptr) {
batch_bias += attention_bias_offsets[b];
}
if constexpr (std::is_same_v<T, float>) {
args.attention_bias = batch_bias;
} else {
Expand All @@ -980,15 +1003,16 @@ class GQAAttentionBase {
// Concatenates new K/V into the FP32 present cache, then runs the tiled
// online-softmax kernel MlasFlashAttentionGQA (QK^T + softmax + S*V fused).
Status ApplyAttentionFlash(
const float* Q, // Q data [B, N, S, H] BNSH
const float* K, // K data [B, N_kv, L, H] or nullptr for packed_qkv
const float* V, // V data [B, N_kv, L, H] or nullptr for packed_qkv
const Tensor* attention_bias, // additive bias [B|1, N|1, S, T] or nullptr
const Tensor* past_key, // past K (float)
const Tensor* past_value, // past V (float)
Tensor* output, // output [B, S, N*H] float
Tensor* present_key, // present K (float)
Tensor* present_value, // present V (float)
const float* Q, // Q data [B, N, S, H] BNSH
const float* K, // K data [B, N_kv, L, H] or nullptr for packed_qkv
const float* V, // V data [B, N_kv, L, H] or nullptr for packed_qkv
const Tensor* attention_bias, // additive bias [B|1, N|1, S, T] or nullptr
const int32_t* attention_bias_offsets, // per-batch absolute KV origin, or nullptr
const Tensor* past_key, // past K (float)
const Tensor* past_value, // past V (float)
Tensor* output, // output [B, S, N*H] float
Tensor* present_key, // present K (float)
Tensor* present_value, // present V (float)
const Tensor* seqlens_k,
GroupQueryAttentionParameters& parameters,
AllocatorPtr allocator,
Expand Down Expand Up @@ -1139,8 +1163,14 @@ class GQAAttentionBase {
min_total_seqlen = std::min(min_total_seqlen, total_sl);
}
const bool ragged_seqlens = (max_total_seqlen != min_total_seqlen);
bool ragged_bias_offsets = false;
if (attention_bias_data != nullptr && attention_bias_offsets != nullptr) {
for (int b = 1; b < batch_size; ++b) {
ragged_bias_offsets = ragged_bias_offsets || attention_bias_offsets[b] != attention_bias_offsets[0];
}
}

if (ragged_seqlens) {
if (ragged_seqlens || ragged_bias_offsets) {
common_past_seqlen = -1; // sentinel: per-batch
} else if (past_key == nullptr || is_prompt) {
common_past_seqlen = 0;
Expand Down Expand Up @@ -1229,7 +1259,9 @@ class GQAAttentionBase {
args.k_cache = present_key_data;
args.v_cache = present_value_data;
args.output = output->MutableData<float>();
args.attention_bias = attention_bias_data;
args.attention_bias = attention_bias_data == nullptr
? nullptr
: attention_bias_data + (attention_bias_offsets == nullptr ? 0 : attention_bias_offsets[0]);
args.attention_bias_seqlen_stride = attention_bias_seqlen_stride;
args.attention_bias_broadcast_batch = attention_bias_broadcast_batch;
args.attention_bias_broadcast_head = attention_bias_broadcast_head;
Expand Down Expand Up @@ -1283,6 +1315,9 @@ class GQAAttentionBase {
const size_t bias_head_extent = attention_bias_broadcast_head ? 1 : static_cast<size_t>(num_heads_);
batch_bias += static_cast<size_t>(b) * bias_head_extent * sequence_length * attention_bias_seqlen_stride;
}
if (batch_bias != nullptr && attention_bias_offsets != nullptr) {
batch_bias += attention_bias_offsets[b];
}
args.attention_bias = batch_bias;
args.attention_bias_seqlen_stride = attention_bias_seqlen_stride;
args.attention_bias_broadcast_batch = true; // batch offset handled above
Expand All @@ -1309,6 +1344,7 @@ class GQAAttentionBase {
const T* head_sink, // smooth softmax sink per head, or nullptr
const int32_t* seqlens_k, // total_sequence_length - 1 per batch
const T* attention_bias, // additive bias [B|1, N|1, S, T], or nullptr
const int32_t* attention_bias_offsets, // per-batch absolute KV origin, or nullptr
const size_t batch_size, // batch size
const size_t sequence_length, // Q sequence length (new tokens)
const size_t kv_sequence_length, // K/V input sequence length; 0 for shared KV
Expand Down Expand Up @@ -1408,6 +1444,9 @@ class GQAAttentionBase {
if (attention_bias_shape[1] != 1) {
attention_bias_offset += SafeInt<ptrdiff_t>(head_index) * attention_matrix_size;
}
if (attention_bias_offsets != nullptr) {
attention_bias_offset += attention_bias_offsets[batch_index];
}

attention_bias_thread = attention_bias + attention_bias_offset;
}
Expand Down
26 changes: 21 additions & 5 deletions onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,8 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids,
attention_bias,
head_sink,
parameters));
parameters,
/*support_windowed_attention_bias=*/true));

// Populate quantization fields in parameters.
parameters.k_quant_type = k_quant_type_;
Expand Down Expand Up @@ -318,6 +319,11 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
"seqlens_k[", b, "] = ", seqlens_k_data[b],
" is out of range [0, ", present_kv_seqlen, ")");
}
if (windowed && attention_bias != nullptr && seqlens_k_data[b] >= attention_bias->Shape()[3]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"seqlens_k[", b, "] = ", seqlens_k_data[b],
" exceeds the attention_bias sequence dimension ", attention_bias->Shape()[3], ".");
}
if ((windowed || !parameters.is_first_prompt) &&
static_cast<int64_t>(seqlens_k_data[b]) + 1 < sequence_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
Expand Down Expand Up @@ -507,9 +513,11 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
Tensor* attention_present_key = present_k;
Tensor* attention_present_value = present_v;
const Tensor* attention_seqlens_k = seqlens_k;
const int32_t* attention_bias_offsets = nullptr;

std::vector<WindowedStep> windowed_steps;
std::vector<int32_t> windowed_cache_seqlens;
std::vector<int32_t> windowed_attention_bias_offsets;
std::optional<Tensor> windowed_seqlens_tensor;
std::optional<Tensor> staged_key_tensor;
std::optional<Tensor> staged_value_tensor;
Expand All @@ -529,6 +537,13 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
PlanWindowedKvCache(seqlens_k->Data<int32_t>(), batch_size, sequence_length, windowed_capacity,
local_window_size_, windowed_steps, windowed_cache_seqlens, windowed_use_staging,
windowed_staged_capacity);
if (attention_bias != nullptr) {
windowed_attention_bias_offsets.resize(batch_size);
for (int b = 0; b < batch_size; ++b) {
windowed_attention_bias_offsets[b] = seqlens_k->Data<int32_t>()[b] - windowed_cache_seqlens[b];
}
Comment thread
tianleiwu marked this conversation as resolved.
attention_bias_offsets = windowed_attention_bias_offsets.data();
}

// Rows are moved verbatim, so the row size is taken from the cache tensor itself and covers
// fp32/fp16 as well as the int8 and (nibble-packed) int4 quantized layouts.
Expand Down Expand Up @@ -623,7 +638,7 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
if (use_flash) {
return ApplyAttentionQuantizedFlash(
q_rotary, k_data_q, v_data_q,
attention_bias,
attention_bias, attention_bias_offsets,
attention_past_key, attention_past_value,
output, attention_present_key, attention_present_value, attention_seqlens_k,
k_scale->Data<float>(), v_scale->Data<float>(),
Expand All @@ -632,7 +647,7 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {

return ApplyAttentionQuantized(
q_rotary, k_data_q, v_data_q, head_sink_data,
attention_bias, attention_past_key, attention_past_value,
attention_bias, attention_bias_offsets, attention_past_key, attention_past_value,
output, attention_present_key, attention_present_value, output_qk, attention_seqlens_k,
k_scale->Data<float>(), v_scale->Data<float>(),
mlas_quant_type, parameters, allocator, context);
Expand Down Expand Up @@ -661,14 +676,15 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
attention_present_key != nullptr && attention_present_value != nullptr;
if (use_flash) {
return ApplyAttentionFlash(q_rotary, k_data, v_data,
attention_bias, attention_past_key, attention_past_value,
attention_bias, attention_bias_offsets, attention_past_key, attention_past_value,
output, attention_present_key, attention_present_value, attention_seqlens_k,
parameters, allocator, context);
}
}

return ApplyAttention(q_rotary, k_data, v_data,
head_sink_data, attention_bias, attention_past_key, attention_past_value, output,
head_sink_data, attention_bias, attention_bias_offsets,
attention_past_key, attention_past_value, output,
attention_present_key, attention_present_value,
output_qk, attention_seqlens_k, parameters, allocator, context);
};
Expand Down
Loading
Loading