onnxruntime
2cd57057 - [CUDA] Support attention_bias in GroupQueryAttention via the unfused path (#29525)

Commit
4 days ago
[CUDA] Support attention_bias in GroupQueryAttention via the unfused path (#29525) ### Description `com.microsoft.GroupQueryAttention`'s optional `attention_bias` input (input #10) is implemented by the CPU EP (#23944) and the WebGPU EP (#25285, #26769), but the CUDA EP rejects it at runtime. This PR wires it through. No new kernel is needed: the unfused GQA fallback added for #28195 calls `LaunchUnfusedAttention`, whose kernel already implements an additive bias with dim-0/dim-1 broadcast, per-batch `seqlens_k`, causal/sliding-window masking and softcap — the op just passed `attn_bias=nullptr`. The change is dispatch plumbing: - **`group_query_attention.cc`** — remove the blanket rejection; validate the bias element type; set `broadcast_attn_bias_dim_0/1` from the bias shape (the fields already exist on `AttentionParameters`); add `!has_attention_bias` to the XQA / cuDNN SDPA / flash / flash-fast-decode / MEA eligibility so bias-carrying nodes dispatch to the unfused fallback; set `data.attention_bias`. - **`attention_data.h`** — add the `attention_bias` pointer to `GroupQueryAttentionData`. - **`group_query_attention_impl.cu`** — pass the real pointer and broadcast flags in `UnfusedGqaAttention` (previously hardcoded `nullptr`/`false`). Why each fused path stays disqualified with a bias: - flash / flash fast-decode: `flash_api.h` has no bias parameter (same exclusion as MHA and the ONNX `Attention`-op CUDA kernel). - XQA: no bias parameter. - cuDNN SDPA: GQA's cuDNN path is bottom-right causal, which cuDNN documents as incompatible with a bias (`multihead_attention.cc` has the same restriction). - cutlass MEA: the wrapper computes the bias row stride from `kv_sequence_length`, which GQA sets to the KV-cache capacity (`seqlen_present_kv_cache`) rather than `total_sequence_length`, so rows would be misaligned under past/present buffer sharing. Left for a follow-up (needs an explicit `attn_bias_strideM` in `MemoryEfficientAttentionParams`). Kept `NOT_IMPLEMENTED` (explicit, clear errors instead of the previous blanket rejection): bias × quantized KV cache (unfused requires `T == U`), bias × smooth-softmax/head_sink. ### Tests New `TestGQAAttentionBias` in `test_gqa.py`: prompt and past/decode parity across packed/unpacked QKV, shared/separate KV buffer, rotary, odd head sizes (40/80), and a subsequent multi-token prompt. The bias is **non-zero random** so a kernel that silently ignores the input fails parity (the harness previously modeled a zeros bias). Also fixes the test graph builder declaring the bias input's last dim as the KV-cache capacity instead of `total_sequence_length` (the shape the op validates). Full `test_gqa.py` suite: 476 tests pass, no regressions (SM89, CUDA 12.8). ### Motivation and Context Fixes #29506. Transformers.js-exported speech models carry non-causal attention patterns as `attention_bias` and currently cannot run on the CUDA EP at all, while running fine on WebGPU and CPU: - [`onnx-community/Voxtral-Mini-4B-Realtime-2602-ONNX`](https://huggingface.co/onnx-community/Voxtral-Mini-4B-Realtime-2602-ONNX) (streaming ASR) - [`onnx-community/cohere-transcribe-03-2026-ONNX`](https://huggingface.co/onnx-community/cohere-transcribe-03-2026-ONNX) End-to-end validation with this patch on an RTX 4070 SUPER: the full Voxtral-Mini-4B-Realtime streaming pipeline (q4f16) transcribes correctly at **RTF 0.23** (31 s clip, 25.6 tok/s sustained decode, 6.4 GB VRAM) — on this workload the unfused-attention path outperforms the same model on the WebGPU EP (RTF 0.26). (While validating, an unrelated pre-existing issue surfaced: `GroupQueryAttentionFusion` breaks graphs whose GQA nodes carry >9 inputs — filed as #29524.) --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Author
Parents
Loading