[BugFix] Fix V2 model runner profile_run crash for Omni Talker stages

[Sy0307]

Purpose

Fix three V2 model runner bugs that crash during profile_run for Omni Talker stages (Qwen3-Omni and Qwen2.5-Omni). These bugs block V2 runner adoption for all Omni models with Talker stages.

Reproduces CI failures in buildkite builds #6704 (H100) and #6691 (L4).

Root causes and fixes:

1. encoder_runner buffer shape mismatch — Talker shares Qwen3OmniMoeForConditionalGeneration architecture with Thinker, so multimodal registry marks supports_mm_inputs=True. encoder_runner allocates inputs_embeds buffer using hf_text_config.hidden_size (2048) but Talker's actual embedding dim is 1024 → RuntimeError: expanded size (2048) must match existing size (1024). Fix: disable MM path for models with has_preprocess (embeddings injected via run_preprocess(), not encoder_runner).

2. OmniOutput not handled in execute_model — OmniOutput is a NamedTuple (len=4), but the tuple intercept only checks len==2, leaving hidden_states as a non-tensor → AssertionError. Fix: check isinstance(OmniOutput) first and extract text_hidden_states.

3. Qwen2.5-Omni M-RoPE signature mismatch — V2 RopeState calls get_mrope_input_positions(input_tokens, mm_features) but the old signature required keyword-only args (hf_config, image_grid_thw, etc.) → TypeError. Fix: new signature accepts all params as optional, delegates to thinker.get_mrope_input_positions() which extracts grid data from mm_features via gather_kwargs. Old implementation preserved as _get_mrope_input_positions_v1.

4. Qwen2.5-Omni not in _OMNI_ARCHITECTURES — V2 fell back to DefaultModelState (no run_preprocess) → AttributeError. Fix: register Qwen2_5OmniForConditionalGeneration.

Test Plan

• Qwen3-Omni 30B V1 text-only (H20 2xGPU) — pass
• Qwen3-Omni 30B V2 text-only (H20 2xGPU) — pass
• Qwen2.5-Omni 7B V2 text-only (H20 2xGPU) — pass
• Qwen3-TTS V1 end2end Base (H20) — pass
• Qwen3-TTS V2 end2end Base (H20) — pass

Test Result

All 5 test configurations pass with returncode=0 and Processed prompts: 100%.

V2 runner previously crashed at profile_run for all Omni Talker stages; after fix all stages initialize and run inference successfully.

cc @Fattysand @tzhouam

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[hsliuustc0106]

For bug 2 (OmniOutput handling), you now set self._last_aux_output = None for OmniOutput but preserve it for the (tensor, dict) tuple path. Is this intentional? If OmniOutput returns aux_dict, should we store it like the tuple path does?

[hsliuustc0106]

BLOCKER scan:

Category	Result
Correctness	PASS
Reliability/Safety	PASS
Breaking Changes	PASS
Test Coverage	PASS - Local test results provided for 5 configurations
Documentation	PASS - PR description is comprehensive
Security	PASS

OVERALL: NO BLOCKERS

VERDICT: COMMENT

---

This is a focused bug fix that addresses four specific issues with the V2 model runner for Omni Talker stages:

1. encoder_runner buffer shape mismatch - Correctly disabled MM path for preprocess models to avoid embedding dimension mismatch
2. OmniOutput handling - Added isinstance check for OmniOutput before the tuple intercept
3. Qwen2.5-Omni M-RoPE signature - Updated get_mrope_input_positions to accept optional parameters
4. Qwen2.5-Omni registration - Added Qwen2_5OmniForConditionalGeneration to _OMNI_ARCHITECTURES

The fixes are well-reasoned and the local test results show all 5 test configurations passing successfully.

Minor suggestion: Consider adding a regression test for these specific bugs to prevent them from reoccurring. While the existing tests cover the happy path, a targeted test that verifies the fix (e.g., checking that profile_run succeeds for Talker stages with has_preprocess=True) would be valuable.

Note: DCO check is currently failing - please add commit sign-off (e.g., git commit --amend --signoff and force push).

[tzhouam]

for the hidden state extraction logic can we do like the following to handle last_aux_output?

if isinstance(model_output, tuple) and len(model_output) == 2:
    first, second = model_output
    if isinstance(first, torch.Tensor):
        self._last_aux_output = second
        if hasattr(self.model, "_last_captured_layers"):
            self.model._last_captured_layers = second
elif isinstance(model_output, OmniOutput) or isinstance(model_output, torch.Tensor):
    first = model_output
else:
    raise TypeError("Error Message")
hidden_states = first if not isinstance(model_output, OmniOutput) else model_output.text_hidden_states
```bash