diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index aceb7a9b81c3..a0aba8bb868c 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -308,6 +308,9 @@ async def handle_request(api: str, request: Request): "remote_block_ids" ] req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] + req_data["kv_transfer_params"]["remote_hosts"] = prefill_kv.get( + "remote_hosts" + ) req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[ "dp_size" @@ -315,6 +318,7 @@ async def handle_request(api: str, request: Request): req_data["kv_transfer_params"]["remote_tp_size"] = prefill_instance_endpoint[ "tp_size" ] + req_data["kv_transfer_params"]["tp_size"] = prefill_instance_endpoint["tp_size"] if selected_prefill_dp_rank is not None: req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index b843c5b59308..6e25b2fa6462 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -320,6 +320,11 @@ class ReqMeta: remote_engine_id: str tp_size: int remote_dp_size: int + # Ordered list of all prefill-instance host IPs for multi-node TP. + # Each decode worker picks remote_hosts[tp_rank // ranks_per_node] as its + # actual peer host for handshake + post-transfer notify. None or len<=1 + # falls back to single-host behaviour (remote_host). + remote_hosts: list[str] | None = None class MoRIIOConnectorMetadata(KVConnectorMetadata): @@ -362,8 +367,15 @@ def add_new_req( remote_port=remote_handshake_port, remote_handshake_port=remote_handshake_port, remote_notify_port=remote_notify_port, - tp_size=kv_transfer_params.get("tp_size", 1), + # Defense-in-depth: callers (e.g. moriio_toy_proxy_server in + # READ-mode multi-node TP) may forward `remote_tp_size` only. + # Read `tp_size` first, fall back to `remote_tp_size`, default 1. + tp_size=kv_transfer_params.get( + "tp_size", + kv_transfer_params.get("remote_tp_size", 1), + ), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), + remote_hosts=kv_transfer_params.get("remote_hosts"), ) if write_mode: self.reqs_to_save[request_id] = _req diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index bf0575dd2ac8..8104ab50ed0c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging import math +import os import queue import threading import time @@ -252,10 +253,26 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.engine_id: EngineId = engine_id self.mode = get_moriio_mode() self.host_ip = get_ip() + # Multi-node TP: VLLM_MORIIO_NODE_HOSTS holds the ordered list of host + # IPs in this engine's TP group (rank 0 first). Surfaced to the peer + # side via request_finished's kv_transfer_params so the consumer + # workers can dial the correct producer host for their tp_rank. For + # single-node TP, falls back to [host_ip]. + node_hosts_env = os.environ.get("VLLM_MORIIO_NODE_HOSTS", "").strip() + if node_hosts_env: + self.node_hosts: list[str] = [ + h.strip() for h in node_hosts_env.split(",") if h.strip() + ] + else: + self.node_hosts = [self.host_ip] self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[ "handshake_port" ] - logger.info("Initializing MoRIIO Scheduler engine_id = %s", engine_id) + logger.info( + "Initializing MoRIIO Scheduler engine_id = %s node_hosts = %s", + engine_id, + self.node_hosts, + ) self.side_notify_port = self.kv_transfer_config.kv_connector_extra_config[ "notify_port" @@ -602,6 +619,10 @@ def request_finished( remote_engine_id=self.engine_id, tp_size=self.vllm_config.parallel_config.tensor_parallel_size, transfer_id=params["transfer_id"], + # Multi-node TP: list of all prefill-instance host IPs in this + # engine's TP group (rank 0 first). Decode workers use this to + # pick the correct producer host per their tp_rank. + remote_hosts=self.node_hosts, ) @@ -1072,6 +1093,25 @@ def _moriio_handshake( return {remote_agent_name} + def _pick_remote_host(self, meta: ReqMeta) -> str: + """Resolve the per-worker peer host for multi-node TP prefill-decode. + + For TP=16 sym 1P1D (2 prefill nodes + 2 decode nodes), the prefill KV + cache for ranks [0..7] lives on prefill-head and for [8..15] lives on + prefill-worker. Each decode worker has a fixed self.tp_rank and must + dial the prefill node that owns its half of the KV. + + meta.remote_hosts is the ordered prefill-side host list (rank 0 + first). If unset or length 1, falls back to meta.remote_host + (single-host behaviour, preserves TP=8 and monolithic correctness). + """ + if meta.remote_hosts and len(meta.remote_hosts) > 1: + ranks_per_node = max(1, int(meta.tp_size) // len(meta.remote_hosts)) + node_idx = self.tp_rank // ranks_per_node + if 0 <= node_idx < len(meta.remote_hosts): + return meta.remote_hosts[node_idx] + return meta.remote_host + def _background_moriio_handshake( self, req_id: ReqId, remote_engine_id: EngineId, meta: ReqMeta ): @@ -1080,7 +1120,9 @@ def _background_moriio_handshake( if remote_engine_id is not None: fut = self._handshake_futures.get(remote_engine_id) if fut is None: - host = meta.remote_host + # Multi-node TP: pick the producer host for this worker's tp_rank. + # For single-node TP this returns meta.remote_host unchanged. + host = self._pick_remote_host(meta) port = int(meta.remote_handshake_port) tp_size = int(meta.tp_size) remote_dp_size = int(meta.remote_dp_size) @@ -1394,12 +1436,19 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): meta.remote_engine_id, req_id, ) + # Multi-node TP: remote_host here is used by _read_blocks to record the + # post-transfer notify callback address (so prefill can free its KV + # blocks). For multi-node prefill the notify must go to the prefill + # node that actually owns this worker's KV slice, not the prefill + # head. _pick_remote_host returns meta.remote_host unchanged for + # single-host setups, preserving TP=8 / monolithic behaviour. + actual_remote_host = self._pick_remote_host(meta) self._read_blocks( request_id=req_id, dst_engine_id=meta.remote_engine_id, local_block_ids=meta.local_block_ids, remote_block_ids=meta.remote_block_ids, - remote_host=meta.remote_host, + remote_host=actual_remote_host, remote_notify_port=meta.remote_notify_port, )