Fix Qwen3-omni accuracy degradation from deepstack inputs under torch.compile (#3885)

Signed-off-by: Dakai An <dakaian108@gmail.com>

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -1205,11 +1205,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     ) -> IntermediateTensors | None:
         if not getattr(self, "deepstack_input_embeds", None):
             return None  # If vision tower is skipped
-        if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
-            return None
-        buf = self.deepstack_input_embeds[0].size(0)
-        if num_tokens > buf:
-            raise ValueError(f"Requested more deepstack tokens than buffer capacity: {num_tokens=} > {buf=}")
+        if num_tokens > self.deepstack_input_embeds[0].size(0):
+            self._resize_deepstack_input_embeds(num_tokens)
         n_valid = self.deepstack_input_embeds_num_tokens
         # embed_input_ids runs on unpadded length; forward uses padded num_tokens
         # (vLLM gpu_model_runner). Tail positions have no vision deepstack and must
@@ -1226,6 +1223,17 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             }
         )
 
+    def _resize_deepstack_input_embeds(self, num_tokens: int) -> None:
+        self.deepstack_input_embeds = [
+            torch.zeros(
+                num_tokens,
+                self.config.text_config.hidden_size,
+                device=self.deepstack_input_embeds[0].device,
+                dtype=self.deepstack_input_embeds[0].dtype,
+            )
+            for _ in range(self.deepstack_num_level)
+        ]
+
     def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None:
         if not getattr(self, "deepstack_input_embeds", None):
             return
@@ -1233,15 +1241,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         # set deepstack_input_embeds to buffer
         num_tokens = deepstack_input_embeds.size(1)
         if num_tokens > self.deepstack_input_embeds[0].size(0):
-            self.deepstack_input_embeds = [
-                torch.zeros(
-                    num_tokens,
-                    self.config.text_config.hidden_size,
-                    device=self.deepstack_input_embeds[0].device,
-                    dtype=self.deepstack_input_embeds[0].dtype,
-                )
-                for _ in range(self.deepstack_num_level)
-            ]
+            self._resize_deepstack_input_embeds(num_tokens)
         for idx in range(self.deepstack_num_level):
             self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx])
         self.deepstack_input_embeds_num_tokens = num_tokens