Assertion failure in graph optimizer Level2 after MatmulTransposeFusion on large MoE model

[justinchuby]

Repro

Loading a Qwen3.5-MoE (40-layer, 256-expert, MoE with shared expert) bf16 ONNX exported with mobius using --ep cuda --dtype f16 (or bf16) aborts during InferenceSession initialization right after MatmulTransposeFusion finishes.

Symptom

[I:onnxruntime:, graph_transformer.cc:15 Apply] GraphTransformer MatmulTransposeFusion modified: 0 with status: OK
/opt/rh/gcc-toolset-14/root/usr/include/c++/14/bits/stl_vector.h:1130:
  constexpr std::vector<_Tp, _Alloc>::reference std::vector<_Tp, _Alloc>::operator[](size_type)
  [... NodeArg* ...]: Assertion '__n < this->size()' failed.
Aborted (core dumped)

Env

• onnxruntime-gpu 1.27.0.dev20260511001 (CUDA 12)
• CUDAExecutionProvider on NVIDIA H200
• Model: Qwen/Qwen3.6-35B-A3B exported via mobius with com.microsoft::GroupQueryAttention, com.microsoft::LinearAttention, com.microsoft::CausalConvWithState
• 40-layer hybrid (full + linear attention) + 256-expert MoE per layer + shared expert
• Decoder model: ~64 GB (f16 with weights, 1.6 M nodes)

What I tried

• disabled_optimizers=['BiasGeluFusion','BiasSoftmaxFusion','BiasDropoutFusion','MatMulScaleFusion','SkipLayerNormFusion','LayerNormFusionL2','QuickGeluFusion','GeluApproximation'] — still aborts after MatmulTransposeFusion.
• graph_optimization_level=ORT_ENABLE_BASIC avoids the abort but then fails EP placement on CausalConvWithState (needs the function-inlining transformer that only runs at higher levels).

Likely cause

Some Level2 GraphTransformer running between MatmulTransposeFusion and the next logged transformer reads past the end of a node-arg list — probably triggered by the unusual graph shape (10240+ MatMul nodes from per-expert MLPs, hybrid attention).

[yuslepukhin]

Perhaps, this is related: #29153

[namgyu-youn]

Reproduced on current main (361184e619) with a 748-byte model — fix in #29707.

GroupQueryAttentionFusion (Level 2, CUDA EP only — which is why ORT_ENABLE_BASIC or CPU EP avoids the abort) matches GQA input producers by OpType() == "RotaryEmbedding" without a domain check, and reads inputs 2/3 assuming the contrib signature (input, position_ids, cos_cache, sin_cache). The ONNX-domain opset-23 op is (X, cos_cache, sin_cache, position_ids?), so a spec-valid 3-input node reads out of bounds:

stl_vector.h:1253: std::vector<_Tp, _Alloc>::operator[] [with _Tp = onnxruntime::NodeArg*; ...]:
Assertion '__n < this->size()' failed.
#5 GroupQueryAttentionFusion::ApplyImpl at group_query_attention_fusion.cc:365

(#29153 is unrelated — different transformer, runs before MatmulTransposeFusion. #29552 covers only >9-input GQA nodes.)

make_repro_model.py

# Minimal repro for microsoft/onnxruntime#28604: a 3-input ONNX-domain (opset 23)
# RotaryEmbedding (optional position_ids omitted) feeding an unfused GQA node makes
# GroupQueryAttentionFusion read RotaryEmbedding InputDefs()[3] out of bounds.
import numpy as np
import onnx
from onnx import TensorProto as T, helper

fp16 = T.FLOAT16
vi = helper.make_tensor_value_info
cache = np.ones((8, 2), dtype=np.float16).tobytes()

graph = helper.make_graph(
    [
        helper.make_node("RotaryEmbedding", ["x", "cos_cache", "sin_cache"], ["q"], num_heads=2),
        helper.make_node(
            "GroupQueryAttention",
            ["q", "key", "value", "past_key", "past_value", "seqlens_k", "total_sequence_length"],
            ["out", "present_key", "present_value"],
            domain="com.microsoft", num_heads=2, kv_num_heads=1),
    ],
    "gqa_rotary3_repro",
    inputs=[
        vi("x", fp16, [1, 1, 8]), vi("key", fp16, [1, 1, 4]), vi("value", fp16, [1, 1, 4]),
        vi("past_key", fp16, [1, 1, 4, 4]), vi("past_value", fp16, [1, 1, 4, 4]),
        vi("seqlens_k", T.INT32, [1]), vi("total_sequence_length", T.INT32, []),
    ],
    outputs=[
        vi("out", fp16, [1, 1, 8]),
        vi("present_key", fp16, [1, 1, "total_seq", 4]), vi("present_value", fp16, [1, 1, "total_seq", 4]),
    ],
    initializer=[
        helper.make_tensor("cos_cache", fp16, [8, 2], cache, raw=True),
        helper.make_tensor("sin_cache", fp16, [8, 2], cache, raw=True),
    ],
)

model = helper.make_model(
    graph, ir_version=10,
    opset_imports=[helper.make_opsetid("", 23), helper.make_opsetid("com.microsoft", 1)])
onnx.checker.check_model(model, full_check=True)
onnx.save(model, "gqa_rotary3_repro.onnx")
print("wrote gqa_rotary3_repro.onnx")

To confirm your export hits this path:

import onnx
m = onnx.load("model.onnx", load_external_data=False)
gqa_in = {i for n in m.graph.node if n.op_type == "GroupQueryAttention" for i in n.input}
for n in m.graph.node:
    if n.op_type == "RotaryEmbedding" and n.domain in ("", "ai.onnx") and set(n.output) & gqa_in:
        print("onnx-domain rotary feeding GQA:", n.name, "inputs:", len(n.input))

Developed with AI assistance (Claude Code); reviewed by the author.