diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 8da1ccfec..df26a1faf 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/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