[Tianlei Wu (tianleiwu)]

Description

Add an optional query_seqlens input (index 16) to com.microsoft::GroupQueryAttention and implement it
in the CUDA execution provider. It carries the number of valid, non-padded tokens of a right-padded
query/key/value chunk per batch entry, so fixed-shape chunked prefill works with FlashAttention,
MemoryEfficientAttention and unfused attention, with a full-length or a windowed KV cache.

This replaces the right_padded_query attribute that earlier revisions of this PR introduced. The attribute
and the CPU heuristic in #33505 are superseded by one input that is part of the operator contract, so every
execution provider can implement identical semantics. The attribute was never in a release, so no released
schema changes.

Stack: this PR is the base of a two-PR stack. #33505 (draft) adds the CPU implementation. Its base is the branch
of this PR, so it shows only the CPU changes, and it is retargeted to main once this PR merges. 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 #33505 gets them after that retarget. This branch is rebased onto the current main.

Contract

For batch entry b: S is the padded sequence_length, T = seqlens_k[b] + 1 is the number of tokens in
the cache after the step, Q = query_seqlens[b] is the number of valid rows, and P = T - Q is the number
of tokens cached before the chunk.

• Rows [0, Q) of query/key/value are valid and rows [Q, S) are right padding.
• Only the valid rows are appended to the cache, at [P, T). Query row s sits at absolute position P + s
  for causal masking, local_window_size and implicit rotary positions. Output rows [Q, S) and cache rows at
  or beyond T are unspecified.
• total_sequence_length is only an upper bound of T (for example the cache capacity). It no longer has to be
  P + S and is not used to tell a first prompt from a later chunk.
• Requirements: causal=1, kv_sequence_length == S, past present, 1 <= Q <= min(S, T), T <= total_sequence_length, S > 1 requires batch_size == 1, and Q == 1 when S == 1. A cache that is not
  windowed also needs P + S <= C (C is the sequence dimension of present_key) and its P cached tokens in
  past_key/past_value. Those buffers may be shorter than present_key, so a cache can grow between steps.
• Omitting the input keeps the existing behavior, so callers that pass the cache capacity as
  total_sequence_length are unaffected.

Key changes

• Schema/docs: new input 16, the right_padded_query attribute is removed, and a "Right-padded query chunks"
  section is added to the operator description. docs/ContribOperators.md and docs/contrib_ops/cuda/gqa.md
  are updated.
• Shared validation (group_query_attention_helper.h): CheckQuerySeqlens validates shape, type and, when
  the tensors are host resident, the per-entry values. With the input present, CheckInputs no longer classifies
  the step by comparing S with total_sequence_length, and it skips the legacy check that ties a growing cache to
  past + S == total_sequence_length, which rejected valid calls such as past rows 10, S = 8, Q = 3,
  total_sequence_length = 32.
• CUDA: GetSequenceLengths takes Q from the input, clamped to [1, min(S, T)], and derives P = T - Q
  from seqlens_k instead of total_sequence_length - S. Padded rows never enter the cache and the unfused
  causal mask is anchored at P. cuDNN SDPA is skipped for padded chunks because it assumes unpadded subsequent
  queries. S == 1 ignores the input so XQA and flash decode stay enabled. query_seqlens with causal=0 is
  rejected. Device-resident values cannot be checked on the host, so out-of-contract values give unspecified but
  memory-safe results.
• CUDA quantized KV cache: an ordinary first prompt attends to the float K/V of the chunk, while later chunks
  read the quantized cache. With the input the host cannot tell them apart, so a chunk is handled as a later one,
  and an entry with no past and more than one valid row is then re-attended to the float K/V to match the unpadded
  first prompt. An entry with one valid row behaves like a single-token step and keeps reading the quantized cache,
  and a call that omits the input keeps the previous classification.
• CPU: temporarily rejects the input with INVALID_ARGUMENT so it is never silently misinterpreted.
  [CPU] Support right-padded prefill chunks in GroupQueryAttention #33505 replaces the rejection with the implementation.
• Other EPs (WebGPU, JS, DML, WebNN, OpenVINO, MIGraphX): reject or decline nodes that set input 16. These
  edits are small, but they were not compiled locally because those EPs are not part of this build, so CI is
  the first compile of them. The OpenVINO check also applies to the com.microsoft domain.
• Latent bug fix (group_query_attention_pre_norm_fusion.cc, group_query_attention_fusion.cc): both fusions
  rebuilt the node input list with a fixed number of slots and dropped any later input, which would have silently
  removed query_seqlens. They now keep all inputs, with unit tests.

Testing

CUDA 13.0 and cuDNN 9, H200 (sm_90), Release build, FlashAttention and MemoryEfficientAttention enabled, on the
head of this branch.

• onnxruntime_provider_test --gtest_filter='*GroupQueryAttention*:*GQA*': 138 passed, 18 skipped, 0 failed.
  The skipped tests need WebGPU. New goldens compare padded chunks with the same call on an unpadded chunk for
  Flash, MEA and unfused in fp16/bf16, with and without packed QKV, for full-length and windowed caches, and for a
  cache that grows between steps. New negative tests cover a bidirectional node, a missing past, a batched chunk
  and a wrong shape.
• onnxruntime_test_all --gtest_filter='*GroupQueryAttentionPreNormFusion*:*GroupQueryAttentionFusion*':
  33 passed, 4 skipped (WebGPU). Reverting the input-count fix makes the new test fail.
• test_gqa.py, the whole file on this build: 1003 tests, 8 skipped, 2 failures. The failures,
  TestFlashGQA.test_gqa_decode_bidirectional and
  TestGQAAttentionBias.test_windowed_cache_attention_bias_is_rejected, also occur on an independent build of
  main and are not related to this change. The right-padded tests cover fp16 and int8-quantized KV (including a
  chunk with one valid row and a call that omits the input), the cuDNN fallback, capacity-bound caches,
  rotary/softcap/head_sink/attention_bias/packed QKV, and S == 1 decode.
• compute-sanitizer --tool memcheck: 0 errors on the head of this branch for the quantized-KV tests
  (TestQuerySeqlensQuantizedCache, 16 tests), the padded windowed-cache tests (7) and the right-padded CUDA
  gtests. An earlier revision also ran on the Flash, MEA and unfused paths for in-contract and out-of-contract
  inputs (Q of 0, negative, above S or T; T above total_sequence_length; negative seqlens_k; P + S
  above the cache) with exact-size buffers separated by spacers, and a deliberately shortened cache is flagged by the
  same harness. That harness was not repeated after the later review fixes.
• lintrunner -m origin/main: clean.

Notes for reviewers:

• docs/ContribOperators.md was refreshed from the schema source with a script that mirrors
  gen_contrib_doc.py for the GQA entry, because my local build is lite-protobuf and strips operator
  descriptions. It matches the generator's output format; the full-protobuf CI docs check is authoritative. The
  DML row of docs/OperatorKernels.md follows the generator's format by hand, because DML cannot be built on Linux.

Motivation and Context

Static-shape (for example compiled or graph-captured) prefill uses a fixed query chunk even when only a few rows
are real tokens. Previously, subsequent chunks derived the past offset and the eviction count from the padded
length, which corrupted the KV cache and shifted causal and RoPE alignment. Inferring the padding 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

The padding heuristic misclassifies existing fixed-bound callers, shifting cache offsets and discarding valid tokens.

1 open issue

• Unpadded chunks misclassified as padded with capacity-valued lengths · New

What changed in this PR

Adds right-padded chunked prefill support to CUDA GroupQueryAttention across FlashAttention, memory-efficient, and unfused backends.

Changes:

• Adjusts sequence lengths, cache eviction, and causal alignment for padding.
• Zeros padded K/V entries.
• Adds padded-prefill coverage, including quantized caches.

File	Description
onnxruntime/​test/​python/​transformers/​test_gqa.py	Adds padded-step parity tests across backends.
onnxruntime/​test/​contrib_ops/​group_query_attention_op_test.cc	Corrects test lengths and covers padded windowed chunks.
onnxruntime/​contrib_ops/​cuda/​bert/​unfused_attention.h	Adds optional per-batch past lengths.
onnxruntime/​contrib_ops/​cuda/​bert/​unfused_attention.cu	Uses explicit past offsets for masking.
onnxruntime/​contrib_ops/​cuda/​bert/​group_query_attention_qkv.cuh	Zeros padded cache entries and adjusts RoPE assertions.
onnxruntime/​contrib_ops/​cuda/​bert/​group_query_attention_impl.cu	Updates padding-aware lengths and backend wiring.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Shared-KV calls retain memory-safety and causal-alignment regressions, and the forced-Flash test fails in quick builds.

3 open issues

• Clamp shared-KV padding length to cache capacity · New
• Preserve signed past offset for shared-KV alignment · New
• Use supported head size in quick-build Flash regression · New

1 resolved since last review

• Unpadded chunks misclassified as padded with capacity-valued lengths

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Unresolved critical findings include incorrect quantized attention and cache handling, possible out-of-bounds reads, and graph-resolution failures.

10 open issues

• Preserve floating-point KV for valid quantized first-chunk rows · New
• Preserve borrowed-KV length and causal alignment without query_seqlens · New
• Copy past KV into non-aliased quantized present buffers · New
• Preserve query_seqlens in GroupQueryAttentionFusion · New
• Reject query_seqlens GQA nodes during OpenVINO partitioning · New
• Guard BF16 golden test on Ampere-or-newer CUDA · New
• Clamp shared-KV padding length to cache capacity
• Use supported head size in quick-build Flash regression
• Preserve signed past offset for shared-KV alignment
• Regenerate OperatorKernels.md for the new GQA input · New

_{Review effort: Max}

[Copilot]

🟡 Changes recommended

Unresolved critical issues break INT8 padded/unpadded parity, change behavior when query_seqlens is omitted, reject valid growing-cache inputs, and cause backend tests to fail on valid builds lacking optional features.

4 open issues

• Allow query_seqlens to bypass the legacy dynamic-cache check · New
• Padded Q=1 calls break zero-past quantized parity · New
• Gate float-KV replay on query_seqlens presence · New
• Gate forced-backend tests on compiled feature availability · New

10 resolved since last review

• Guard BF16 golden test on Ampere-or-newer CUDA
• Reject query_seqlens GQA nodes during OpenVINO partitioning
• Preserve query_seqlens in GroupQueryAttentionFusion
• Copy past KV into non-aliased quantized present buffers
• Preserve borrowed-KV length and causal alignment without query_seqlens
• Preserve floating-point KV for valid quantized first-chunk rows
• Clamp shared-KV padding length to cache capacity
• Use supported head size in quick-build Flash regression
• Preserve signed past offset for shared-KV alignment
• Regenerate OperatorKernels.md for the new GQA input

_{Review effort: Max}

[Ti-Tai Wang (titaiwangms)]

Full static review of 657cc52 -> 823a83b, including the stacked CPU implementation in #33505.

No concrete core correctness, memory-safety, or integration defect was established. I am withholding approval for the focused JSEP regression-coverage requirement below.

Required before approval

Missing focused JSEP rejection coverage -- onnxruntime/contrib_ops/js/bert/group_query_attention.h:28.

The constructor now rejects an existing query_seqlens input at slot 16. Please add focused JSEP tests showing that supplying the input fails with the intended error, while omitting it or supplying an empty optional slot preserves existing admission. The matching .github/instructions/jsep-deprecated.instructions.md explicitly requires focused regression coverage for accepted correctness fixes.

The native WebGPU rejection test added by #33505 is useful, but it does not exercise JSEP session construction. This is a coverage/policy gap, not an established runtime failure.

Nonblocking follow-ups

• Avoid probing availability with the regression case itself -- onnxruntime/test/python/transformers/test_gqa.py:575. has_sdpa_kernel() probes S=8, T=3, C=9, which is also the negative-diagonal boundary tested by test_query_longer_than_borrowed_kv(). A regression that makes Flash/MEA ineligible specifically for T < S can therefore turn the test into a skip. Probe availability with an ordinary supported control, such as S=T=8, C=9, and assert boundary-case dispatch separately.
• Clarify the quantized first-chunk reference -- onnxruntime/core/graph/contrib_ops/bert_defs.cc:1510 and onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu:1333-1337. For P=0, Q>1, the new-input path attends to float new K/V. Omitting input 16 while retaining a capacity-valued total_sequence_length preserves the legacy subsequent-chunk classification and quantized reads. Please specify whether the unpadded reference retains input 16 or uses the true total length. The added tests deliberately distinguish these paths; this is a documentation ambiguity, not a proven numerical bug.
• Update the padded-length comment -- onnxruntime/contrib_ops/cuda/bert/attention_data.h:144-145. padded_seq_lens is no longer first-prompt-only: it also carries P+S for padded continuations and cache-relative bounded lengths for windowed calls.

One nonblocking hardening question remains at group_query_attention_impl.cu:698: the initial signed-int P+S expression can overflow for an extreme windowed absolute length (T=INT_MAX, Q=1, S=8), before that value is overwritten by cache-relative geometry. No observable failure was established. Consider calculating it only where consumed, with widened arithmetic and clamping before narrowing.

Integration and validation scope

The review found agreement on input 16's type/placement, unsupported-provider rejection, trailing-input preservation in both fusions, and windowed-cache alias/staging admission. The absolute RoPE versus cache-relative causal-mask coordinates are consistent in the reviewed stack.

This was a five-role static review. No local build, CUDA execution, memcheck, performance/concurrency qualification, or exhaustive Flash/MEA-internal proof was performed. The observed CI snapshot was 99 successful and 1 neutral check; CI and author-reported tests are not independent local reproduction. No Major/Critical finding was overruled.

[Tianlei Wu (tianleiwu)]

Ti-Tai Wang (@titaiwangms) Addressed the feedback from your review summary in 21ff2d096a:

• Focused JSEP coverage: Added and wired three browser unit tests exercising actual JSEP session construction with CPU fallback disabled. A present input 16 fails with the intended JS-specific error; an omitted input and an explicit empty optional slot are accepted. The suite is loaded only for JSEP builds. All three tests passed on software WebGPU using this PR head’s CI-built JSEP WASM from Web CI run 41390541574; the production JSEP guard is unchanged.
• Availability probe: has_sdpa_kernel now probes the ordinary S=T=8, C=9 control. The T<S regression continues to assert its own dispatch, so it cannot hide a boundary regression by skipping on that case.
• Quantized reference: The schema, generated ContribOperators documentation, and CUDA GQA notes now specify that the unpadded reference retains query_seqlens and the same scalar bound. They also document float versus quantized first-chunk reads and the legacy tight-bound singleton precision rule.
• Length comment and hardening: Updated padded_seq_lens and past_seq_lens comments. Absolute P+S is now computed only for non-windowed calls, using widened arithmetic and clamping before narrowing. Added a windowed INT_MAX absolute-length regression checking both valid-row attention and the resident cache layout.

Local validation: CUDA 13.0 build with FlashAttention enabled; 43 focused Python tests passed, 6 non-CUDA variants skipped; 6 existing query_seqlens C++ tests passed; 3 actual JSEP browser tests passed; TypeScript, lint, and whitespace checks passed.

[Ti-Tai Wang (titaiwangms)]

Follow-up review of 823a83b -> 21ff2d0, building on the prior five-role full static review.

The blocking JSEP coverage gap is resolved: the new browser suite constructs actual sessions with CPU fallback disabled, tests present/omitted/empty input 16, asserts the JS-specific rejection error, and is wired only into JSEP builds. The production rejection guard is unchanged.

The remaining feedback is also addressed:

• The SDPA availability probe now uses S=T=8 rather than the T<S regression case.
• The schema/generated documentation and CUDA notes explicitly retain query_seqlens and the same scalar bound in the unpadded reference.
• The singleton replay threshold correctly distinguishes the legacy tight-bound first prompt from capacity-bound quantized decode; the new tests cover packed/unpacked and windowed cases.
• The length comments are updated. Absolute P+S is calculated only on the nonwindowed path, in int64_t, with capacity clamping before narrowing for explicit-length calls. The new INT_MAX windowed regression checks valid-row attention and resident cache layout.

No remaining code-review blocker was established in this follow-up. Approving the reviewed head.

Scope and validation limitations: this follow-up was direct static inspection of the complete revision delta and controlling code, not another five-role review or a local build/test run. The author's reported CUDA, C++, JSEP browser and lint results were not independently reproduced.

CI on this new head is not yet an all-green validation snapshot: the current PR check listing contains license/cla SUCCESS and Linux Android Emulator QNN CI Pipeline ACTION_REQUIRED. This approval does not waive required CI or authorize merging before required checks are satisfied.