Fix GQA out-of-bounds read in past KV buffer (CPU)
Description
Fixes an out-of-bounds (OOB) read in the CPU GroupQueryAttention (GQA) operator. During token generation (decode), ConcatStateChunkGQA copies seqlens_k + 1 - sequence_length rows out of the past key/value buffers. The existing validation only bounded the present buffer write (seqlens_k < present_kv_seqlen), but the present buffer can be larger than the past buffer when total_sequence_length exceeds the past sequence length. A large seqlens_k combined with a small past buffer therefore read past the end of the past key/value tensors.
Motivation and Context
present_kv_seqlen = max(total_sequence_length, past_sequence_length), so the pre-existing seqlens_k < present_kv_seqlen check does not bound the past-side read when total_sequence_length > past_sequence_length. With a crafted seqlens_k, the decode path in ConcatStateChunkGQA reads seqlens_k + 1 - sequence_length rows from the smaller past buffer, causing an OOB read.
Key Changes
| File | Change |
|---|---|
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc | Add a per-batch bound in the seqlens_k validation block: for the decode case (past_kv_seqlen > 0 && kv_sequence_length != 0 && !is_first_prompt), reject when seqlens_k + 1 - sequence_length > past_kv_seqlen. |
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | Add regression test SeqlensKExceedsPastBuffer_OOBRead exercising a large seqlens_k against a small past buffer. |
Shared KV (kv_sequence_length == 0) appends no new KV and its past read is already bounded by the present-buffer check together with the total_sequence_length <= seqlen_past_kv_cache enforcement in the apply-attention paths, so it needs no additional check.
Testing Notes
- Build and run the GQA tests:
cmake --build build/Linux/Debug --target onnxruntime_provider_test -j$(nproc) ./build/Linux/Debug/onnxruntime_provider_test --gtest_filter="*GroupQueryAttention*" - All 41 CPU GQA tests pass, including the new
SeqlensKExceedsPastBuffer_OOBReadregression and the existing shared-KV CPU cases. lintrunnerreports no issues on the changed files.