Clamp derived sequence lengths and KV-cache index in CUDA GroupQueryAttention

Merged
JI
4mo ago
+113-3

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 both total_seq_lens and the append offset past_seq_lens stay 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, where past_seq_lens points directly at the raw input and bypasses GetSequenceLengths.
  • Added NegativeSeqlensK_CacheAppend_NoOOB_CUDA regression test exercising the KV-cache append path with an out-of-range seqlens_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