Support attention bias with windowed CPU GQA

[Tianlei Wu (tianleiwu)]

Description

Enables the CPU GroupQueryAttention implementation to use attention_bias with sliding_window_cache. This is needed by speculative decoding with sliding-window attention, including calls that provide explicit position_ids for RoPE.

Summary of Changes

CPU GQA

File	Change
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc	Derives the absolute KV origin for each cache-relative batch and forwards it to attention implementations.
onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h	Applies the absolute bias-column offset in quantized, unquantized, flash, and non-flash paths, with per-batch fallback for differing origins.
onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h	Allows windowed attention bias only for callers that explicitly support absolute bias offsets; CUDA and WebGPU behavior is unchanged.

Tests

• Adds a deterministic CPU regression covering post-eviction bias indexing with explicit position_ids.
• Extends windowed-cache parity coverage to combine attention bias, non-default explicit position IDs, repeated eviction, and forced non-flash dispatch.

Testing

• cmake --build build/ci_cpu/Release --target onnxruntime_provider_test -j 8
• build/ci_cpu/Release/onnxruntime_provider_test --gtest_filter=GroupQueryAttentionTest.WindowedCacheAttentionBiasWithPositionIds_CPU
• ORT_GQA_DISABLE_FLASH_ATTENTION=1 build/ci_cpu/Release/onnxruntime_provider_test --gtest_filter=GroupQueryAttentionTest.WindowedCacheAttentionBiasWithPositionIds_CPU
• clang-format --dry-run --Werror onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc onnxruntime/test/contrib_ops/group_query_attention_op_test.cc
• python3 -m py_compile onnxruntime/test/python/transformers/test_gqa.py

Motivation and Context

A windowed KV cache stores resident rows in cache-relative coordinates after eviction, while attention_bias remains indexed by absolute sequence position. The previous validation rejected the combination to avoid silently reading incorrect bias columns. This change carries the per-batch absolute cache origin into the CPU attention paths so resident column zero maps to the correct absolute bias column. Explicit position_ids remain consumed by RoPE before the cache-relative transition.

Checklist

• Tests added/updated
• No breaking changes
• Documentation updated (not applicable; no public API change)

[Copilot]

Pull request overview

Adds absolute attention-bias indexing for CPU GQA with windowed KV caches.

Changes:

• Propagates per-batch absolute KV origins through all CPU attention paths.
• Keeps CUDA and WebGPU validation unchanged.
• Adds CPU regression and parity tests.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
group_query_attention.cc	Derives and forwards bias offsets.
gqa_attention_base.h	Applies offsets across attention implementations.
group_query_attention_helper.h	Adds CPU opt-in validation support.
group_query_attention_op_test.cc	Adds deterministic regression coverage.
test_gqa.py	Adds flash/non-flash parity coverage.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.