[Feature] Support Prefill-Decode disaggregation via vLLM KV transfer (#2220)

Signed-off-by: LiuBingyu <liubingyu62@gmail.com>

tests/e2e/online_serving/test_qwen3_omni.py

@@ -26,9 +26,15 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 QWEN3_OMNI_CONFIG_PATH = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
 QWEN3_OMNI_XPU_CONFIG_PATH = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")
 
+_STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs"
+_PD_SEP_CONFIG = str(_STAGE_CONFIGS_DIR / "qwen3_omni_moe_pd_ci.yaml")
 
-def get_chunk_config(config_path: str):
-    path = modify_stage_config(
+
+def get_chunk_config(config_path: str | None = None):
+    """Load qwen3_omni_ci.yaml with async_chunk modifications for streaming mode."""
+    if config_path is None:
+        config_path = str(_STAGE_CONFIGS_DIR / "qwen3_omni_ci.yaml")
+    return modify_stage_config(
         config_path,
         updates={
             "async_chunk": True,
@@ -43,7 +49,6 @@ def get_chunk_config(config_path: str):
         },
         deletes={"stage_args": {2: ["custom_process_input_func"]}},
     )
-    return path
 
 
 def get_prefix_caching_config(config_path: str):
@@ -59,11 +64,20 @@ def get_prefix_caching_config(config_path: str):
     return path
 
 
-if current_omni_platform.is_xpu():
-    stage_configs = [QWEN3_OMNI_XPU_CONFIG_PATH]
-    prefix_caching_stage_configs = [get_prefix_caching_config(QWEN3_OMNI_XPU_CONFIG_PATH)]
-else:  # MI325 GPU should share the same config as H100
-    stage_configs = [get_chunk_config(QWEN3_OMNI_CONFIG_PATH)]
+# Set VLLM_TEST_PD_MODE=1 to test PD disaggregation, default tests async_chunk mode.
+_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
+
+# Stage configs for H100/CUDA, ROCm MI325, and XPU platforms
+if current_omni_platform.is_rocm():
+    rocm_config = str(_STAGE_CONFIGS_DIR / "rocm" / "qwen3_omni_ci.yaml")
+    stage_configs = [rocm_config]
+    prefix_caching_stage_configs = [get_prefix_caching_config(rocm_config)]
+elif current_omni_platform.is_xpu():
+    xpu_config = str(_STAGE_CONFIGS_DIR / "xpu" / "qwen3_omni_ci.yaml")
+    stage_configs = [xpu_config]
+    prefix_caching_stage_configs = [get_prefix_caching_config(xpu_config)]
+else:
+    stage_configs = [_PD_SEP_CONFIG if _USE_PD else get_chunk_config(QWEN3_OMNI_CONFIG_PATH)]
     prefix_caching_stage_configs = [get_prefix_caching_config(QWEN3_OMNI_CONFIG_PATH)]
 
 # Create parameter combinations for model and stage config
@@ -116,7 +130,8 @@ def get_max_batch_size(size_type="few"):
 @pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
     """
@@ -155,7 +170,8 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_001(omni_server, openai_client) -> None:
     """

vllm_omni/engine/async_omni_engine.py

@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any
 import janus
 import torch
 from omegaconf import OmegaConf
+from vllm import envs as vllm_envs
 from vllm.engine.arg_utils import EngineArgs
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
@@ -79,6 +80,7 @@ from vllm_omni.engine.stage_init_utils import (
     setup_stage_devices,
     terminate_alive_proc,
 )
+from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import (
     inject_omni_kv_config,
     load_and_resolve_stage_configs,
@@ -902,6 +904,7 @@ class AsyncOmniEngine:
             self._initialize_janus_queues()
 
             self._initialize_stages(stage_init_timeout)
+            pd_config = self._detect_pd_config()
             orchestrator = Orchestrator(
                 request_async_queue=self.request_queue.async_q,
                 output_async_queue=self.output_queue.async_q,
@@ -910,6 +913,7 @@ class AsyncOmniEngine:
                 stage_clients=self.stage_clients,
                 output_processors=self.output_processors,
                 stage_vllm_configs=self.stage_vllm_configs,
+                pd_config=pd_config,
             )
             if not startup_future.done():
                 startup_future.set_result(asyncio.get_running_loop())
@@ -1131,6 +1135,48 @@ class AsyncOmniEngine:
             cache_config = AsyncOmniEngine._get_default_cache_config(cache_backend)
         return cache_config
 
+    def _detect_pd_config(self) -> dict[str, Any] | None:
+        """Detect PD (Prefill-Decode) disaggregation config from stage_configs.
+        Returns a dict with 'pd_pair' and 'bootstrap_addr', or None.
+        """
+        pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs)
+        if pd_pair is None:
+            return None
+        prefill_idx, decode_idx = pd_pair
+
+        # Extract bootstrap address from prefill stage engine_args
+        bootstrap_addr: str | None = None
+        try:
+            prefill_cfg = self.stage_configs[prefill_idx]
+            ea = getattr(prefill_cfg, "engine_args", None)
+            kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None
+            if kv_cfg is not None:
+                port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT
+                kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1"
+                bootstrap_addr = f"http://{kv_ip}:{port}"
+        except Exception as exc:
+            logger.warning("[AsyncOmniEngine] Could not extract PD bootstrap address: %s", exc)
+
+        logger.info(
+            "[AsyncOmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s",
+            prefill_idx,
+            decode_idx,
+            bootstrap_addr,
+        )
+        prefill_engine_id: str | None = None
+        try:
+            prefill_client = self.stage_clients[prefill_idx]
+            kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None)
+            prefill_engine_id = getattr(kv_cfg, "engine_id", None)
+        except Exception as exc:
+            logger.warning("[AsyncOmniEngine] Could not extract prefill engine_id: %s", exc)
+
+        return {
+            "pd_pair": (prefill_idx, decode_idx),
+            "bootstrap_addr": bootstrap_addr,
+            "prefill_engine_id": prefill_engine_id,
+        }
+
     @staticmethod
     def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list:
         """Create a default single-stage diffusion config from kwargs."""

vllm_omni/engine/orchestrator.py

@@ -42,6 +42,7 @@ def build_engine_core_request_from_tokens(
     params: SamplingParams | PoolingParams,
     arrival_time: float | None = None,
     model_config: ModelConfig | None = None,
+    mm_features: list | None = None,
 ) -> OmniEngineCoreRequest:
     """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt.
 
@@ -76,7 +77,7 @@ def build_engine_core_request_from_tokens(
     return OmniEngineCoreRequest(
         request_id=request_id,
         prompt_token_ids=prompt_token_ids,
-        mm_features=None,
+        mm_features=mm_features,
         sampling_params=sampling_params,
         pooling_params=pooling_params,
         arrival_time=arrival_time,
@@ -104,6 +105,8 @@ class OrchestratorRequestState:
 
     # Metrics: timestamp when request was submitted to each stage
     stage_submit_ts: dict[int, float] = field(default_factory=dict)
+    mm_processor_kwargs: dict | None = None
+    mm_features: list | None = None
 
 
 class Orchestrator:
@@ -123,6 +126,7 @@ class Orchestrator:
         stage_vllm_configs: list[Any],
         *,
         async_chunk: bool = False,
+        pd_config: dict[str, Any] | None = None,
     ) -> None:
         self.request_async_queue = request_async_queue
         self.output_async_queue = output_async_queue
@@ -135,6 +139,16 @@ class Orchestrator:
         self.output_processors: list[Any] = output_processors
         self.stage_vllm_configs: list[Any] = stage_vllm_configs
 
+        # PD disaggregation state
+        self._pd_pair: tuple[int, int] | None = None
+        self._pd_bootstrap_addr: str | None = None
+        self._pd_prefill_engine_id: str | None = None
+        self._pd_kv_params: dict[str, Any] = {}
+        if pd_config is not None:
+            self._pd_pair = pd_config.get("pd_pair")
+            self._pd_bootstrap_addr = pd_config.get("bootstrap_addr")
+            self._pd_prefill_engine_id = pd_config.get("prefill_engine_id")
+
         # Per-request state
         self.request_states: dict[str, OrchestratorRequestState] = {}
 
@@ -359,6 +373,23 @@ class Orchestrator:
                 }
             )
 
+        # PD disaggregation: extract KV transfer params from prefill stage output
+        if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]:
+            kv_params = getattr(output, "kv_transfer_params", None)
+            if kv_params is not None:
+                self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params)
+                logger.debug(
+                    "[Orchestrator][PD] stored kv_transfer_params for req=%s (keys=%s)",
+                    req_id,
+                    list(self._pd_kv_params[req_id].keys()),
+                )
+            else:
+                logger.warning(
+                    "[Orchestrator][PD] prefill stage output for req=%s has no kv_transfer_params; "
+                    "KV transfer may fail. Ensure apply_mooncake_connector_patch() was called.",
+                    req_id,
+                )
+
         if (
             finished
             and stage_id < req_state.final_stage_id
@@ -371,6 +402,8 @@ class Orchestrator:
                 await self._forward_to_next_stage(req_id, stage_id, output, req_state)
 
         if finished and stage_id == req_state.final_stage_id:
+            # PD: clean up any lingering KV params for this request
+            self._pd_kv_params.pop(req_id, None)
             self._cfg_tracker.cleanup_parent(req_id)
             self.request_states.pop(req_id, None)
 
@@ -418,6 +451,51 @@ class Orchestrator:
             else:
                 await self._forward_to_next_stage(req_id, stage_id, raw_output, req_state)
 
+    def _build_pd_decode_params(self, req_id: str, sp: Any) -> Any:
+        """Build decode-side sampling params with KV transfer params for PD routing.
+
+        Clones the sampling params and injects kv_transfer_params that tell the
+        decode engine where to pull the KV cache from (prefill engine's bootstrap addr).
+        """
+        sp = sp.clone()
+        if sp.extra_args is None:
+            sp.extra_args = {}
+
+        # Get KV params captured from the prefill output (must include remote_request_id).
+        kv_prefill_params = self._pd_kv_params.pop(req_id, None)
+        if not kv_prefill_params or "remote_request_id" not in kv_prefill_params:
+            raise RuntimeError(
+                f"[Orchestrator][PD] Missing prefill kv_transfer_params.remote_request_id for req={req_id}"
+            )
+
+        decode_kv_params: dict[str, Any] = {
+            "transfer_id": f"xfer-{req_id}",
+        }
+
+        if self._pd_bootstrap_addr:
+            decode_kv_params["remote_bootstrap_addr"] = self._pd_bootstrap_addr
+
+        if self._pd_prefill_engine_id:
+            decode_kv_params["remote_engine_id"] = self._pd_prefill_engine_id
+
+        # Overlay params from prefill side (includes remote_request_id set by monkey patch).
+        decode_kv_params.update(kv_prefill_params)
+
+        # Ensure these flags are set correctly after any overlay.
+        decode_kv_params["do_remote_prefill"] = True
+        decode_kv_params["do_remote_decode"] = False
+        if not decode_kv_params.get("transfer_id"):
+            decode_kv_params["transfer_id"] = f"xfer-{req_id}"
+
+        sp.extra_args["kv_transfer_params"] = decode_kv_params
+
+        logger.debug(
+            "[Orchestrator][PD] decode kv_transfer_params for req=%s: %s",
+            req_id,
+            decode_kv_params,
+        )
+        return sp
+
     def _build_stage_metrics(
         self,
         stage_id: int,
@@ -540,6 +618,52 @@ class Orchestrator:
             req_state.stage_submit_ts[next_stage_id] = _time.time()
             return
 
+        # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params
+        if self._pd_pair is not None and (stage_id, next_stage_id) == self._pd_pair:
+            # Save prefill stage outputs so thinker2talker can merge embeddings later
+            self.stage_clients[stage_id].set_engine_outputs([output])
+
+            params = self._build_pd_decode_params(req_id, params)
+
+            # Use the original user prompt for the decode stage (not processed embeddings)
+            original_prompt = req_state.prompt
+            raw_decode_inputs = [original_prompt] if not isinstance(original_prompt, list) else original_prompt
+
+            decode_inputs: list[dict[str, Any]] = []
+            for decode_input in raw_decode_inputs:
+                if isinstance(decode_input, dict):
+                    decode_inputs.append(decode_input)
+                    continue
+                prompt_token_ids = getattr(decode_input, "prompt_token_ids", None)
+                if prompt_token_ids is None:
+                    raise TypeError(
+                        "[Orchestrator][PD] decode input must be dict or have prompt_token_ids, "
+                        f"got {type(decode_input).__name__} for req={req_id}"
+                    )
+                decode_inputs.append({"prompt_token_ids": list(prompt_token_ids)})
+
+            for decode_input in decode_inputs:
+                request = build_engine_core_request_from_tokens(
+                    request_id=req_id,
+                    prompt=decode_input,
+                    params=params,
+                    model_config=self.stage_vllm_configs[next_stage_id].model_config,
+                    mm_features=req_state.mm_features,  # Pass mm_features for M-RoPE
+                )
+                request.external_req_id = request.request_id
+
+                self.output_processors[next_stage_id].add_request(
+                    request=request,
+                    prompt=None,
+                    parent_req=None,
+                    request_index=0,
+                    queue=None,
+                )
+                await next_client.add_request_async(request)
+
+            req_state.stage_submit_ts[next_stage_id] = _time.time()
+            return
+
         self.stage_clients[stage_id].set_engine_outputs([output])
 
         # Process inputs for next stage
@@ -558,11 +682,16 @@ class Orchestrator:
 
         # Build and submit requests for each input
         for next_input in next_inputs:
+            # Only AR thinker stages consume encoder mm_features; downstream
+            # (talker/code2wav/…) must not see them (avoids encoder-cache misses).
+            _ms = getattr(next_client, "model_stage", None)
+            _mm_features = req_state.mm_features if _ms == "thinker" else None
             request = build_engine_core_request_from_tokens(
                 request_id=req_id,
                 prompt=next_input,
                 params=params,
                 model_config=self.stage_vllm_configs[next_stage_id].model_config,
+                mm_features=_mm_features,
             )
 
             # TODO: Here we directly use the req id to assign.
@@ -644,6 +773,7 @@ class Orchestrator:
             prompt=original_prompt,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
+            mm_features=getattr(prompt, "mm_features", None),  # Save mm_features for PD
         )
         req_state.stage_submit_ts[stage_id] = _time.time()
         self.request_states[request_id] = req_state

vllm_omni/engine/output_processor.py

@@ -233,10 +233,9 @@ class OmniRequestState(RequestState):
         # Reuse base text/logprobs logic, then annotate with pooling_result.
         base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
         try:
+            if not hasattr(base_output, "multimodal_output"):
+                setattr(base_output, "multimodal_output", {})
             if self.mm_accumulated is not None:
-                # Attach accumulated multimodal dict on the completion output
-                if not hasattr(base_output, "multimodal_output"):
-                    setattr(base_output, "multimodal_output", {})
                 mm_out = getattr(base_output, "multimodal_output")
                 if isinstance(mm_out, dict):
                     for k, v in self.mm_accumulated.items():

vllm_omni/entrypoints/async_omni.py

@@ -78,7 +78,6 @@ class AsyncOmni(EngineClient, OmniBase):
         self.final_output_task: asyncio.Task | None = None
 
         self.config_path = self.engine.config_path
-        self.stage_configs = self.engine.stage_configs
         self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None)
         self.input_processor = self.engine.input_processor
 
@@ -209,6 +208,13 @@ class AsyncOmni(EngineClient, OmniBase):
             # Start final output dispatcher on the first call to generate()
             self._final_output_handler()
 
+            # Expand sampling params for PD disaggregation (user may provide N-1 params)
+            if (
+                sampling_params_list is not None
+                and isinstance(sampling_params_list, Sequence)
+                and not isinstance(sampling_params_list, (str, bytes))
+            ):
+                sampling_params_list = self._maybe_expand_sampling_params(list(sampling_params_list))
             sampling_params_list = self.resolve_sampling_params_list(sampling_params_list)
 
             # Track per-request metrics
@@ -228,20 +234,27 @@ class AsyncOmni(EngineClient, OmniBase):
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
 
+            # PD disaggregation: modify prefill-stage sampling params per request
+            req_sp_list = list(sampling_params_list)
+            pd_pair = self._get_pd_separation_pair()
+            if pd_pair is not None:
+                p_id = pd_pair[0]
+                req_sp_list[p_id] = self._prepare_prefill_sampling_params(request_id, req_sp_list[p_id])
+
             # Add request(s) to stage 0. For streaming inputs, submit
             # chunks incrementally through streaming_update.
             if isinstance(prompt, AsyncGenerator):
                 input_stream_task = await self._add_streaming_input_request(
                     request_id=request_id,
                     input_stream=prompt,
-                    sampling_params_list=sampling_params_list,
+                    sampling_params_list=req_sp_list,
                     final_stage_id=final_stage_id_for_e2e,
                 )
             else:
                 await self.engine.add_request_async(
                     request_id=request_id,
                     prompt=prompt,
-                    sampling_params_list=sampling_params_list,
+                    sampling_params_list=req_sp_list,
                     final_stage_id=final_stage_id_for_e2e,
                 )
             submit_ts = time.time()

vllm_omni/entrypoints/omni.py

@@ -66,6 +66,13 @@ class Omni(OmniBase):
         py_generator: bool = False,
         use_tqdm: bool | Callable[..., tqdm] = True,
     ) -> Generator[OmniRequestOutput, None, None] | list[OmniRequestOutput]:
+        # Expand sampling params for PD disaggregation (user may provide N-1 params)
+        if (
+            sampling_params_list is not None
+            and isinstance(sampling_params_list, Sequence)
+            and not isinstance(sampling_params_list, (str, bytes))
+        ):
+            sampling_params_list = self._maybe_expand_sampling_params(list(sampling_params_list))
         sampling_params_list = self.resolve_sampling_params_list(sampling_params_list)
         try:
             if py_generator:
@@ -125,10 +132,17 @@ class Omni(OmniBase):
                 req_state.metrics = metrics
                 self.request_states[req_id] = req_state
 
+                # PD disaggregation: modify stage-0 (prefill) sampling params per request
+                req_sp_list = list(sampling_params_list)
+                pd_pair = self._get_pd_separation_pair()
+                if pd_pair is not None:
+                    p_id = pd_pair[0]
+                    req_sp_list[p_id] = self._prepare_prefill_sampling_params(req_id, req_sp_list[p_id])
+
                 self.engine.add_request(
                     request_id=req_id,
                     prompt=prompt,
-                    sampling_params_list=sampling_params_list,
+                    sampling_params_list=req_sp_list,
                     final_stage_id=final_stage_id,
                 )
                 submit_ts = time.time()

vllm_omni/entrypoints/omni_base.py

@@ -14,6 +14,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
+from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -65,7 +66,7 @@ def omni_snapshot_download(model_id: str) -> str:
 OutputMessageHandleResult = tuple[Literal[True], None, None, None] | tuple[Literal[False], str, int, ClientRequestState]
 
 
-class OmniBase:
+class OmniBase(PDDisaggregationMixin):
     """Shared runtime foundation for AsyncOmni and Omni."""
 
     def __init__(
@@ -84,6 +85,7 @@ class OmniBase:
         if "log_requests" in kwargs:
             raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.")
         model = omni_snapshot_download(model)
+        self._name = self.__class__.__name__
         self.model = model
         self.log_stats = log_stats
         self.async_chunk = async_chunk
@@ -125,10 +127,18 @@ class OmniBase:
             model,
         )
 
+        # PD disaggregation state (detects if a prefill/decode stage pair is configured)
+        self._init_pd_state()
+
     @property
     def num_stages(self) -> int:
         return self.engine.num_stages
 
+    @property
+    def stage_configs(self) -> list:
+        """Expose engine stage configs for PD disaggregation detection and validation."""
+        return self.engine.stage_configs
+
     @property
     def is_running(self) -> bool:
         return self.engine.is_alive()

vllm_omni/entrypoints/pd_utils.py

@@ -23,9 +23,19 @@ logger = logging.getLogger(__name__)
 class PDDisaggregationMixin:
     """Mixin supplying PD disaggregation helpers to OmniBase."""
 
+    def _get_pd_separation_pair(self) -> tuple[int, int] | None:
+        """PD prefill/decode indices when ``_init_pd_state`` ran; else ``None``.
+
+        Partial test doubles may skip ``OmniBase.__init__``; treat missing state as
+        no PD disaggregation instead of raising ``AttributeError``.
+        """
+        return getattr(self, "_pd_separation_pair", None)
+
     def _init_pd_state(self) -> None:
         """Initialise PD disaggregation state."""
-        self._pd_separation_pair: tuple[int, int] | None = self._detect_pd_separation()
+        self._pd_separation_pair: tuple[int, int] | None = self.detect_pd_separation_from_stage_configs(
+            self.stage_configs
+        )
         self._pd_connector_info: dict[str, Any] | None = None
         self._pd_kv_params_by_req: dict[str, dict[str, Any]] = {}
         self._pd_kv_params_lock = threading.Lock()
@@ -40,11 +50,19 @@ class PDDisaggregationMixin:
                 d_id,
             )
 
-    def _detect_pd_separation(self) -> tuple[int, int] | None:
-        """Scan stage_list for a prefill/decode pair. Returns (p_id, d_id) or None."""
+    @staticmethod
+    def detect_pd_separation_from_stage_configs(stage_configs: list[Any]) -> tuple[int, int] | None:
+        """Scan stage configs for a prefill/decode pair.
+
+        Returns:
+            (prefill_idx, decode_idx) if one pair exists, None if not found.
+
+        Raises:
+            ValueError: if multiple candidate PD pairs are found.
+        """
         prefill_by_id: dict[int, int] = {}
         decode_indices: list[int] = []
-        for i, stage in enumerate(self.stage_list):
+        for i, stage in enumerate(stage_configs):
             if getattr(stage, "is_prefill_only", False):
                 prefill_by_id[i] = i
                 sid = getattr(stage, "stage_id", i)
@@ -55,7 +73,7 @@ class PDDisaggregationMixin:
 
         pd_pairs: list[tuple[int, int]] = []
         for j in decode_indices:
-            source_ids = getattr(self.stage_list[j], "engine_input_source", [])
+            source_ids = getattr(stage_configs[j], "engine_input_source", [])
             for src in source_ids:
                 if src in prefill_by_id:
                     pd_pairs.append((prefill_by_id[src], j))
@@ -107,10 +125,11 @@ class PDDisaggregationMixin:
 
     def _validate_pd_separation_config(self) -> None:
         """Validate PD stage configurations are consistent."""
-        assert self._pd_separation_pair is not None
-        p_id, d_id = self._pd_separation_pair
-        p_stage = self.stage_list[p_id]
-        d_stage = self.stage_list[d_id]
+        pair = self._get_pd_separation_pair()
+        assert pair is not None
+        p_id, d_id = pair
+        p_stage = self.stage_configs[p_id]
+        d_stage = self.stage_configs[d_id]
 
         def _get_kv_cfg(stage: "OmniStage") -> dict[str, Any]:
             ea = stage.engine_args
@@ -158,11 +177,12 @@ class PDDisaggregationMixin:
 
     def _get_pd_connector_info(self) -> dict[str, Any] | None:
         """Extract prefill engine KV connector info."""
-        if self._pd_separation_pair is None:
+        pair = self._get_pd_separation_pair()
+        if pair is None:
             return None
 
-        p_id, _ = self._pd_separation_pair
-        p_stage = self.stage_list[p_id]
+        p_id, _ = pair
+        p_stage = self.stage_configs[p_id]
 
         ea = p_stage.engine_args
         kv_cfg = getattr(ea, "kv_transfer_config", None)
@@ -241,18 +261,17 @@ class PDDisaggregationMixin:
 
     def _is_pd_routing(self, stage_id: int, next_stage_id: int) -> bool:
         """True when edge stage_id → next_stage_id is the prefill→decode boundary."""
-        return self._pd_separation_pair is not None and self._pd_separation_pair == (
-            stage_id,
-            next_stage_id,
-        )
+        pair = self._get_pd_separation_pair()
+        return pair is not None and pair == (stage_id, next_stage_id)
 
     def _maybe_expand_sampling_params(self, sampling_params_list: list) -> list:
         """Auto-duplicate thinker SP for decode stage when user provides N-1 params."""
-        if self._pd_separation_pair is None:
+        pair = self._get_pd_separation_pair()
+        if pair is None:
             return sampling_params_list
-        if len(sampling_params_list) != len(self.stage_list) - 1:
+        if len(sampling_params_list) != len(self.stage_configs) - 1:
             return sampling_params_list
-        p_id, d_id = self._pd_separation_pair
+        p_id, d_id = pair
         sp_list = list(sampling_params_list)
         sp_list.insert(d_id, sp_list[p_id])
         return sp_list

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -891,10 +891,11 @@ class Qwen3OmniMoeForConditionalGeneration(
         Returns:
             (input_ids, input_embeds) for talker
         """
+        target_len = thinker_result_ids.shape[-1]
         im_start_indexes = torch.cat(
             (
                 torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
-                torch.tensor([thinker_result_ids.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+                torch.tensor([target_len], device=input_ids.device, dtype=input_ids.dtype),
             ),
             dim=-1,
         )  # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
@@ -1029,8 +1030,35 @@ class Qwen3OmniMoeForConditionalGeneration(
         return last_talker_hidden, text_step, update_dict
 
     def _get_talker_user_parts(self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed):
+        clamped = min(
+            segment_end_index,
+            multimodal_mask.shape[0],
+            thinker_hidden.shape[0],
+            thinker_embed.shape[0],
+        )
+        if clamped < segment_end_index:
+            logger.warning(
+                "_get_talker_user_parts: segment_end_index %d clamped to %d "
+                "(embed=%d, hidden=%d, mask=%d). "
+                "This usually means _merge_pd_embeddings failed to merge "
+                "prefill embeddings – check PD prefill_mm keys.",
+                segment_end_index,
+                clamped,
+                thinker_embed.shape[0],
+                thinker_hidden.shape[0],
+                multimodal_mask.shape[0],
+            )
+        segment_end_index = clamped
+        seg_len = segment_end_index - im_start_index
+        if seg_len <= 0:
+            return torch.empty(
+                (0, self.config.talker_config.text_config.hidden_size),
+                device=thinker_hidden.device,
+                dtype=torch.bfloat16,
+            )
+
         user_talker_part = torch.empty(
-            (segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size),
+            (seg_len, self.config.talker_config.text_config.hidden_size),
             device=thinker_hidden.device,
             dtype=torch.bfloat16,
         )

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -3,6 +3,7 @@
 # Copyright 2025 The Qwen team.
 """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""
 
+import logging
 from typing import Any
 
 import torch
@@ -18,6 +19,12 @@ from vllm_omni.model_executor.stage_input_processors.tts_utils import (
     extract_speaker_from_request,
 )
 
+logger = logging.getLogger(__name__)
+
+# Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer
+_EMBED_LAYER_KEY = "0"
+_HIDDEN_LAYER_KEY = "24"
+
 
 def _compute_talker_prompt_ids_length(info, device: torch.device | str = "cuda") -> int:
     im_start_token_id = 151644
@@ -84,6 +91,95 @@ def _validate_stage_inputs(stage_list, engine_input_source):
     return stage.engine_outputs
 
 
+# =========================
+# PD disaggregation helpers
+# =========================
+
+
+def _get_prefill_stage(stage_list: list[Any], source_stage_id: int) -> Any | None:
+    if source_stage_id <= 0:
+        return None
+    source_stage = stage_list[source_stage_id]
+    if not getattr(source_stage, "is_decode_only", False):
+        return None
+    prev_stage = stage_list[source_stage_id - 1]
+    if getattr(prev_stage, "is_prefill_only", False) and prev_stage.engine_outputs is not None:
+        return prev_stage
+    return None
+
+
+def _merge_pd_embeddings(
+    decode_emb: torch.Tensor,
+    decode_hid: torch.Tensor,
+    prefill_mm: dict[str, Any],
+    device: torch.device,
+    expected_total: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Merge prefill prompt embeddings with decode generated embeddings.
+
+    In PD mode the prefill engine processes the prompt and the decode engine
+    generates tokens starting from position 1.  This function concatenates
+    them, removing the overlapping token(s):
+
+        merged = prefill[:P] + decode[overlap:]
+
+    where overlap = P + D - expected_total.
+    """
+    try:
+        p_emb = prefill_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float)
+        p_hid = prefill_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float)
+    except (KeyError, AttributeError, TypeError) as exc:
+        available_keys = list(prefill_mm.keys()) if isinstance(prefill_mm, dict) else type(prefill_mm).__name__
+        logger.error(
+            "_merge_pd_embeddings: failed to extract prefill embeddings (%s). "
+            "Expected keys %r and %r, got: %s. "
+            "Falling back to decode-only embeddings – talker user-segment will be degraded.",
+            exc,
+            _EMBED_LAYER_KEY,
+            _HIDDEN_LAYER_KEY,
+            available_keys,
+        )
+        return decode_emb, decode_hid
+
+    if p_emb.shape[0] == 0 or decode_emb.shape[0] == 0:
+        return decode_emb, decode_hid
+
+    raw_total = p_emb.shape[0] + decode_emb.shape[0]
+    overlap = max(0, raw_total - expected_total) if expected_total is not None else 0
+
+    merged_emb = torch.cat([p_emb, decode_emb[overlap:]], dim=0)
+    merged_hid = torch.cat([p_hid, decode_hid[overlap:]], dim=0)
+    return merged_emb, merged_hid
+
+
+def _get_prefill_multimodal_output(prefill_stage: Any, output_index: int) -> dict[str, Any] | None:
+    """Return multimodal_output dict from the PD prefill stage for a given batch index."""
+    try:
+        prefill_eos = prefill_stage.engine_outputs
+        prefill_eo = prefill_eos[min(output_index, len(prefill_eos) - 1)]
+        return prefill_eo.outputs[0].multimodal_output
+    except Exception:
+        return None
+
+
+def _resolve_tts_token_embedding(
+    key: str,
+    *,
+    thinker_mm: dict[str, Any],
+    prefill_mm: dict[str, Any] | None,
+    device: torch.device,
+) -> torch.Tensor | None:
+    """Return TTS BOS/EOS/PAD embedding tensors for the talker projection path.
+
+    Values are taken from the current thinker (decode) ``multimodal_output``; in
+    PD mode, missing keys may be filled from the paired prefill stage output.
+    """
+    val = thinker_mm.get(key)
+    if val is None and prefill_mm is not None:
+        val = prefill_mm.get(key)
+    return val.detach().to(device=device, dtype=torch.float) if val is not None else None
+
+
 # =========================
 # Thinker -> Talker
 # =========================
@@ -111,8 +207,8 @@ def thinker2talker_async_chunk(
         all_token_ids = _ensure_list(all_token_ids)
         prompt_token_ids = _ensure_list(prompt_token_ids)
         talker_additional_info = {
-            "thinker_prefill_embeddings": pooling_output.get("0").detach().cpu(),
-            "thinker_hidden_states": pooling_output.get("24").detach().cpu(),
+            "thinker_prefill_embeddings": pooling_output.get(_EMBED_LAYER_KEY).detach().cpu(),
+            "thinker_hidden_states": pooling_output.get(_HIDDEN_LAYER_KEY).detach().cpu(),
             "thinker_sequences": all_token_ids,
             "thinker_input_ids": prompt_token_ids,
             # Provide thinker-side TTS token embeddings for talker projection
@@ -161,7 +257,7 @@ def thinker2talker_async_chunk(
 
         if output_token_ids:
             talker_additional_info["override_keys"] = ["thinker_decode_embeddings", "thinker_output_token_ids"]
-            talker_additional_info["thinker_decode_embeddings"] = pooling_output.get("0").detach().cpu()
+            talker_additional_info["thinker_decode_embeddings"] = pooling_output.get(_EMBED_LAYER_KEY).detach().cpu()
             talker_additional_info["thinker_output_token_ids"] = output_token_ids
         else:
             # When prefilling a chunked thinker, thinker_hidden_states needs to be updated.
@@ -185,6 +281,9 @@ def thinker2talker(
     2. Split hidden states into: prompt embeddings + generated embeddings
     3. Package for talker with additional information
 
+    In PD disaggregation mode, merges prefill-stage prompt embeddings with
+    decode-stage generated embeddings before handing off to the talker.
+
     Args:
         stage_list: List of stage objects
         engine_input_source: Source stage IDs (typically [0] for thinker)
@@ -199,21 +298,49 @@ def thinker2talker(
 
     device = torch.device(current_platform.device_type)
 
+    # PD disaggregation: look up the preceding prefill stage (if any)
+    source_stage_id = engine_input_source[0]
+    prefill_stage = _get_prefill_stage(stage_list, source_stage_id)
+
     # Process each thinker output
     for i, thinker_output in enumerate(thinker_outputs):
         output = thinker_output.outputs[0]
+        thinker_mm = output.multimodal_output
+        # Full thinker embedding sequence for the talker: single thinker engine in the
+        # non-PD path; after optional merge with prefill-side tensors in PD mode.
+        thinker_emb = thinker_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float)
+        thinker_hid = thinker_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float)
+
+        prefill_mm: dict[str, Any] | None = None
+        if prefill_stage is not None:
+            prefill_mm = _get_prefill_multimodal_output(prefill_stage, i)
+
+        if prefill_mm is not None:
+            expected_total = len(thinker_output.prompt_token_ids) + len(output.token_ids)
+            try:
+                thinker_emb, thinker_hid = _merge_pd_embeddings(
+                    thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total
+                )
+            except Exception as exc:
+                logger.warning("[PD] Could not merge prefill embeddings: %s", exc)
 
         info = {
-            "thinker_prefill_embeddings": output.multimodal_output["0"].detach().to(device=device, dtype=torch.float),
-            "thinker_hidden_states": output.multimodal_output["24"].detach().to(device=device, dtype=torch.float),
+            "thinker_prefill_embeddings": thinker_emb,
+            "thinker_hidden_states": thinker_hid,
             "thinker_sequences": (
                 thinker_output.prompt_token_ids + output.token_ids
             ),  # the thinker_sequences is the whole ids
             "thinker_input_ids": thinker_output.prompt_token_ids,
             # Provide thinker-side TTS token embeddings for talker projection
-            "tts_bos_embed": output.multimodal_output["tts_bos_embed"].detach().to(device=device, dtype=torch.float),
-            "tts_eos_embed": output.multimodal_output["tts_eos_embed"].detach().to(device=device, dtype=torch.float),
-            "tts_pad_embed": output.multimodal_output["tts_pad_embed"].detach().to(device=device, dtype=torch.float),
+            "tts_bos_embed": _resolve_tts_token_embedding(
+                "tts_bos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+            ),
+            "tts_eos_embed": _resolve_tts_token_embedding(
+                "tts_eos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+            ),
+            "tts_pad_embed": _resolve_tts_token_embedding(
+                "tts_pad_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+            ),
         }
         speaker = extract_speaker_from_prompt(prompt, index=i)
         if speaker is not None: