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