SIGN IN SIGN UP

Fix Qwen2.5-Omni / Qwen3-Omni-MoE generation with a compilable cache (#47872)

* Fix Qwen2.5-Omni / Qwen3-Omni-MoE generation with a compilable cache

generate() replaces the 2D attention mask by a dict of 4D masks when the cache is compilable, but these models rebuild the multimodal rope positions from the 2D padding mask inside forward(), which then crashes. Override create_masks_for_generate to keep the 2D mask and let the text model create the 4D masks as it already does.

* Compute mrope positions in `generate` instead of overriding masks

The thinker's `forward` recomputed the 3D mrope positions from the 2D
padding mask. When a compilable cache is used, `generate` replaces the 2D
mask with a dict of 4D masks, so that code crashed with
`TypeError: unsupported operand type(s) for -: 'int' and 'dict'`.

The previous workaround kept the 2D mask alive through
`create_masks_for_generate`, which meant the mask was rebuilt on every
decoding step and defeated CUDA graphs.

Instead, sync with Qwen2.5-VL: prepare the packed `[4, batch, seq]`
position ids once in `generate` via `_prepare_position_ids_for_generation`
and move the mrope computation into `compute_3d_position_ids`. The text
model already accepted that layout, so the mask no longer needs to be
overridden.

`_expand_inputs_for_generation` is overridden as well because the batch
dimension of the packed position ids is 1, not 0.

This lets the thinker declare `_can_compile_fullgraph = True` and enables
the static cache and compile tests that were skipped until now.

* Shorten comments
J
jiqing-feng committed
279dbcf4d66051c6e184d8d6960e96294d830f11
Parent: b67f702
Committed by GitHub <noreply@github.com> on 8/19/2026, 10:56:44 AM