onnxruntime

3fc465dc - [CPU] Support right-padded prefill chunks in GroupQueryAttention (#33505)

[CPU] Support right-padded prefill chunks in GroupQueryAttention (#33505) ### Description <!-- Describe your changes. --> Implement the optional `query_seqlens` input (index 16) of `com.microsoft::GroupQueryAttention` in the CPU execution provider. The input carries the number of valid, non-padded rows of a right-padded query/key/value chunk per batch entry, so fixed-shape chunked prefill works on CPU with a full-length or a windowed KV cache (`sliding_window_cache=1`), with float, fp16 and int8-quantized KV, on both the MLAS flash and the unfused paths. This PR is stacked on #33503, which adds the input to the schema, the shared validation and the CUDA implementation, and makes the CPU provider reject the input until this PR lands. The input's semantics are specified there and are only summarized below. This revision supersedes the earlier one of this PR, which inferred the valid length of a chunk from `total_sequence_length` (`ResolveGQAStepLengths`). That heuristic is removed. **Stack:** the base of this PR is the branch of #33503, so "Files changed" shows only the CPU changes. Once #33503 merges, this PR is retargeted to `main` and rebased onto it. ORT's build, test and CodeQL workflows run only for pull requests that target `main` or `rel-*` (the lint workflow runs for any base), so this PR gets them after that retarget. The local results below were produced on the combined branch. #### CPU behavior For batch entry `b`: `S` is the padded `sequence_length`, `T = seqlens_k[b] + 1`, `Q = query_seqlens[b]` and `P = T - Q`. - Only the `Q` valid rows are appended to the cache, for full-length and windowed caches alike. Query row `s` sits at absolute position `P + s` for causal masking, `local_window_size` and the implicit rotary positions. The padded rows of the implicit rotary positions get position 0, which is always a valid rotary-cache index. - Every contract violation is rejected with `INVALID_ARGUMENT` (CUDA clamps device-side values instead): a shape other than `(batch_size)`, `Q` outside `[1, min(S, T)]`, `T > total_sequence_length`, a missing past, key/value that are not a full chunk of `S` rows (including the shared-KV case `kv_sequence_length == 0`), `S > 1` with `batch_size > 1`, `P` above the sequence dimension of `past_key` (a cache may grow between steps, so the past buffers can be shorter than the present buffers), `P + S > C` for a cache that is not windowed (`C` is the sequence dimension of `present_key`), `causal=0`, and missing `present_key`/`present_value` outputs. The outputs are required for every step, also for `S == 1`, because the contract of the input does not depend on the chunk length. - `S == 1` treats the input as absent (it is still validated), so a model can feed it at every step. - With the input, an explicit `position_ids` is honored for every row, also on the first chunk, as on any later chunk. Without it, a first prompt reads only `position_ids[0]`. The entries of padded rows must still be valid rotary-cache indices. - A windowed cache accepts `C >= local_window_size` and `attention_bias` together with padded chunks. CUDA is limited to `C == local_window_size` and no `attention_bias`. - `output_qk`, when requested, keeps its shape `(batch_size, num_heads, S, total_sequence_length)`. The columns `[T, total_sequence_length)` of a batch entry are zero. They used to hold uninitialized scratch memory, which already happened for ragged batches without the new input, because the kernel never writes them. - Without the input the behavior is unchanged except for the `output_qk` columns above: the helper below falls back to the previous per-batch logic, and the one changed constant (the padded-row rotary position) is unreachable then. #### Key changes - **`gqa_attention_base.h`**: the per-batch lengths of a step (total, past rows to copy, causal past, new rows to append) were open-coded at 8 sites. They now come from `ResolveBatchStepLengths`, which takes an optional per-batch past length and reproduces the previous logic when none is given. `ConcatStateChunkGQA` and its quantized variant append `valid_new_rows` instead of `kv_sequence_length`, so padded rows never overwrite valid entries or run past the capacity. The flash paths take their past length per batch entry. `WriteOutputQKHeadChunk` copies only the columns the kernel wrote and zero-fills the rest of the `output_qk` row. - **`group_query_attention.cc`**: reads input 16 and performs the checks above. `PlanWindowedKvCache` sizes eviction and the staging cache with `Q` instead of `S`, and also returns the kept past length per entry, which is passed down as `attention_past_seqlens`. The default rotary positions use `Q`. - **`group_query_attention_helper.h`**: the earlier `ResolveGQAStepLengths` is removed. The file is identical to the #33503 branch, which holds the validation shared by CPU and CUDA. - **Schema and docs**: input 16 is documented as honored by the CPU and CUDA providers (other providers must reject the node). `docs/ContribOperators.md` is regenerated and the "Right-Padded Prefill Chunks" section of `docs/contrib_ops/cpu/gqa.md` is rewritten for `query_seqlens`. Comments in the MIGraphX, OpenVINO and JS providers that said only CUDA implements the input now say CPU and CUDA. - **Tests**: - C++ (`group_query_attention_op_test.cc`): `QuerySeqlensRightPaddedChunk_CPU` and `QuerySeqlensRotaryRightPaddedChunkMatchesUnpadded_CPU` compare a padded chunk with the same call on an unpadded chunk (float and fp16, flash on and off, packed and unpacked QKV, and the appended cache rows). `WindowedCacheRightPaddedChunkMatchesUnpadded_CPU` does the same for a windowed cache over many prefill and valid lengths. Every rejection above has a negative test, including a missing `present_key`/`present_value` for `S == 1` and a growing cache whose past is longer than the bound past buffers. A WebGPU-only test checks that the WebGPU provider rejects a node with the input. It is compiled only with `USE_WEBGPU`, which my build does not enable. - Python (`TestGQAWindowedKvCacheCpu`, shared with the CUDA suite): padded chunks run through `query_seqlens` with a float and an int8 cache, flash and non-flash, `C == W` and `C == W + 8`, a capacity-valued `total_sequence_length`, `head_sink`, `position_ids`, softcap and `attention_bias`, and `S == 1` decode with the input present. `TestQuerySeqlensQkOutputCpu` checks `output_qk` for a padded chunk against the unpadded call, including the zeroed columns. #### Testing H200 host with a CUDA 13.0 Release build, so the gtest filter below also runs the CUDA tests from #33503. The Python classes below run on the CPU provider only. All results are from the head of this branch. - `onnxruntime_provider_test --gtest_filter='*GroupQueryAttention*:*GQA*'`: **149 passed, 18 skipped, 0 failed** (167 run). The skipped tests need an execution provider or hardware that this build does not have. - `test_gqa.py TestGQAWindowedKvCacheCpu`: **25 passed, 1 skipped** (the CUDA-only cuDNN fallback test). - `test_gqa.py TestQuerySeqlensQkOutputCpu` (`output_qk` of padded chunks): **8 passed**. - `test_gqa.py TestBorrowedKVCache` (cache inputs stay unchanged across the refactored copy paths): **46 passed**. - Mutation checks: four bugs were injected one at a time into a scratch copy of the kernel (non-windowed past length taken from `S` instead of `Q`, default rotary positions ignoring `Q`, windowed planning ignoring `Q`, and padded chunk handling disabled). Each is caught by both a C++ test and the Python suite. The first round showed that no C++ test covered the rotary one, which is why `QuerySeqlensRotaryRightPaddedChunkMatchesUnpadded_CPU` exists. Two more were checked for the later fixes: reverting the `output_qk` zero-fill produces 24 failures, and limiting the present-output requirement to `S > 1` fails exactly `QuerySeqlensSingleTokenWithoutPresentOutputsRejected_CPU`. - `lintrunner -m origin/main`: clean. - The CUDA results, including the two unrelated failures that also fail on a `main` build, are in #33503. Notes for reviewers: - `docs/ContribOperators.md` is refreshed the same way as in #33503, with a script that mirrors `gen_contrib_doc.py` because my local build is lite-protobuf. The full-protobuf CI docs check is authoritative. - The new WebGPU rejection test and the three comment edits in the MIGraphX, OpenVINO and JS providers are not compiled locally, because those providers are not part of my build. - Padded output rows are unspecified. A padded row that sees no key (a local window combined with heavy padding) may hold NaN. The contract allows that, but no test exercises it. - This PR gets only the lint workflow while it is stacked (see above). Once #33503 merges, I will retarget it to `main` and rebase. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Static-shape (compiled or graph-captured) prefill uses a fixed query chunk even when only a few rows are real tokens, for example the last chunk of a prompt padded to `S = 512`. Then `sequence_length` is the padded size while `seqlens_k[b] = P + Q - 1` reflects the real length. On the CPU provider this: - rejected a subsequent step with `seqlens_k + 1 < sequence_length`, - derived the past length and the number of rows to evict from the padded length, which dropped valid cache rows, - appended all `S` rows to the cache, which could overwrite valid entries or write past the capacity when `C > local_window_size`, and - anchored the default rotary positions and the causal mask at `T - S` instead of `P`. Inferring the valid length from `total_sequence_length` cannot work in general, because existing callers pass the cache capacity there. An explicit per-entry valid length removes the ambiguity, needs no heuristic, and gives the CPU and CUDA implementations the same contract. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>