From 88caaf182b2239d52ce4363987b9350ae834cd56 Mon Sep 17 00:00:00 2001 From: Jinheng Date: Mon, 16 Mar 2026 14:05:19 +0800 Subject: [PATCH] [Feature] Split #1303 Part 1: PD disaggregation scaffolding (#1863) Signed-off-by: Jinheng Li --- docs/configuration/README.md | 4 + docs/configuration/pd_disaggregation.md | 176 +++++++++ vllm_omni/distributed/kv_transfer/__init__.py | 13 + .../distributed/kv_transfer/monkey_patch.py | 197 +++++++++++ vllm_omni/entrypoints/pd_utils.py | 333 ++++++++++++++++++ 5 files changed, 723 insertions(+) create mode 100644 docs/configuration/pd_disaggregation.md create mode 100644 vllm_omni/distributed/kv_transfer/__init__.py create mode 100644 vllm_omni/distributed/kv_transfer/monkey_patch.py create mode 100644 vllm_omni/entrypoints/pd_utils.py diff --git a/docs/configuration/README.md b/docs/configuration/README.md index e2073a96a..7e32806ea 100644 --- a/docs/configuration/README.md +++ b/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 diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md new file mode 100644 index 000000000..5cd5d0e7f --- /dev/null +++ b/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. diff --git a/vllm_omni/distributed/kv_transfer/__init__.py b/vllm_omni/distributed/kv_transfer/__init__.py new file mode 100644 index 000000000..f8914a428 --- /dev/null +++ b/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. +""" diff --git a/vllm_omni/distributed/kv_transfer/monkey_patch.py b/vllm_omni/distributed/kv_transfer/monkey_patch.py new file mode 100644 index 000000000..455ad8065 --- /dev/null +++ b/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 diff --git a/vllm_omni/entrypoints/pd_utils.py b/vllm_omni/entrypoints/pd_utils.py new file mode 100644 index 000000000..0e3d65f55 --- /dev/null +++ b/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