Clamp derived sequence lengths and KV-cache index in CUDA GroupQueryAttention - #29240
Merged
jiafatom merged 5 commits intoJun 24, 2026
Merged
Conversation
…ttention On the CUDA EP the seqlens_k input is device-resident, so the host-side range check in the operator is skipped. Guard the device path instead: - GetSequenceLengths: clamp the negative case at the source so total_seq_lens and the KV append offset (past_seq_lens) stay non-negative. - UnpackRoPEAppend: make the KV-cache store bound two-sided, mirroring the position-index guard in the same kernel. This also covers the fast-decode path, where past_seq_lens points directly at the raw input. Add a CUDA regression test exercising the KV-cache append path with an out-of-range seqlens_k. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…_sequence_length Mark total_sequence_length as a graph initializer so shape inference computes present_kv = max(past_seq_len, total_sequence_length), matching the declared present_key/present_value outputs and resolving the model-build shape mismatch. Use a controlled minimal underflow seqlens_k=0 (with sequence_length=2) so the derived past length is -1 before neutralization, keeping all other internal lengths in valid range and the output finite. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The negative seqlens_k under-specifies the kv length relative to the query, so the resulting degenerate attention may legitimately yield non-finite values. The regression point is that the cache append and attention complete without indexing outside their buffers, so verify only that a correctly shaped result is produced rather than asserting finiteness. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
output_size derives from constexpr locals, so it is a constant expression usable inside the verifier lambda without capture. Capturing it triggers -Werror=unused-lambda-capture under the Android clang build. Reference it directly instead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens the CUDA (and hipified ROCm) GroupQueryAttention implementation against malformed seqlens_k values that can otherwise produce negative KV-cache append offsets and lead to out-of-bounds device writes. It brings the CUDA path closer to the CPU path’s safety guarantees by adding device-side clamping/guards and introduces a CUDA regression test for the scenario.
Changes:
- Clamp derived sequence lengths in the CUDA
GetSequenceLengthskernel so downstream offsets stay non-negative. - Add a two-sided bounds check for KV-cache writes in
UnpackRoPEAppend(cache_s >= 0 && cache_s < max_seqlen). - Add a CUDA regression test covering negative
seqlens_kin the KV-cache append path.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu |
Clamps total_len/past_seq_lens to non-negative values in GetSequenceLengths to prevent negative derived offsets. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh |
Adds a lower-bound check to KV-cache store indices to prevent negative-index stores. |
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc |
Adds a CUDA regression test intended to catch the negative-seqlens_k KV-cache append OOB write. |
Reword the verifier comment to state explicitly that the degenerate attention implied by the negative derived past length may yield non-finite outputs, which is expected and intentionally not asserted; the test confirms completion with a correctly shaped result and no out-of-bounds indexing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
apsonawane
approved these changes
Jun 24, 2026
This was referenced Jul 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
The CUDA
GroupQueryAttentionkernel derives a KV-cache append offset from theseqlens_kinput (past_seq_lens = (seqlens_k + 1) - sequence_length). On the CUDA EPseqlens_kis device-resident (onlytotal_sequence_lengthis a CPU input), so the host-side range validation in the operator/helper is skipped. The device kernelUnpackRoPEAppendthen guarded the cache store with only a one-sided upper bound (cache_s < max_seqlen), so an out-of-rangeseqlens_kcould produce a negative offset that is sign-extended into the cache-index arithmetic.The CPU operator already validates
seqlens_khost-side; this change brings the CUDA path to parity by guarding on the device.Changes
group_query_attention_impl.cu(GetSequenceLengths): clamp the negative case at the source so bothtotal_seq_lensand the append offsetpast_seq_lensstay non-negative for all downstream consumers.group_query_attention_qkv.cuh(UnpackRoPEAppend): make the KV-cache store bound two-sided (cache_s >= 0 && cache_s < max_seqlen), mirroring the existing position-index guard a few lines above. This also covers the fast-decode path, wherepast_seq_lenspoints directly at the raw input and bypassesGetSequenceLengths.NegativeSeqlensK_CacheAppend_NoOOB_CUDAregression test exercising the KV-cache append path with an out-of-rangeseqlens_k(CUDA-guarded; skips when CUDA EP is unavailable).Notes
total_sequence_length. The CUDA implementation is shared with ROCm via hipify.compute-sanitizer; the test asserts the run completes with finite outputs.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com