[Tianlei Wu (tianleiwu)]

Description

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
  [CUDA] Support right-padded prefill chunks in GroupQueryAttention #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 [CUDA] Support right-padded prefill chunks in GroupQueryAttention #33503.

Notes for reviewers:

• docs/ContribOperators.md is refreshed the same way as in [CUDA] Support right-padded prefill chunks in GroupQueryAttention #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 [CUDA] Support right-padded prefill chunks in GroupQueryAttention #33503 merges, I will retarget it to
  main and rebase.

Motivation and Context

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.

[Copilot]

🟡 Changes recommended

Regression coverage misses the in-place cache-overflow boundary, and the public padded-length contract needs clarification.

2 open issues

• Add tests covering in-place cache overflow boundaries · New
• Document padded versus valid sequence length conventions · New

What changed in this PR

Adds right-padded prefill support to CPU GroupQueryAttention for full-length and sliding-window KV caches.

Changes:

• Resolve valid token counts and past lengths for cache planning and rotary positions.
• Limit cache writes to valid rows across float and quantized attention paths.
• Add padded-versus-unpadded parity tests.

File	Description
onnxruntime/​test/​python/​transformers/​test_gqa.py	Adds padded-step helpers and CPU parity tests.
onnxruntime/​test/​contrib_ops/​group_query_attention_op_test.cc	Adds a padded windowed-cache regression test.
onnxruntime/​contrib_ops/​cpu/​bert/​group_query_attention.cc	Updates validation, cache planning, and rotary positions.
onnxruntime/​contrib_ops/​cpu/​bert/​group_query_attention_helper.h	Adds shared step-length resolution.
onnxruntime/​contrib_ops/​cpu/​bert/​gqa_attention_base.h	Uses valid row counts and resolved causal offsets.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The static-cache changes regress previously accepted unpadded calls by dropping real tokens, shifting causal/RoPE offsets, and rejecting some capacity-valued inputs.

1 open issue

• Preserve default interpretation of unpadded capacity-valued calls · New

2 previously missed issues

In code that hasn't changed since last review

Test ragged batched windowed prompts with flash on and off

onnxruntime/​contrib_ops/​cpu/​bert/​group_query_attention.cc:342

Removing the windowed condition also enables first windowed prompts with batch_size > 1 and ragged valid lengths, but all new padded tests use a single batch. Add a ragged-batch test comparing each entry's valid outputs and resident cache with an unpadded single-entry run, with flash enabled and disabled. Include unequal lengths on opposite sides of capacity to exercise the shared staging decision and per-entry lengths.

Pad attention bias to match padded dimensions and test position IDs

onnxruntime/​test/​python/​transformers/​test_gqa.py:4477

_check_parity allocates the bias for the valid total, so this slice truncates padded dimensions: the third step of the new sequence produces (10, 43) instead of (32, 65) and fails input validation. Slice the valid bias region and pad both axes to the required sizes. Add a padded case with attention bias and explicit position IDs, with flash enabled and disabled; none of the new cases covers these inputs.

2 resolved since last review

• Add tests covering in-place cache overflow boundaries
• Document padded versus valid sequence length conventions

_{Review effort: Max}

[Copilot]

🟡 Changes recommended

Critical correctness and memory-safety defects remain, including out-of-bounds K/V reads, uninitialized cache and QK data, and fusion that drops query_seqlens.

10 open issues

• Initialize unwritten QK output columns in both paths · New
• Reject missing outputs whenever query_seqlens is supplied · New
• Copy past cache rows into distinct present buffers · New
• Clamp shared-KV key length to the cache size · New
• Preserve query_seqlens when fusing CUDA attention nodes · New
• Regenerate OperatorKernels.md with query_seqlens support · New
• Preserve default interpretation of unpadded capacity-valued calls
• Skip the legacy cache-length equality for query_seqlens · New
• Add focused JSEP and WebGPU query_seqlens rejection tests · New
• Test packed-QKV CPU query_seqlens with rotary and both attention paths · New

_{Review effort: Max}

[Ti-Tai Wang (titaiwangms)]

Approved the incremental CPU change, reviewed as 823a83b -> 5a499a8.

No new Major/Critical correctness or integration issue was established in this PR. The per-batch step-length resolution, Q-only cache append, window planning, absolute RoPE positions versus cache-relative masks, flash/unfused consumers, and initialized output-QK tail are consistent with the shared contract. Explicit padded-row CPU position_ids are range-checked by RunRotaryEmbedding. The native WebGPU rejection regression is also present.

One nonblocking comment clarification: at onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h:136-141, describe causal_past_seqlen as the attention-buffer/cache-relative index of the first query row, not an unqualified absolute position. For windowed calls it is the retained resident-past count; absolute P remains the RoPE coordinate. The implementation uses these coordinates correctly.

This approval is for the incremental CPU implementation. It does not resolve or waive the focused JSEP regression-coverage requirement reported on the base PR #33503, or its nonblocking test/documentation follow-ups.

Validation scope: five-role static review only; no local build/tests, CUDA execution, memcheck, performance/concurrency qualification, or exhaustive Flash/MEA-internal proof. The observed CI snapshot was 99 successful and 2 neutral checks, not independent local reproduction. No Major/Critical finding was overruled.