onnxruntime
125ea215 - Support bidirectional GroupQueryAttention on CPU and CUDA (#31704)

Commit
60 days ago
Support bidirectional GroupQueryAttention on CPU and CUDA (#31704) ### 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. ```cpp 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. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>
Author
Parents
Loading