Support bidirectional GroupQueryAttention on CPU and CUDA

[Copilot]

Description

GroupQueryAttention previously always applied a causal mask. This adds a causal attribute, defaulting to 1 for backward compatibility.

• CPU
  • Uses bidirectional masking when causal=0.
  • Routes bidirectional execution through the compatible unfused path.
  • Rejects local_window_size != -1 with bidirectional attention because local-window alignment is defined only for causal attention.
• CUDA
  • Propagates the attribute across Flash Attention, memory-efficient attention, cuDNN SDPA, and unfused paths.
  • Excludes causal-only XQA for bidirectional attention.
  • Rejects local_window_size != -1 with bidirectional attention.
  • Quantized bidirectional KV-cache execution requires Flash Attention. MEA and unfused fallbacks do not consume quantized KV caches and return NOT_IMPLEMENTED instead of reading them incorrectly.
• Other EPs
  • WebGPU and JS report NOT_IMPLEMENTED for causal=0.
  • DML rejects causal=0 during kernel creation, and WebNN declines the node during capability checks, avoiding silent causal output.
• Coverage
  • Adds default-causal and bidirectional CPU/CUDA mask tests with identity-sensitive Q/K logits.
  • Adds non-quantized and quantized bidirectional past-KV parity coverage.
  • Adds local-window rejection and WebGPU rejection tests.

tester.AddAttribute<int64_t>("causal", 0);

Motivation and Context

Bidirectional models require each query token to attend to the full valid key sequence. The new attribute enables this on CPU and CUDA while preserving existing causal behavior by default. Generation conversion stamps causal=1 explicitly because its decoder attention is unidirectional by definition.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Adds a new causal attribute to the com.microsoft::GroupQueryAttention contrib op to support bidirectional (non-causal) attention on CPU and CUDA while preserving the existing causal-by-default behavior for backward compatibility. The change propagates the attribute through backend selection and masking logic, and updates documentation plus test coverage across EPs.

Changes:

• Introduce causal attribute (default 1) in the operator schema and documentation; treat 0 as bidirectional attention.
• CPU/CUDA: wire causal into masking and backend eligibility (e.g., disable causal-only kernels like XQA when causal=0).
• WebGPU/JS: explicitly reject causal=0 with NOT_IMPLEMENTED, with corresponding tests.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_gqa.py	Adds causal to test config/node attributes and expands CUDA parity + rejection coverage (including quantized bidirectional cases).
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc	Adds focused unit tests for default-causal behavior, bidirectional masking, invalid values, and WebGPU NOT_IMPLEMENTED.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Extends the GroupQueryAttention schema with the causal attribute (default 1) and updates operator doc text.
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h	Enforces causal ∈ {0,1} and rejects causal=0 as not implemented for WebGPU.
onnxruntime/contrib_ops/js/bert/group_query_attention.h	Rejects causal=0 as not implemented for the JS implementation.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc	Parses/validates causal, propagates it to parameters, blocks causal-only XQA when bidirectional, and threads it into cuDNN SDPA.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu	Propagates causal into MEA/cuDNN calls; improves the unfused-path error for quantized KV-cache unsupported cases.
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc	Propagates causal into runtime parameters and prevents using the CPU flash path for bidirectional attention.
onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h	Adds causal_ parsing/validation and adjusts masking bounds for bidirectional vs causal behavior.
docs/ContribOperators.md	Documents the new causal attribute for GroupQueryAttention.
docs/contrib_ops/cuda/gqa.md	Updates CUDA GQA documentation to describe causal and backend support/eligibility for bidirectional attention.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Automated review synthesis

Verdict: request changes. Core CPU/CUDA masking math for causal=0 is correct, but there are a few real defects.

🔴 Critical

1. Bidirectional prompt test validates against a causal reference — onnxruntime/test/python/transformers/test_gqa.py:1202
parity_check_gqa_prompt now stamps causal=0 onto the ORT node (via config.causal), but its attention_ref(...) call still hardcodes causal=True — only the _past variant (line 1539) was updated. attention_ref overrides window_size to (left, 0) when causal=True (test_gqa.py:950), so test_gqa_prompt_bidirectional_attention_bias_broadcast compares a bidirectional ORT output against a causal PyTorch reference. The test either fails, or passes only because the ORT side silently ignored causal=0 — both are bugs. Fix: causal=causal at line 1202.

🟠 Major

2. local_window_size >= 0 combined with causal=0 gives three different, unvalidated answers across backends

• CPU (gqa_attention_base.h:461,1485): start_off = total_seqlen - local_window_size is the same for every query row — this is not a sliding window at all when bidirectional.
• CUDA Flash: bottom-right-aligned left window.
• CUDA MEA: top-left-aligned (causal_diagonal_offset stays 0 when causal=false).

Nothing rejects this attribute combination. Recommend either rejecting local_window_size >= 0 && causal == 0 (matching the ORT_NOT_IMPLEMENTED precedent already set for JS/WebGPU) or specifying and implementing one consistent alignment.

3. DML and WebNN silently ignore the new causal attribute → silent wrong output for bidirectional models

• onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorGroupQueryAttention.cpp never reads causal.
• onnxruntime/core/providers/webnn/builders/impl/gqa_op_builder.cc reads local_window_size etc. but never causal, and IsOpSupportedImpl doesn't reject causal=0.
• MIGraphX/OpenVINO also list GQA support without attribute inspection (lower priority, delegated frontends).

These EPs will now compute causal output for a model author who explicitly requested bidirectional attention, with no error. Recommend adding a rejection for causal == 0 in DML kernel creation and WebNN IsOpSupportedImpl, mirroring the JS/WebGPU pattern this PR already introduces.

4. Test coverage gaps

• No non-quantized bidirectional decode/past-KV parity test (only prompt + quantized-past are covered).
• The new C++ unit tests (RunGQACausalMaskTest) use all-zero query and key, so every logit is identical; they can only detect how many keys were averaged, not which — an off-by-one or wrong-alignment bug would still pass. Suggest adding a case with distinct Q/K values and a ragged batch (seqlens_k = {3, 1}) to catch exactly the boundary issues in finding Remove vsts test runner in cmake file #2.

🟡 Minor

• Naming inconsistency: the new CPU-side member is causal_, while CUDA already uses (and this PR reuses) the cross-EP convention is_unidirectional_/parameters.is_unidirectional. The new schema attribute is also named causal while the existing MultiHeadAttention schema uses unidirectional (defaulting to 0, the opposite default). Contrib-op attribute names are effectively permanent — worth reconciling now rather than after release.
• effective_causal_length (gqa_attention_base.h) is a misleading name once causal_=false — it's just "effective visible length" in that branch, and the surrounding comment still explains only the causal case. Consider clarifying why the bidirectional branch is safe (it equals the QK GEMM's N, so it can't read past the filled region).
• onnxruntime/python/tools/transformers/convert_generation.py replaces MultiHeadAttention nodes with GroupQueryAttention but doesn't map unidirectional → causal; since the defaults are opposite (0 vs 1), this silently forces causal behavior during that conversion. Likely fine in practice (text-gen models are causal) but breaks the attribute contract silently.
• Scope creep, undocumented: the new !is_inputs_quantized guard on MEA eligibility plus the new quantized-KV NOT_IMPLEMENTED in QkvToContext aren't mentioned in the PR description. They look like a necessary/correct fix (disabling XQA for causal=0 routes quantized decode into MEA, which can't read a quantized cache), but should be called out explicitly or split into a separate PR/commit.
• The cuDNN SDPA eligibility comment lost the "bottom-right causal" detail explaining why an attention bias can't compose with cuDNN's fused mask.

✅ Cleared (hypotheses checked and refuted)

• Flash Attention was not hardcoded causal — group_query_attention_impl.cu:1156/1221/1334/1409 (and unfused :1615) already read parameters.is_unidirectional on main; they were simply dormant because is_unidirectional_ was pinned to true. This PR correctly activates them.
• No uninitialized/padding-memory read for causal=0: total_seqlen is the per-batch-item valid KV length and equals the QK GEMM's N, so effective_causal_length ≤ total_seqlen still holds and the zero-fill loop is empty.
• Schema addition is backward-compatible (contrib ops are pinned at v1; default 1 reproduces prior hardcoded behavior).
• smooth_softmax / head_sink / softcap carry no causal assumption — math holds for causal=0.

Open questions (need execution to settle; not run as part of this review)

• CPU-vs-CUDA output divergence for causal=0 && local_window_size>0 (needs a CUDA build).
• Non-quantized bidirectional decode parity on CUDA.

[Tianlei Wu (tianleiwu)]

Addressed the actionable feedback in commit 9eab39cdce:

• Fixed parity_check_gqa_prompt to pass the requested causal value to the reference.
• Defined local windows as causal-only and reject causal=0 with local_window_size != -1 in both CPU and CUDA constructors; updated schema/CUDA docs and added rejection tests.
• Added explicit DML kernel-creation rejection and WebNN capability rejection for causal=0, so neither silently produces causal output.
• Strengthened the C++ mask tests with distinct Q/K logits and added non-quantized bidirectional past-KV parity coverage. Since non-causal local windows are now rejected, the backend-alignment ambiguity no longer reaches execution.
• Renamed the CPU state to is_unidirectional_, clarified the visible-length logic, and restored the cuDNN bottom-right causal-mask rationale.
• Made generation conversion stamp causal=1 explicitly. That converter is intentionally decoder-specific, so propagating MHA's generic unidirectional=0 default would change generation semantics.
• Updated the PR description to call out the quantized bidirectional fallback guard and NOT_IMPLEMENTED behavior.

I left the OpenVINO/MIGraphX allowlists unchanged: those EPs delegate GQA to external backend frontends and do not have an in-repo GQA lowering where this attribute can be gated. Their bidirectional support needs backend-specific capability confirmation separately.

[Ti-Tai Wang (titaiwangms)]

Re-review after update

Re-checked the latest push (commits 44c5dd7…aa07930) against my prior review. All previously flagged Critical/Major issues are resolved:

• ✅ Critical (test validates bidirectional output against causal reference) — fixed. parity_check_gqa_prompt's attention_ref call now uses causal=causal (was hardcoded True) at what's now test_gqa.py:1201, matching the earlier fix already present in parity_check_gqa_past.
• ✅ Major (local_window_size >= 0 + causal=0 cross-backend divergence) — fixed. Both CPU (gqa_attention_base.h) and CUDA (group_query_attention.cc) constructors now ORT_ENFORCE(is_unidirectional_ || local_window_size_ == -1, ...), rejecting the combination outright rather than silently producing backend-dependent results. New tests BidirectionalLocalWindowRejected_CPU/_CUDA cover this.
• ✅ Major (DML/WebNN silently ignore causal) — fixed. DML now reads the new AttrName::Causal attribute and rejects causal=0 via ML_CHECK_VALID_ARGUMENT; WebNN's IsOpSupportedImpl now rejects nodes with causal != 1. convert_generation.py's replace_mha_with_gqa now explicitly passes causal=1, removing the silent-default-mismatch concern versus MultiHeadAttention's unidirectional default of 0.
• ✅ Major (test coverage gaps) — improved. New test_gqa_decode_bidirectional adds non-quantized bidirectional decode/past-KV parity coverage on CUDA, and a new TestQuantizedBidirectionalGQA class covers the quantized-KV-cache rejection path. The C++ RunGQACausalMaskTest helper now uses distinct, non-uniform Q/K values instead of all-zeros, so it can actually detect masking/alignment bugs rather than just counting attended keys.
• ✅ Minor (naming split causal_ vs is_unidirectional_) — resolved; the CPU-side member was renamed to is_unidirectional_, consistent with the CUDA kernel and the shared AttentionParameters::is_unidirectional field.
• ✅ Minor (misleading effective_causal_length name/comment) — resolved; renamed to visible_length with an updated comment explaining the bidirectional case explicitly.
• ✅ Minor (cuDNN comment lost rationale) — resolved; the comment now explicitly states cuDNN's bottom-right causal mask can't compose with an attention bias, regardless of causal/bidirectional.

Remaining, non-blocking:

• MIGraphX/OpenVINO still list GroupQueryAttention support without inspecting causal (pre-existing, lower-priority EPs not covered in this pass — worth a follow-up issue if those EPs are actively used with this op).
• CPU-vs-CUDA output parity for causal=0 combined with other exotic feature combinations beyond local_window_size hasn't been exhaustively cross-validated, but the rejected combination was the one with divergent semantics, so this is low risk.

No remaining blocking issues from this round. LGTM pending CI.