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"

[jiangzhuo]

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:

onnxruntime/onnxruntime/core/optimizer/group_query_attention_fusion.cc

Lines 496 to 513 in d7efada

	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};
	auto& gqa_input_args = node.MutableInputArgsCount();
	gqa_input_args[7] = 1;
	gqa_input_args[8] = 1;
	// Switch GQA input defs from unfused into the fused form.
	auto& gqa_node_input_defs = node.MutableInputDefs();
	gqa_node_input_defs.assign(gqa_input_defs.begin(), gqa_input_defs.end());

    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):

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 @ d7efada (also 1.28 dev build), Linux x64, CUDA 12.8, RTX 4070 SUPER
• Related: [Feature Request] CUDA EP: support attention_bias in GroupQueryAttention (last EP missing it) #29506 (CUDA GQA kernel itself also rejects attention_bias; both need addressing before bias-carrying GQA models can run on CUDA)