[CUDA] Support attention_bias in GroupQueryAttention via the unfused path - #29525
Conversation
…path The GQA CUDA op rejected the optional attention_bias input (input 10) outright, while the CPU and WebGPU EPs implement it. The unfused fallback kernel (LaunchUnfusedAttention, issue microsoft#28195) already takes an additive bias with dim-0/dim-1 broadcast and per-batch seqlens - the op just never passed it. This wires the bias through: - group_query_attention.cc: drop the blanket rejection; validate bias element type; set broadcast_attn_bias_dim_0/1 from the bias shape; disqualify the bias-incapable fused paths (XQA, cuDNN SDPA, flash, flash fast-decode, cutlass MEA) so dispatch reaches the unfused fallback; keep bias x quantized-KV and bias x smooth-softmax/head_sink as NOT_IMPLEMENTED. - attention_data.h: add attention_bias pointer to GroupQueryAttentionData. - group_query_attention_impl.cu: pass the real bias pointer and broadcast flags in UnfusedGqaAttention. MEA is left disqualified for now: the cutlass wrapper derives the bias row stride from kv_sequence_length, which GQA sets to the KV-cache capacity rather than total_sequence_length, so it would read misaligned rows under past/present buffer sharing. Can be enabled later with an explicit bias stride. Tests: TestGQAAttentionBias (prompt + past/decode, packed/unpacked, shared/separate KV buffer, rotary, odd head sizes, subsequent prompt) with non-zero random bias so a kernel ignoring the input fails parity. Also fixes the test builder declaring the bias input with the KV-cache capacity instead of total_sequence_length as its last dim. Fixes microsoft#29506 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR adds CUDA EP support for the optional attention_bias input of com.microsoft.GroupQueryAttention by routing bias-carrying nodes to the existing unfused attention fallback (the fused CUDA paths don’t accept a bias parameter).
Changes:
- Plumbs
attention_biasthrough CUDA GQA dispatch and data structures, and disables fused-path eligibility when bias is present so execution falls back to the unfused implementation. - Extends the unfused GQA CUDA path to pass the bias pointer and broadcast flags into
LaunchUnfusedAttention. - Updates
test_gqa.pyto declare the bias shape correctly for past/present and adds focused CUDA parity coverage with non-zero random bias.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc | Accepts/validates attention_bias, carries broadcast flags, and ensures fused paths are skipped when bias is present. |
| onnxruntime/contrib_ops/cuda/bert/attention_data.h | Adds attention_bias pointer to GroupQueryAttentionData. |
| onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu | Passes attention_bias and broadcast flags into the unfused attention launch. |
| onnxruntime/test/python/transformers/test_gqa.py | Fixes bias input shape declaration, uses non-zero bias generation, and adds focused CUDA parity tests. |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
tianleiwu
left a comment
tianleiwu
left a comment
There was a problem hiding this comment.
Reviewed the dispatch plumbing that routes attention_bias-carrying GQA nodes to the unfused fallback. The approach is sound and the key correctness claims hold up against the code:
- All four fused paths correctly gate on
!has_attention_bias(XQA, cuDNN SDPA, flash — and flash-fast-decode transitively viause_flash_attention, plus MEA), so a bias-carrying node deterministically reachesUnfusedGqaAttention. - The unfused softmax indexes the bias row with stride
total_kv_length(== parameters.total_sequence_length), which matches the bias tensor's last dim, and the per-batchseqlens_kcutoff keeps reads within[0, total_sequence_length). So the past/present buffer-sharing hazard that disqualifies MEA (row stride derived from the KV-cache capacity) does not apply on this path — the MEA exclusion rationale is accurate. masked_bias_valuestays at its default 0 for GQA (no composed mask sentinel), so the fully-masked-row guard degrades to the intended zero-row behavior via thes_max == -infbranch.
One test-coverage gap is worth addressing before merge (inline comment). I also concur with the two review notes already open on the current head, which still apply:
TestGQAAttentionBiasis missing a CUDA skip guard; sibling CUDA classes such asTestGQAQKNormall carry@unittest.skipIf(not has_cuda_device(80), ...). Without it, CPU-only / no-GPU runs will fail allocating CUDA tensors.- No
NumDimensions() == 4guard exists before the 4-Dattention_biasshape is indexed (TensorShape::operator[]is unchecked, so a malformed rank-<4 bias is UB rather than a cleanINVALID_ARGUMENT). This is a pre-existing gap in the sharedCheckCustomAttentionInputshelper (also reached by the CPU/WebGPU EPs), but this PR is the first to exercise it on CUDA.
…coverage - CheckCustomAttentionInputs: validate attention_bias is 4D before indexing its dims (TensorShape::operator[] is unchecked; a rank<4 bias was UB instead of a clean INVALID_ARGUMENT). Shared helper, so the CPU and WebGPU EPs get the same guard. - TestGQAAttentionBias: add the missing skipIf gate (has_cuda_device(53), matching TestMemoryEfficientGQA — the unfused path has no arch floor). Fixes the CPU-torch CI failures. - Cover both untested broadcast-flag combinations in BOTH suites: batch-broadcast biases (dim0 == 1, prompt + past incl. shared and separate KV buffer) and per-head biases (dim1 == num_heads, prompt + past). A batch-shared bias cannot carry the harness's per-batch -10000 tails, so the dim0-broadcast past cases skip them — passing parity there also verifies the kernel's seqlens_k cutoff and the reference mask never read past each batch's valid length.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
Can the C++ unit tests be updated to also validate the new code? AFAIK the python tests are only run post-merge, so a change that broke would not be detected in a PR. onnxruntime\test\contrib_ops\group_query_attention_op_test.cc |
The Python parity tests need a CUDA-enabled torch that the PR CI agents don't have, so they only exercise the CUDA kernel post-merge. Add a C++ gtest that runs in the GPU CI on every PR: it runs the same prompt case with a non-null attention_bias on CUDA (fp16) and CPU (fp32, the reference EP that already implements the input) and compares the outputs, covering the three bias shapes that drive broadcast_attn_bias_dim_0/1 — [batch,1,S,S] (default), [1,1,S,S] (dim0 broadcast) and [batch,heads,S,S] (per-head). Skips cleanly when no CUDA EP is available. Verified on SM 8.9: the new test passes at fp16 tolerance 0.02, full GroupQueryAttentionTest suite 56 passed / 12 skipped (WebGPU) / 0 failed.
|
@tianleiwu the earlier review feedback (broadcast-flag coverage, the Added in b3d4a79 — Verified locally on SM 8.9: passes at fp16 tolerance 0.02, and the full |
Description
com.microsoft.GroupQueryAttention's optionalattention_biasinput (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-batchseqlens_k, causal/sliding-window masking and softcap — the op just passedattn_bias=nullptr. The change is dispatch plumbing:group_query_attention.cc— remove the blanket rejection; validate the bias element type; setbroadcast_attn_bias_dim_0/1from the bias shape (the fields already exist onAttentionParameters); add!has_attention_biasto the XQA / cuDNN SDPA / flash / flash-fast-decode / MEA eligibility so bias-carrying nodes dispatch to the unfused fallback; setdata.attention_bias.attention_data.h— add theattention_biaspointer toGroupQueryAttentionData.group_query_attention_impl.cu— pass the real pointer and broadcast flags inUnfusedGqaAttention(previously hardcodednullptr/false).Why each fused path stays disqualified with a bias:
flash_api.hhas no bias parameter (same exclusion as MHA and the ONNXAttention-op CUDA kernel).multihead_attention.cchas the same restriction).kv_sequence_length, which GQA sets to the KV-cache capacity (seqlen_present_kv_cache) rather thantotal_sequence_length, so rows would be misaligned under past/present buffer sharing. Left for a follow-up (needs an explicitattn_bias_strideMinMemoryEfficientAttentionParams).Kept
NOT_IMPLEMENTED(explicit, clear errors instead of the previous blanket rejection): bias × quantized KV cache (unfused requiresT == U), bias × smooth-softmax/head_sink.Tests
New
TestGQAAttentionBiasintest_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 oftotal_sequence_length(the shape the op validates).Full
test_gqa.pysuite: 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_biasand currently cannot run on the CUDA EP at all, while running fine on WebGPU and CPU:onnx-community/Voxtral-Mini-4B-Realtime-2602-ONNX(streaming ASR)onnx-community/cohere-transcribe-03-2026-ONNXEnd-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:
GroupQueryAttentionFusionbreaks graphs whose GQA nodes carry >9 inputs — filed as #29524.)