[BugFix] Fix Qwen3 Omni talker mtp torch.compile startup error (#1104)

Signed-off-by: ram16g <anlianfengjie@163.com>
Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Co-authored-by: ram16g <anlianfengjie@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/worker/test_omni_gpu_model_runner.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from types import SimpleNamespace
 
 import torch
 
@@ -86,6 +87,12 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
 
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
 
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
     # Initialize per-request embeds (batch-major inside talker_mtp_inputs_embeds)
     runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
     runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -458,10 +458,6 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
                 TopPLogitsWarper(top_p=0.8),
             ]
         )
-        max_batch_size = vllm_config.scheduler_config.max_num_seqs
-        self.position_ids_buffer = []
-        for layer_idx in range(self.num_code_groups - 1):
-            self.position_ids_buffer.append(torch.arange(layer_idx + 2, dtype=torch.int64).repeat(max_batch_size))
 
         compilation_config = get_current_vllm_config().compilation_config
         if prefix in compilation_config.static_forward_context:
@@ -508,13 +504,11 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
         batch_size = current_input.shape[0]
 
         # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
-        begin_pos = 0
         for layer_idx in range(self.num_code_groups - 1):
-            # Input for this layer: [last_talker_hidden, prev_layer_embed]
-
+            seq_len = layer_idx + 2
+            # Compute position_ids dynamically to avoid torch.compile specializing batch_size
+            position_ids = torch.arange(seq_len, device=current_input.device, dtype=torch.int64).repeat(batch_size)
             # Forward through code_predictor model
-            position_ids = self.position_ids_buffer[layer_idx][: (layer_idx + 2) * batch_size]  # [1, seq_len]
-            begin_pos += layer_idx + 2
             outputs = self.model(
                 inputs_embeds=current_input,
                 attention_mask=None,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -135,8 +135,11 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
         )
+        max_batch_size = max(
+            vllm_config.scheduler_config.max_num_seqs, vllm_config.compilation_config.max_cudagraph_capture_size
+        )
         self.layer0_embed_buffer = torch.zeros(
-            (vllm_config.scheduler_config.max_num_seqs, 1, self.config.text_config.hidden_size),
+            (max_batch_size, 1, self.config.text_config.hidden_size),
             dtype=vllm_config.model_config.dtype,
         )
 

vllm_omni/worker/gpu_model_runner.py

@@ -52,12 +52,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                     self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
                 )
             hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size
-            self.talker_mtp_input_ids = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
+            max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
+            self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
             self.talker_mtp_inputs_embeds = self._make_buffer(
-                self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False
+                max_batch_size, hidden_size, dtype=self.dtype, numpy=False
             )
-            self.last_talker_hidden = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
-            self.text_step = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
+            self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+            self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
 
     def _init_mrope_positions(self, req_state: CachedRequestState):
         """Initialize M-RoPE positions for multimodal inputs.
@@ -645,7 +646,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
                     num_tokens_padded_talker_mtp = num_tokens_padded
                     if num_tokens_padded_talker_mtp == self.max_num_tokens:
-                        num_tokens_padded_talker_mtp = min(num_tokens_padded, self.max_num_reqs)
+                        num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
                     outputs = self.talker_mtp(
                         self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
                         self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
@@ -1075,10 +1076,11 @@ class OmniGPUModelRunner(GPUModelRunner):
             max_num_scheduled_tokens=1,
             use_cascade_attn=False,
         )
-        req_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size]
-        req_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size]
-        last_talker_hidden = self.last_talker_hidden.gpu[:decode_batch_size]
-        text_step = self.text_step.gpu[:decode_batch_size]
+        num_tokens_padded = batch_desc.num_tokens
+        req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded]
+        req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded]
+        last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded]
+        text_step = self.text_step.gpu[:num_tokens_padded]
         with set_forward_context(
             None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):