[Feature] Split #1303 Part 1: PD disaggregation scaffolding (#1863)

Signed-off-by: Jinheng Li <ahengljh@gmail.com>

docs/configuration/README.md

@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
+## Multi-Stage Recipes
+
+- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
+
 ## Optimization Features
 
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/configuration/pd_disaggregation.md

@@ -0,0 +1,176 @@
+# Prefill-Decode (PD) Disaggregation
+
+PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
+stages so prompt processing and token generation can run on different workers.
+
+This is documented as a stage-config recipe instead of a bundled YAML because the
+deployment-specific values usually change per environment:
+
+- GPU placement
+- `tensor_parallel_size`
+- connector backend and connector ports
+- connector IPs or bootstrap addresses
+
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
+and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
+changes below.
+
+## Requirements
+
+- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
+- A KV connector supported by vLLM, such as `MooncakeConnector`
+- Matching `tensor_parallel_size` on the prefill and decode thinker stages
+
+## 1. Split the thinker into prefill and decode stages
+
+Replace the original thinker stage with two stages:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_prefill_only: true
+    runtime:
+      devices: "0"
+      max_batch_size: 16
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_producer"
+        kv_rank: 0
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25201
+    final_output: false
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm
+    is_decode_only: true
+    runtime:
+      devices: "1"
+      max_batch_size: 64
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_consumer"
+        kv_rank: 1
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25202
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+```
+
+Notes:
+
+- `is_prefill_only: true` marks the thinker stage that only saves KV.
+- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
+- `kv_transfer_config` is required on both stages.
+- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
+  prefill side only processes the prompt and exports KV.
+
+## 2. Shift the downstream stages by one index
+
+After inserting the extra thinker stage, renumber the remaining stages:
+
+```yaml
+  - stage_id: 2
+    runtime:
+      devices: "2"
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+
+  - stage_id: 3
+    runtime:
+      devices: "2"
+      max_batch_size: 1
+    engine_input_source: [2]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+```
+
+Compared with the default Qwen3-Omni config:
+
+- the talker becomes stage `2` instead of stage `1`
+- the code2wav stage becomes stage `3` instead of stage `2`
+- the talker now reads from decode stage `1`
+
+## 3. Add runtime edges for the four-stage pipeline
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1
+    - from: 2
+      to: 3
+      window_size: -1
+```
+
+## 4. Launch with your custom config
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --stage-configs-path /path/to/qwen3_omni_pd.yaml
+```
+
+## Operational Notes
+
+- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
+  pair. Keep prefill and decode at the same `tensor_parallel_size`.
+- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
+  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
+- Choose connector ports and addresses that match your deployment. The values
+  shown above are examples only.

vllm_omni/distributed/kv_transfer/__init__.py

@@ -0,0 +1,13 @@
+"""Patched KV transfer connectors for PD disaggregation.
+
+This package provides monkey-patched versions of vLLM's native KV transfer
+connectors (e.g. MooncakeConnector) that fix the request-ID mismatch problem
+in prefill-decode disaggregation.
+
+vLLM's ``InputProcessor.assign_request_id()`` appends a random 8-char suffix
+to each request ID internally.  The prefill engine stores KV under its own
+suffix, but the decode engine generates a *different* suffix — so it can never
+find the KV data.  The patched connector threads the prefill engine's internal
+``remote_request_id`` through ``kv_transfer_params`` so the decode side can
+reference the correct KV entry.
+"""

vllm_omni/distributed/kv_transfer/monkey_patch.py

@@ -0,0 +1,197 @@
+"""Monkey-patch vLLM's MooncakeConnector to fix request-ID mismatch in PD disaggregation.
+
+vLLM's InputProcessor appends a random suffix to each request ID. The prefill
+engine stores KV under its suffix, but the decode engine generates a different
+suffix. This patch threads ``remote_request_id`` through ``kv_transfer_params``
+so the decode side references the correct KV entry.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_patched: bool = False
+
+
+@dataclass
+class PatchedRecvReqMeta:
+    """Receive-request metadata carrying the prefill engine's request ID."""
+
+    request_id: str
+    remote_request_id: str
+    local_block_ids: list[int]
+    kv_transfer_params: dict[str, Any]
+
+
+def _import_mooncake_module():
+    """Import MooncakeConnector module, supporting both vLLM >=0.16 and older."""
+    try:
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector
+
+        return mooncake_connector
+    except ImportError:
+        pass
+    try:
+        from vllm.distributed.kv_transfer.kv_connector.v1 import mooncake_connector
+
+        return mooncake_connector
+    except ImportError:
+        return None
+
+
+def _create_patched_mooncake_connector():
+    """Return a subclass of MooncakeConnector with remote_request_id support."""
+    try:
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
+            MooncakeConnector as _OriginalMooncakeConnector,
+        )
+    except (ImportError, AttributeError):
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector import (
+            MooncakeConnector as _OriginalMooncakeConnector,
+        )
+
+    class PatchedMooncakeConnector(_OriginalMooncakeConnector):
+        """Fixes request-ID mismatch in PD disaggregation by injecting
+        remote_request_id on the prefill side and using it for KV lookup
+        on the decode side.
+        """
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.remote_to_local_req: dict[str, str] = {}
+            logger.info("[PatchedMooncakeConnector] Initialized")
+
+        def request_finished(
+            self,
+            request: Any,
+            block_ids: list[int],
+        ) -> tuple[bool, dict[str, Any] | None]:
+            result = super().request_finished(request, block_ids)
+
+            if isinstance(result, tuple) and len(result) == 2:
+                delay_free, kv_params = result
+            else:
+                delay_free, kv_params = False, result
+
+            # Normalise _reqs_need_send values
+            req_id = getattr(request, "request_id", None)
+            if req_id and hasattr(self, "_reqs_need_send"):
+                entry = self._reqs_need_send.get(req_id)
+                if isinstance(entry, tuple) and len(entry) == 2:
+                    self._reqs_need_send[req_id] = entry[1]
+
+            # Inject remote_request_id into kv_transfer_params
+            if kv_params is not None and isinstance(kv_params, dict):
+                kv_params["remote_request_id"] = req_id or "NOT_SET"
+                if hasattr(self, "side_channel_host"):
+                    kv_params.setdefault("remote_host", self.side_channel_host)
+                if hasattr(self, "side_channel_port"):
+                    kv_params.setdefault("remote_port", self.side_channel_port)
+
+            return delay_free, kv_params
+
+        def add_new_req(
+            self,
+            request_id: str,
+            local_block_ids: list[int],
+            kv_transfer_params: dict[str, Any] | None = None,
+            **kwargs: Any,
+        ) -> None:
+            super().add_new_req(request_id, local_block_ids, kv_transfer_params, **kwargs)
+
+            kv_transfer_params = kv_transfer_params or {}
+            load_remote_cache = kv_transfer_params.get(
+                "do_remote_prefill",
+                kv_transfer_params.get("load_remote_cache", False),
+            )
+
+            if load_remote_cache:
+                remote_request_id = kv_transfer_params.get("remote_request_id", request_id)
+                meta = PatchedRecvReqMeta(
+                    request_id=request_id,
+                    remote_request_id=remote_request_id,
+                    local_block_ids=local_block_ids,
+                    kv_transfer_params=kv_transfer_params,
+                )
+                if not hasattr(self, "_reqs_need_recv"):
+                    self._reqs_need_recv = {}
+                self._reqs_need_recv[request_id] = meta
+
+        def group_kv_pull(self, metadata: Any | None = None) -> None:
+            """Use remote_request_id as ZMQ lookup key via save-patch-restore."""
+            if not hasattr(self, "_reqs_need_recv") or not self._reqs_need_recv:
+                return
+
+            original_recv = self._reqs_need_recv.copy()
+            patched_recv: dict[str, Any] = {}
+
+            for local_id, meta in original_recv.items():
+                if isinstance(meta, PatchedRecvReqMeta):
+                    remote_id = meta.remote_request_id
+                    self.remote_to_local_req[remote_id] = local_id
+                    patched_meta = type(meta)(
+                        request_id=remote_id,
+                        remote_request_id=remote_id,
+                        local_block_ids=meta.local_block_ids,
+                        kv_transfer_params=meta.kv_transfer_params,
+                    )
+                    patched_recv[remote_id] = patched_meta
+                else:
+                    patched_recv[local_id] = meta
+
+            self._reqs_need_recv = patched_recv
+            super().group_kv_pull(metadata)
+
+            # Restore unconsumed entries to original local keys
+            for remote_id, local_id in list(self.remote_to_local_req.items()):
+                if remote_id in self._reqs_need_recv:
+                    entry = self._reqs_need_recv.pop(remote_id)
+                    self._reqs_need_recv[local_id] = original_recv.get(local_id, entry)
+
+        def receive_kv(self, path: Any = None, req_blocks: Any = None) -> Any:
+            result = super().receive_kv(path, req_blocks)
+
+            if self.remote_to_local_req:
+                completed = [
+                    rid
+                    for rid, lid in self.remote_to_local_req.items()
+                    if not hasattr(self, "_reqs_need_recv") or lid not in self._reqs_need_recv
+                ]
+                for remote_id in completed:
+                    self.remote_to_local_req.pop(remote_id, None)
+
+            return result
+
+    PatchedMooncakeConnector.__qualname__ = _OriginalMooncakeConnector.__qualname__
+
+    return PatchedMooncakeConnector
+
+
+def apply_mooncake_connector_patch() -> bool:
+    """Replace vLLM's MooncakeConnector with the patched version."""
+    global _patched
+    if _patched:
+        return True
+
+    _mc_module = _import_mooncake_module()
+    if _mc_module is None:
+        logger.warning("[monkey_patch] Cannot import MooncakeConnector — patch NOT applied.")
+        return False
+
+    _OriginalClass = _mc_module.MooncakeConnector
+
+    PatchedClass = _create_patched_mooncake_connector()
+
+    _mc_module.MooncakeConnector = PatchedClass
+    for _, module in sys.modules.items():
+        if hasattr(module, "MooncakeConnector") and module.MooncakeConnector is _OriginalClass:
+            module.MooncakeConnector = PatchedClass
+
+    _patched = True
+    logger.info("[monkey_patch] MooncakeConnector patch applied")
+    return True

vllm_omni/entrypoints/pd_utils.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PD (Prefill-Decode) disaggregation helpers.
+
+Mixin for OmniBase — keeps omni.py focused on orchestration.
+"""
+
+import logging
+import threading
+from dataclasses import asdict, is_dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from vllm import SamplingParams
+
+    from vllm_omni.entrypoints.omni_stage import OmniStage
+
+_DEFAULT_MOONCAKE_BOOTSTRAP_PORT = 25201
+
+logger = logging.getLogger(__name__)
+
+
+class PDDisaggregationMixin:
+    """Mixin supplying PD disaggregation helpers to OmniBase."""
+
+    def _init_pd_state(self) -> None:
+        """Initialise PD disaggregation state."""
+        self._pd_separation_pair: tuple[int, int] | None = self._detect_pd_separation()
+        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()
+        if self._pd_separation_pair is not None:
+            self._validate_pd_separation_config()
+            self._pd_connector_info = self._get_pd_connector_info()
+            p_id, d_id = self._pd_separation_pair
+            logger.info(
+                "[%s] PD disaggregation detected: prefill=stage-%d, decode=stage-%d",
+                self._name,
+                p_id,
+                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."""
+        prefill_by_id: dict[int, int] = {}
+        decode_indices: list[int] = []
+        for i, stage in enumerate(self.stage_list):
+            if getattr(stage, "is_prefill_only", False):
+                prefill_by_id[i] = i
+                sid = getattr(stage, "stage_id", i)
+                if sid != i:
+                    prefill_by_id[sid] = i
+            if getattr(stage, "is_decode_only", False):
+                decode_indices.append(i)
+
+        pd_pairs: list[tuple[int, int]] = []
+        for j in decode_indices:
+            source_ids = getattr(self.stage_list[j], "engine_input_source", [])
+            for src in source_ids:
+                if src in prefill_by_id:
+                    pd_pairs.append((prefill_by_id[src], j))
+                    break
+
+        if len(pd_pairs) > 1:
+            raise ValueError(
+                f"Multiple PD pairs detected ({pd_pairs}); only a single PD pair per pipeline is supported"
+            )
+        return pd_pairs[0] if pd_pairs else None
+
+    @staticmethod
+    def _to_dict(obj: Any, default: Any = None) -> dict[str, Any] | None:
+        """Convert obj to a plain dict, trying several strategies."""
+        if obj is None:
+            return default
+        if isinstance(obj, dict):
+            return obj
+        if is_dataclass(obj):
+            try:
+                return asdict(obj)
+            except Exception:
+                return default
+        for attr in ("model_dump", "dict"):
+            if hasattr(obj, attr):
+                try:
+                    return getattr(obj, attr)()
+                except Exception:
+                    pass
+        if hasattr(obj, "items"):
+            try:
+                return dict(obj)
+            except Exception:
+                pass
+        try:
+            return dict(obj)
+        except Exception:
+            try:
+                return vars(obj)
+            except Exception:
+                logger.debug("Unable to convert %s to dict", type(obj))
+                return default
+
+    def _kv_cfg_to_dict(self, kv_cfg: Any) -> dict[str, Any]:
+        return self._to_dict(kv_cfg, default={}) or {}
+
+    def _normalize_kv_transfer_params(self, kv_params: Any) -> dict[str, Any] | None:
+        return self._to_dict(kv_params)
+
+    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]
+
+        def _get_kv_cfg(stage: "OmniStage") -> dict[str, Any]:
+            ea = stage.engine_args
+            cfg = getattr(ea, "kv_transfer_config", None)
+            if cfg is None:
+                cfg = ea.get("kv_transfer_config", None) if hasattr(ea, "get") else None
+            if cfg is None:
+                raise ValueError(f"Stage-{stage.stage_id} is marked for PD but has no 'kv_transfer_config'")
+            cfg_dict = self._kv_cfg_to_dict(cfg)
+            if not cfg_dict:
+                raise ValueError(f"Stage-{stage.stage_id} kv_transfer_config could not be parsed")
+            return cfg_dict
+
+        p_cfg = _get_kv_cfg(p_stage)
+        d_cfg = _get_kv_cfg(d_stage)
+
+        p_role = p_cfg.get("kv_role")
+        d_role = d_cfg.get("kv_role")
+        if p_role not in ("kv_producer", "kv_both"):
+            raise ValueError(f"Prefill stage-{p_id} kv_role must be 'kv_producer' or 'kv_both', got '{p_role}'")
+        if d_role not in ("kv_consumer", "kv_both"):
+            raise ValueError(f"Decode stage-{d_id} kv_role must be 'kv_consumer' or 'kv_both', got '{d_role}'")
+
+        d_sources = list(getattr(d_stage, "engine_input_source", []) or [])
+        if p_id not in d_sources and p_stage.stage_id not in d_sources:
+            raise ValueError(f"Decode stage-{d_id} must list prefill stage-{p_id} in engine_input_source")
+
+        p_conn = p_cfg.get("kv_connector")
+        d_conn = d_cfg.get("kv_connector")
+        if p_conn != d_conn:
+            raise ValueError(f"PD connector mismatch: prefill uses '{p_conn}', decode uses '{d_conn}'")
+        if not p_conn:
+            raise ValueError("PD requires kv_connector in kv_transfer_config")
+
+        for key in ("kv_buffer_device", "kv_buffer_size"):
+            p_val = p_cfg.get(key)
+            d_val = d_cfg.get(key)
+            if p_val is not None and d_val is not None and p_val != d_val:
+                raise ValueError(f"PD {key} mismatch: prefill='{p_val}', decode='{d_val}'")
+
+        p_tp = getattr(getattr(p_stage, "engine_args", None), "tensor_parallel_size", 1)
+        d_tp = getattr(getattr(d_stage, "engine_args", None), "tensor_parallel_size", 1)
+        if p_tp != d_tp:
+            raise ValueError(f"PD stages must have matching tensor_parallel_size: prefill={p_tp}, decode={d_tp}")
+
+    def _get_pd_connector_info(self) -> dict[str, Any] | None:
+        """Extract prefill engine KV connector info."""
+        if self._pd_separation_pair is None:
+            return None
+
+        p_id, _ = self._pd_separation_pair
+        p_stage = self.stage_list[p_id]
+
+        ea = p_stage.engine_args
+        kv_cfg = getattr(ea, "kv_transfer_config", None)
+        if kv_cfg is None and hasattr(ea, "get"):
+            kv_cfg = ea.get("kv_transfer_config")
+        if kv_cfg is None:
+            return None
+
+        kv_cfg_dict = self._kv_cfg_to_dict(kv_cfg)
+        if not kv_cfg_dict:
+            return None
+
+        kv_connector = str(kv_cfg_dict.get("kv_connector", "") or "")
+        extra_cfg = kv_cfg_dict.get("kv_connector_extra_config", {}) or {}
+        if not isinstance(extra_cfg, dict):
+            extra_cfg = self._kv_cfg_to_dict(extra_cfg)
+
+        info: dict[str, Any] = {}
+
+        if "mooncake" in kv_connector.lower():
+            bootstrap_port = extra_cfg.get("mooncake_bootstrap_port", _DEFAULT_MOONCAKE_BOOTSTRAP_PORT)
+            kv_ip = kv_cfg_dict.get("kv_ip") or "127.0.0.1"
+            info["prefill_bootstrap_addr"] = f"{kv_ip}:{bootstrap_port}"
+
+        return info
+
+    def _prepare_prefill_sampling_params(self, req_id: str, sp: "SamplingParams") -> "SamplingParams":
+        sp = sp.clone()
+        sp.max_tokens = 1
+        if hasattr(sp, "min_tokens"):
+            try:
+                sp.min_tokens = 1
+            except Exception:
+                pass
+        # MooncakeConnector cancels KV transfer unless finish_reason='length'
+        sp.stop = []
+        sp.stop_token_ids = []
+        sp.include_stop_str_in_output = False
+        if sp.extra_args is None:
+            sp.extra_args = {}
+        kv_params = self._normalize_kv_transfer_params(sp.extra_args.get("kv_transfer_params"))
+        merged: dict[str, Any] = {}
+        if kv_params:
+            merged.update(kv_params)
+        merged.update(
+            {
+                "do_remote_decode": True,
+                "do_remote_prefill": False,
+                "transfer_id": f"xfer-{req_id}",
+            }
+        )
+        sp.extra_args["kv_transfer_params"] = merged
+        logger.debug("[PD] prefill SP: req=%s kv_transfer_params=%s", req_id, merged)
+        return sp
+
+    def _pop_pd_kv_params(self, req_id: str, fallback: Any | None = None) -> dict[str, Any] | None:
+        kv_params = self._normalize_kv_transfer_params(fallback)
+        with self._pd_kv_params_lock:
+            stored = self._pd_kv_params_by_req.pop(req_id, None)
+        if kv_params is None:
+            kv_params = stored
+        return kv_params
+
+    def _drop_pd_kv_params(self, req_id: str) -> None:
+        with self._pd_kv_params_lock:
+            self._pd_kv_params_by_req.pop(req_id, None)
+
+    def _extract_kv_transfer_params(self, engine_outputs: Any) -> dict[str, Any] | None:
+        """Extract kv_transfer_params from engine outputs (if available)."""
+        outputs = engine_outputs if isinstance(engine_outputs, list) else [engine_outputs]
+        for output in outputs:
+            kv_params = getattr(output, "kv_transfer_params", None)
+            if kv_params is not None:
+                return self._normalize_kv_transfer_params(kv_params)
+        return None
+
+    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,
+        )
+
+    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:
+            return sampling_params_list
+        if len(sampling_params_list) != len(self.stage_list) - 1:
+            return sampling_params_list
+        p_id, d_id = self._pd_separation_pair
+        sp_list = list(sampling_params_list)
+        sp_list.insert(d_id, sp_list[p_id])
+        return sp_list
+
+    def _build_decode_kv_params(
+        self,
+        req_id: str,
+        sp_next: "SamplingParams",
+        prefill_kv_params: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Build decode-side kv_transfer_params dict."""
+        decode_kv_params: dict[str, Any] = {
+            "do_remote_decode": False,
+            "do_remote_prefill": True,
+            "transfer_id": f"xfer-{req_id}",
+        }
+
+        if sp_next.extra_args:
+            existing = self._normalize_kv_transfer_params(sp_next.extra_args.get("kv_transfer_params"))
+            if existing:
+                decode_kv_params.update(existing)
+
+        if self._pd_connector_info:
+            baddr = self._pd_connector_info.get("prefill_bootstrap_addr")
+            if baddr is not None and "remote_bootstrap_addr" not in decode_kv_params:
+                decode_kv_params["remote_bootstrap_addr"] = baddr
+
+        if prefill_kv_params:
+            decode_kv_params.update(prefill_kv_params)
+
+        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}"
+
+        return decode_kv_params
+
+    def _prepare_pd_decode_routing(
+        self,
+        req_id: str,
+        prompt: Any,
+        sp_template: "SamplingParams",
+        prefill_kv_fallback: Any | None = None,
+        engine_outputs: Any | None = None,
+    ) -> tuple[list, "SamplingParams"]:
+        """Prepare next_inputs and sampling params for prefill→decode routing.
+
+        Used by both sync (omni.py) and async (async_omni.py) paths.
+
+        Returns:
+            (next_inputs, sp_next) ready to submit to the decode stage.
+        """
+        next_inputs = [prompt] if not isinstance(prompt, list) else prompt
+
+        sp_next = sp_template.clone()
+        if sp_next.extra_args is None:
+            sp_next.extra_args = {}
+
+        if engine_outputs is not None:
+            prefill_kv = self._extract_kv_transfer_params(engine_outputs)
+        else:
+            prefill_kv = self._pop_pd_kv_params(req_id, prefill_kv_fallback)
+
+        decode_kv_params = self._build_decode_kv_params(req_id, sp_next, prefill_kv)
+        sp_next.extra_args["kv_transfer_params"] = decode_kv_params
+
+        logger.debug(
+            "[PD] routing req=%s, remote_request_id=%s",
+            req_id,
+            decode_kv_params.get("remote_request_id", "NOT SET"),
+        )
+        if "remote_request_id" not in decode_kv_params:
+            logger.warning(
+                "[PD] remote_request_id NOT SET for req %s. Apply mooncake_connector.py patch to fix.",
+                req_id,
+            )
+
+        return next_inputs, sp_next