[Bugfix][Qwen3-Omni] Fix synchronous Thinker pipeline parallelism (#7345)

Signed-off-by: Deep Shah <deep@socratic.co>

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -34,6 +34,7 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import (
 from vllm.model_executor.models.utils import (
     WeightsMapper,
     init_vllm_registered_model,
+    make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -53,6 +54,7 @@ from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     apply_outer_quant_config_mapping,
 )
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
+    PP_CAPTURE_PREFIX,
     Qwen3OmniMoeThinkerDummyInputsBuilder,
     Qwen3OmniMoeThinkerForConditionalGeneration,
     Qwen3OmniMoeThinkerMultiModalProcessor,
@@ -274,6 +276,18 @@ class Qwen3OmniMoeForConditionalGeneration(
         self.make_empty_intermediate_tensors = (
             self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None
         )
+        if self.model_stage == "thinker":
+            capture_indices = self._thinker_capture_layer_indices()
+            start_layer = self.thinker.language_model.model.start_layer
+            incoming_captures = [f"{PP_CAPTURE_PREFIX}{index}" for index in capture_indices if index < start_layer]
+            if incoming_captures:
+                self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+                    ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size
+                )
+
+    def _thinker_capture_layer_indices(self) -> list[int]:
+        accept_layer = getattr(self.talker_config, "accept_hidden_layer", None)
+        return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else []
 
     @classmethod
     async def buffer_realtime_audio(
@@ -441,11 +455,11 @@ class Qwen3OmniMoeForConditionalGeneration(
             # Only staged runs have a talker stage to consume the capture; in a
             # plain vLLM run capturing would waste a clone per layer per step
             # and make the thinker return a tuple stock vLLM cannot handle.
-            accept_layer = getattr(self.talker_config, "accept_hidden_layer", None)
+            capture_indices = self._thinker_capture_layer_indices()
             capture_kwargs = {}
-            if accept_layer is not None and self.is_staged_run:
+            if capture_indices:
                 capture_kwargs = {
-                    "capture_layer_indices": [0, int(accept_layer)],
+                    "capture_layer_indices": capture_indices,
                     "return_hidden_states": True,
                 }
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -118,6 +118,7 @@ from vllm.multimodal.utils import set_mm_embedding_modality
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.processor import cached_processor_from_config
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerMultiModalDataParser,
@@ -140,6 +141,8 @@ except (ImportError, ModuleNotFoundError):
 
 logger = init_logger(__name__)
 
+PP_CAPTURE_PREFIX = "capture_"
+
 _THINKER_ARCHITECTURE = "Qwen3OmniMoeThinkerForConditionalGeneration"
 
 
@@ -568,7 +571,15 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
         capture_set = set(capture_layer_indices) if capture_layer_indices else None
-        captured_hidden_states: dict[str, torch.Tensor] | None = {} if return_hidden_states else None
+        captured_hidden_states: OmniPayload | None = {} if return_hidden_states else None
+
+        if captured_hidden_states is not None and capture_set and intermediate_tensors is not None:
+            for layer_idx in capture_set:
+                if layer_idx < self.start_layer:
+                    hs = captured_hidden_states.setdefault("hidden_states", {})
+                    layers = hs.setdefault("layers", {})
+                    # Receive buffers are reused on the next step; retain an independent snapshot.
+                    layers[layer_idx] = intermediate_tensors[f"{PP_CAPTURE_PREFIX}{layer_idx}"].clone()
 
         for layer_idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
             layer_idx = layer_idx + self.start_layer
@@ -589,7 +600,15 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
                 hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
 
         if not get_pp_group().is_last_rank:
-            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+            tensors = {"hidden_states": hidden_states, "residual": residual}
+            if captured_hidden_states:
+                tensors.update(
+                    {
+                        f"{PP_CAPTURE_PREFIX}{index}": value
+                        for index, value in captured_hidden_states["hidden_states"]["layers"].items()
+                    }
+                )
+            return IntermediateTensors(tensors)
         hidden_states, _ = self.norm(hidden_states, residual)
         if captured_hidden_states is not None:
             return hidden_states, captured_hidden_states