[Chi Lo (chilo-ms)]

Description

Factors the sequence-length-independent parts of the CUDA GroupQueryAttention backend-selection logic out of GroupQueryAttention::ComputeInternal into a new header-only module, group_query_attention_eligibility.h. It provides four pure "seq-free" route-eligibility predicates:

Helper	Route
IsGQAXqaEligibleSeqFree<U>	XQA (Extreme Query Attention) decode kernel
IsGQACudnnSdpaCoreEligibleSeqFree<T, U>	cuDNN SDPA
IsGQAFlashEligibleSeqFree<T>	FlashAttention
IsGQAMemoryEfficientEligibleSeqFree<T>	CUTLASS memory-efficient attention

A seq-free predicate depends only on partition-time facts (geometry, device capability, quantization, feature flags) and not on the per-step sequence length or phase (prompt vs decode). The runtime keeps the sequence-length-dependent gates and dispatch priorities at the call sites:

• XQA — phase gates (is_first_prompt / sequence_length == 1 / kv_sequence_length) and the per-node shared-memory capability probe.
• cuDNN — cudnn_sdpa::is_stable() and the seq-dependent is_supported(...) probe.
• Flash / MEA — the !use_xqa / !use_cudnn_sdpa / !use_flash_attention dispatch priority.

No functional change

The eligibility booleans are computed from the same expressions in the same short-circuit-equivalent order. Every moved term is a side-effect-free read, so the selected backend is identical for every input. The GQA operator tests exercise all four routes and validate this (a wrong predicate would dispatch the wrong kernel and fail the tests).

Why header-only

All four helpers are templated on the Q/KV element types (T/U), so the module is header-only rather than the usual .h/.cc split.

Motivation

This is a small NFC step toward policy-based, capacity-aware CPU offloading (see the design doc in #32944). The shared seq-free core is what the Level-1 workspace estimator will later reuse, so runtime route selection and workspace estimation cannot drift.

Relationship to sibling PRs

• [docs] Policy-based, route-aware workspace estimation design #32944 — design doc (policy-based workspace estimation).
• [CUDA] Consolidate GQA scratch sizing (MEA, unfused, XQA, Flash) into shared workspace recipes #32945 — consolidates GQA scratch sizing into recipe helpers. This PR is the sibling that consolidates route eligibility. Both are independent NFC refactors; whichever merges second will need a trivial merge in the XQA/Flash region of group_query_attention.cc.

[Copilot]

Copilot review overview

🟡 Changes recommended

Runtime-only aliasing is included in a partition-time helper, and the new predicates lack direct regression tests.

Review effort: Balanced
Findings: 2

Open (2)

• Keep cache-alias checks at runtime, not in the seq-free policy · New
• Add focused tests for all eligibility predicates and rejection gates · New

What changed in this PR

Extracts CUDA GQA backend eligibility into reusable header-only predicates for future workspace estimation.

Changes:

• Adds predicates for XQA, cuDNN SDPA, FlashAttention, and MEA.
• Refactors runtime dispatch to use the predicates.

File	Description
group_query_attention.cc	Uses shared eligibility helpers.
group_query_attention_eligibility.h	Defines backend eligibility predicates.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Reviewed head 860b3906df9962c596c62770837a12f7ce5cd7f9 against base 967a8e064d08c98b910ed13e9f05604342327df1, in the context of #32944.

The runtime extraction preserves the phase gates, shared-memory fallback, cuDNN probes, and dispatch priority. I did not establish a current runtime route regression. Two actionable issues remain:

1. Minor — the shared helper's partition-time contract includes runtime aliasing

Location: onnxruntime/contrib_ops/cuda/bert/group_query_attention_eligibility.h:10-18,55,70-78.

The header defines "seq-free" as depending only on partition-time facts, but past_present_share_buffer is actual K/V pointer equality calculated during ComputeInternal(). L1 cannot generally know that fact.

This is not evidence of current L1 underestimation: the estimator does not call the helper in this PR. It is a reuse-contract problem. A future caller that leaves unknown aliasing at the default false could incorrectly exclude a reachable XQA route.

Please keep actual aliasing beside the runtime phase gates, or explicitly distinguish runtime eligibility from AOT alias possibility. Also document that these geometry arguments are exact, validated values. Exact eligibility at componentwise maximum geometry is not a proof of reachability over a bounded domain; smaller supported head sizes or groups may still be reachable.

2. Minor — add direct predicate regression tests

Location: group_query_attention_eligibility.h:70-188; this PR adds no direct predicate tests.

Output-only attention tests can pass through a different correct backend after an eligibility regression. Please add table-driven positive cases and individual rejection gates for the four predicates, including:

• Nonquantized group 5 versus quantized group 5.
• FP8 SM80 rejection versus SM89/SM90 support, under the appropriate build guards.
• Bias, quantization, QK normalization, and smooth-softmax restrictions.
• cuDNN type/window restrictions and Flash/MEA feature gates.

Scope and recommendation

This is runtime-only extraction, not yet a shared bounded L1 reachability implementation. Please address the contract wording/boundary and focused test coverage before using it as that implementation. No build, GPU execution, numerical, or plugin/minimal-build validation was performed for this PR.

[Chi Lo (chilo-ms)]

Thanks Ti-Tai Wang (@titaiwangms) — addressed both points in 7b19676.

1. Actual past/present aliasing now stays at the runtime call site alongside the phase gates, outside the seq-free inputs. The header documents exact, validated positive geometry fitting the runtime int32 parameters, with a valid head ratio. It explicitly states that this is not bounded-domain L1 reachability: rejecting componentwise maximum geometry cannot exclude smaller supported shapes. A regression case demonstrates that distinction for XQA head size and group size.
2. Added direct table-driven/typed tests for all four predicates, covering the requested group-5 distinction, FP8 SM80 rejection versus SM89/SM90 support under build guards, individual bias/quantization/QK-Norm/smooth-softmax gates, cuDNN type/window restrictions, and Flash/MEA feature gates.

Validation: the updated GQA dispatch and all predicate tests compiled with Flash/MEA/FP8 enabled. The host-only predicate suites passed 15 tests with Flash disabled, 13 with Flash and FP8 disabled, and 11 with Flash/MEA/FP8 disabled. Full Flash test execution remains unverified because the standalone Windows harness failed before Google Test startup; no GPU numerical or plugin/minimal-build validation is claimed.

[Ti-Tai Wang (titaiwangms)]

Follow-up review at 7b19676: both issues from my previous review are addressed.

Actual past/present aliasing is now checked at the runtime dispatch site rather than included in the seq-free predicate inputs. The new direct eligibility tests cover independent feature gates, storage types, architecture differences, and supported/unsupported geometry.

The header now explicitly states that these predicates take exact, validated geometry and are not a bounded-domain Level-1 reachability implementation. That distinction is important: rejecting a componentwise maximum does not exclude a route reachable at smaller supported geometry, and the added test documents this nonmonotonicity.

No additional actionable findings from this updated source review. This slice looks ready from the code-review perspective; the checks currently have no failed or pending entries. This COMMENT is not a formal approval.

Validation scope: static review of the updated implementation and tests, plus inspection of CI status; no new local build or runtime execution at this head.

[Ti-Tai Wang (titaiwangms)]

Reviewed 7b19676. The previous aliasing and direct eligibility-test concerns are addressed; no remaining actionable findings from this updated source review. Approval is limited to this exact-geometry refactor, not bounded-domain estimator reachability. No new local build/runtime execution was performed.