[Core][KV-transfer] MoRIIO multi-node TP prefill→decode dispatch

examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py

@@ -308,13 +308,17 @@ 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"
         ]
         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

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

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,
         )