microsoft / microsoft/onnxruntime

GroupQueryAttentionFusion produces an invalid graph when a GQA node has more than 9 inputs (e.g. attention_bias) — "sum of input arg count is not equal to size of input defs"

Open
#29,524 0 comments 1 reaction 2 assignees Claimed by @tianleiwu View on GitHub
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

### Describe the issue

`GroupQueryAttentionFusion` (CUDA-EP L2 optimizer) rewrites a matched GQA node's inputs to a fixed 9-element list and only updates the arg counts for the two inputs it adds:

https://github.com/microsoft/onnxruntime/blob/d7efadab308ed9bde763d8b014e30c2e5664b614/onnxruntime/core/optimizer/group_query_attention_fusion.cc#L496-L513

```cpp
const std::array gqa_input_defs{
&matmul_or_nbits_output,
&empty_node_arg,
&empty_node_arg,
past_key_values_key_arg,
past_key_values_value_arg,
seqlens_k,
total_seq_len,
cos_cache_arg,
sin_cache_arg}; // <- exactly 9 defs, up to input #8 (sin_cache)

auto& gqa_input_args = node.MutableInputArgsCount();
gqa_input_args[7] = 1;
gqa_input_args[8] = 1; // <- entries for inputs 9+ keep their old count

auto& gqa_node_input_defs = node.MutableInputDefs();
gqa_node_input_defs.assign(gqa_input_defs.begin(), gqa_input_defs.end());
```

If the original node uses any optional input beyond `sin_cache` — `position_ids` (#9), `attention_bias` (#10), `head_sink` (#11) — the input-arg-count array still contains a `1` for those slots while the input defs were truncated to 9, so graph resolve fails at session creation:

```
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL :
This is an invalid model. The sum of input arg count is not equal to size of input
defs in node (/model/layers.0/attn/GroupQueryAttention)
```

Note the fusion also silently **drops** those inputs, so even if the count array were fixed up, the fused node would compute the wrong thing — the fusion should skip nodes that use inputs it does not preserve.

### To reproduce

Any model whose GQA nodes carry `attention_bias`, loaded with the CUDA EP at default optimization level. Concrete public model (its audio encoder uses bidirectional attention expressed via `attention_bias`):

```python
import onnxruntime as ort

M = "onnx/audio_encoder_q4f16.onnx" # from onnx-community/Voxtral-Mini-4B-Realtime-2602-ONNX

ort.InferenceSession(M, providers=["CUDAExecutionProvider"]) # -> invalid model (above)

ort.InferenceSession(M, providers=["CUDAExecutionProvider"], # -> loads fine
disabled_optimizers=["GroupQueryAttentionFusion"])
```

`graphOptimizationLevel = basic` also avoids it. CPU EP is unaffected (the fusion is CUDA-only), which makes this a silent CUDA-only load failure for affected models.

### Suggested fix

In `GroupQueryAttentionFusion::ApplyImpl`, skip fusion when the GQA node has a bound input def beyond index 8 (or preserve inputs 9+ in the rewritten def list and size the arg-count array accordingly).

### Environment

- ONNX Runtime: main @ d7efadab (also 1.28 dev build), Linux x64, CUDA 12.8, RTX 4070 SUPER
- Related: #29506 (CUDA GQA kernel itself also rejects `attention_bias`; both need addressing before bias-carrying GQA models can run on CUDA)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.