microsoft / microsoft/onnxruntime
Qwen3 RotaryEmbedding fusion produces invalid position_ids shape when batch_size > 1
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
**Describe the issue**
`FusionRotaryEmbeddings.fuse()` (added in #27590 for Qwen3's on-the-fly RoPE) wires whatever tensor feeds the RoPE `MatMul` directly into the new `com.microsoft.RotaryEmbedding` node's `position_ids` input, without checking its shape.
For a normal `optimum-cli export onnx` of a Qwen3 model (no explicit `position_ids` graph input), HF's `Qwen3Model.forward()` computes `position_ids = cache_position.unsqueeze(0)` — batch axis hardcoded to size 1, independent of the real batch size. This is harmless in the *unfused* graph (broadcasting handles it), but `RotaryEmbedding`'s kernel requires `position_ids.shape[0] == batch_size` exactly (`rotary_embedding_helper.h:136-149`; only a scalar/1-element tensor or a true `(batch_size, sequence_length)` tensor is accepted — no broadcasting). So any request with `batch_size > 1` fails:
```
Input 'position_ids' dimension 0 should be of size batch_size, got 1
```
#27590's own tests don't catch this because `qwen3_model_generator.py` declares `position_ids` as an already-correctly-shaped `(batch_size, seq_len)` input — unlike real Qwen3 export.
**To reproduce**
1. `optimum-cli export onnx --model ` (no explicit `position_ids` input)
2. `python -m onnxruntime.transformers.optimizer --input /model.onnx --output /model.onnx --model_type qwen3 --float16`
3. Run inference with `batch_size > 1` → crash. `batch_size == 1` works.
Repro'd with `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`; should reproduce for any Qwen3 export without an explicit `position_ids` input.
**Urgency**: Moderate — blocks using this fusion for any serving setup with batching enabled; we're carrying a local graph-patch workaround meanwhile.
**System info**: onnxruntime-gpu 1.27.0 (released package), also present on `main`; Python API; Linux x64; CUDA EP (issue is in shared `rotary_embedding_helper.h`, not EP-specific).
Contributor guide
Research direction
Start in FusionRotaryEmbeddings.fuse() and inspect how the RoPE MatMul input is connected to position_ids. Read rotary_embedding_helper.h lines 136-149 and compare this behavior with qwen3_model_generator.py; reproduce using the stated optimum-cli export, optimizer command, and batch_size greater than 1. Done means the fused graph accepts the real Qwen3 export with batched inference and tests cover the missing explicit position_ids case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100