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

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

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -22,6 +22,10 @@ both stages on vLLM's Triton unquantized MoE path and avoids the unstable FlashI
 some GPUs.
 Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
 
+For Thinker pipeline parallelism, configure `pipeline_parallel_size` and `devices`
+on the Thinker stage and set `async_scheduling: false`. Keep Talker and Code2Wav
+at `pipeline_parallel_size: 1`. Audio chunk streaming can remain enabled.
+
 To explicitly utilize a custom deployment YAML, specify the configuration path:
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \

tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml

@@ -0,0 +1,66 @@
+async_chunk: true
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      initial_codec_chunk_frames: 4
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+stages:
+- stage_id: 0
+  max_num_batched_tokens: 4096
+  max_num_seqs: 2
+  gpu_memory_utilization: 0.65
+  moe_backend: triton
+  trust_remote_code: true
+  enable_prefix_caching: false
+  devices: 0,1
+  default_sampling_params:
+    temperature: 0.0
+    max_tokens: 128
+  enforce_eager: true
+  max_model_len: 4096
+  async_scheduling: false
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 2
+- stage_id: 1
+  max_num_batched_tokens: 4096
+  max_num_seqs: 2
+  gpu_memory_utilization: 0.2
+  moe_backend: triton
+  trust_remote_code: true
+  enable_prefix_caching: false
+  devices: '1'
+  input_connectors:
+    from_stage_0: connector_of_shared_memory
+  default_sampling_params:
+    temperature: 0.0
+    top_k: 50
+    max_tokens: 1024
+    repetition_penalty: 1.05
+  enforce_eager: true
+  max_model_len: 4096
+  async_scheduling: false
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 1
+- stage_id: 2
+  max_num_batched_tokens: 4096
+  max_num_seqs: 2
+  gpu_memory_utilization: 0.05
+  enforce_eager: true
+  trust_remote_code: true
+  enable_prefix_caching: false
+  enable_chunked_prefill: false
+  async_scheduling: false
+  devices: '1'
+  input_connectors:
+    from_stage_1: connector_of_shared_memory
+  default_sampling_params:
+    temperature: 0.0
+    top_p: 1.0
+    top_k: -1
+    max_tokens: 65536
+    repetition_penalty: 1.1
+  max_model_len: 4096
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 1

tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Required Thinker snapshots must survive consecutive pipeline ranks.
+
+Actual forward methods with deterministic decoder doubles test the contract
+without model weights. Distributed serving remains a separate integration gate.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.qwen3_omni import qwen3_omni as combined_module
+from vllm_omni.model_executor.models.qwen3_omni import qwen3_omni_moe_thinker as thinker_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel]
+
+
+class _Layer(nn.Module):
+    def forward(self, positions, hidden_states, residual):
+        total = hidden_states if residual is None else hidden_states + residual
+        return total * 0.01 + 0.125, total
+
+
+class _Norm(nn.Module):
+    def forward(self, hidden_states, residual):
+        return hidden_states + residual, None
+
+
+def _forward(monkeypatch, start, end, embeddings, intermediate=None, *, capture=True):
+    rank = SimpleNamespace(is_first_rank=start == 0, is_last_rank=end == 48)
+    monkeypatch.setattr(thinker_module, "get_pp_group", lambda: rank)
+    model = object.__new__(thinker_module.Qwen3MoeLLMModel)
+    nn.Module.__init__(model)
+    model.start_layer = start
+    model.end_layer = end
+    model.layers = nn.ModuleList([_Layer() for _ in range(48)])
+    model.norm = _Norm()
+    return model.forward(
+        input_ids=None,
+        positions=torch.arange(embeddings.shape[0]),
+        inputs_embeds=embeddings,
+        intermediate_tensors=intermediate,
+        capture_layer_indices=[0, 24] if capture else None,
+        return_hidden_states=capture,
+    )
+
+
+@pytest.mark.parametrize("tokens", [1, 4])
+@pytest.mark.parametrize("boundaries", [(0, 12, 48), (0, 24, 48), (0, 36, 48), (0, 16, 32, 48)])
+def test_required_captures_match_unsplit_forward(monkeypatch, tokens, boundaries):
+    embeddings = torch.arange(tokens * 8, dtype=torch.float32).reshape(tokens, 8) / 8
+    reference, captures = _forward(monkeypatch, 0, 48, embeddings)
+    intermediate = None
+    for start, end in zip(boundaries, boundaries[1:]):
+        result = _forward(monkeypatch, start, end, embeddings, intermediate)
+        if end < 48:
+            assert isinstance(result, IntermediateTensors)
+            assert all(isinstance(value, torch.Tensor) for value in result.tensors.values())
+            intermediate = result
+    actual, actual_captures = result
+    torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+    actual_layers = actual_captures["hidden_states"]["layers"]
+    assert set(actual_layers) == {0, 24}
+    for key in (0, 24):
+        torch.testing.assert_close(actual_layers[key], captures["hidden_states"]["layers"][key], rtol=0, atol=0)
+    # A later receive into the rank's persistent buffers cannot rewrite snapshots.
+    assert intermediate is not None
+    for value in intermediate.tensors.values():
+        value.zero_()
+    for key in (0, 24):
+        torch.testing.assert_close(actual_layers[key], captures["hidden_states"]["layers"][key], rtol=0, atol=0)
+
+
+def test_capture_disabled_keeps_original_intermediate_schema(monkeypatch):
+    embeddings = torch.ones(4, 8)
+    first = _forward(monkeypatch, 0, 24, embeddings, capture=False)
+    assert set(first.tensors) == {"hidden_states", "residual"}
+    actual, captures = _forward(monkeypatch, 24, 48, embeddings, first, capture=False)
+    reference, _ = _forward(monkeypatch, 0, 48, embeddings, capture=False)
+    assert captures is None
+    torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "start,expected", [(0, set()), (16, {"capture_0"}), (24, {"capture_0"}), (32, {"capture_0", "capture_24"})]
+)
+@pytest.mark.parametrize("staged", [False, True])
+def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, expected, staged):
+    thinker_config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=8))
+    config = SimpleNamespace(
+        thinker_config=thinker_config,
+        talker_config=SimpleNamespace(accept_hidden_layer=24),
+        code2wav_config=SimpleNamespace(),
+        tts_bos_token_id=1,
+        tts_eos_token_id=2,
+        tts_pad_token_id=3,
+    )
+    model_config = SimpleNamespace(hf_config=config, multimodal_config=None)
+    if staged:
+        model_config.model_stage = "thinker"
+    vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None)
+    thinker = nn.Module()
+    thinker.weight = nn.Parameter(torch.ones(1))
+    thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start))
+    thinker.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states", "residual"], 8)
+    monkeypatch.setattr(combined_module, "init_vllm_registered_model", lambda **kwargs: thinker)
+    model = combined_module.Qwen3OmniMoeForConditionalGeneration(vllm_config=vllm_config)
+    buffers = model.make_empty_intermediate_tensors(4, dtype=torch.float32, device="cpu")
+    assert set(buffers.tensors) == {"hidden_states", "residual"} | (expected if staged else set())
+    assert all(value.shape == (4, 8) for value in buffers.tensors.values())

tests/worker/test_pp_intermediate_output.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Non-final pipeline outputs must reach vLLM's existing PP sender unchanged."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel]
+
+
+def test_nonfinal_forward_bypasses_final_output_conversion(monkeypatch):
+    intermediate = IntermediateTensors({"hidden_states": torch.ones(4, 8), "residual": torch.ones(4, 8)})
+    monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: intermediate)
+
+    def final_only_converter(*args, **kwargs):
+        pytest.fail("A non-final rank must not convert intermediate tensors to OmniOutput")
+
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(make_omni_output=final_only_converter, have_multimodal_outputs=True)
+    runner._build_model_kwargs_extra = lambda: {}
+    result = runner._model_forward()
+    assert result is intermediate
+    hidden, multimodal = runner.extract_multimodal_outputs(result)
+    assert hidden is intermediate
+    assert multimodal == {}
+
+
+def test_final_tensor_still_uses_model_converter(monkeypatch):
+    tensor = torch.ones(4, 8)
+    monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: tensor)
+    converted = object()
+    calls = []
+
+    def convert(value, **kwargs):
+        calls.append(value)
+        return converted
+
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(make_omni_output=convert)
+    runner._build_model_kwargs_extra = lambda: {}
+    assert runner._model_forward() is converted
+    assert calls == [tensor]

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

vllm_omni/worker/gpu_model_runner.py

@@ -34,6 +34,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 
 from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding
 from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs
@@ -892,8 +893,8 @@ class OmniGPUModelRunner(GPUModelRunner):
 
     @torch.inference_mode()
     def extract_multimodal_outputs(
-        self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput
-    ) -> tuple[Any, Any]:
+        self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput | IntermediateTensors
+    ) -> tuple[torch.Tensor | IntermediateTensors, OmniPayload | None]:
         if (
             hasattr(self.model, "have_multimodal_outputs")
             and self.model.have_multimodal_outputs
@@ -902,7 +903,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             text_hidden_states = hidden_states.text_hidden_states
             multimodal_outputs = hidden_states.multimodal_outputs
 
-        elif isinstance(hidden_states, torch.Tensor):
+        elif isinstance(hidden_states, (torch.Tensor, IntermediateTensors)):
             text_hidden_states = hidden_states
             multimodal_outputs = {}
         elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple):
@@ -1225,6 +1226,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             else:
                 hidden_states = outputs
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
+            if isinstance(hidden_states, IntermediateTensors):
+                # Profiling/graph capture needs a tensor to select dummy sampler rows.
+                hidden_states = hidden_states["hidden_states"]
             if self.speculative_config and (
                 self.speculative_config.use_eagle()
                 or self.speculative_config.uses_draft_model()
@@ -2072,7 +2076,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             **model_kwargs,
             **model_kwargs_extra,
         )
-        if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
+        if not isinstance(model_output, (OmniOutput, IntermediateTensors)) and hasattr(self.model, "make_omni_output"):
             model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra)
         # Cache model output so later sample_tokens can consume multimodal results.
         self._omni_last_model_output = model_output