Clamp derived sequence lengths and KV-cache index in CUDA GroupQueryAttention
Merged
JI
4mo ago+113+113+012345678901234567890123456789-3-3-0123456789
Description
The CUDA GroupQueryAttention kernel derives a KV-cache append offset from the seqlens_k input (past_seq_lens = (seqlens_k + 1) - sequence_length). On the CUDA EP seqlens_k is device-resident (only total_sequence_length is a CPU input), so the host-side range validation in the operator/helper is skipped. The device kernel UnpackRoPEAppend then guarded the cache store with only a one-sided upper bound (cache_s < max_seqlen), so an out-of-range seqlens_k could produce a negative offset that is sign-extended into the cache-index arithmetic.
The CPU operator already validates seqlens_k host-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.- Added
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
- The two-sided guard matches the pattern introduced for the rotary position index in #27597.
- CPU is unaffected (already validated host-side); WebGPU relies on the CPU-validated
total_sequence_length. The CUDA implementation is shared with ROCm via hipify. - The regression is a device-memory write best observed under
compute-sanitizer; the test asserts the run completes with finite outputs.
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com