karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Core][KV-transfer] MoRIIO: heterogeneous TP<->DP prefill/decode read routing (#46116)

Browse Source 
Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io>
This commit is contained in:
limeward
2026-07-27 02:10:00 +00:00
committed by GitHub
parent 50aa830482
commit ffc4f08c8e
3 changed files with 467 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_moriio_routing_fairness.py
+210
View File
@@ -0,0 +1,210 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Hardware-fair KV-read routing across MoRIIO heterogeneous P/D configs.
Dependency-light (no GPU / ROCm / mori): builds bare ``MoRIIOConnectorWorker``
instances via ``object.__new__`` and drives the REAL read-source routing
decision (``_resolve_read_source`` / ``_next_flex_tp_rank``) directly -- no
routing logic is re-implemented here, so a future change to those functions is
what these tests exercise.
Scope -- the CONNECTOR (decode-side) read routing. The connector never selects
the prefill instance; the proxy does and hands the connector a ``remote_host`` +
``remote_dp_rank`` per request. The connector's only fairness lever is WHICH
prefill (dp, tp) rank each read targets, and the RFC's deployments collapse to
three real connector behaviours:
symmetric TP (TP prefill + TP decode) -> decode tp_k reads prefill tp_k
flexible (TP prefill + DP decode) -> round-robin over prefill tp0..N-1
owner DP (DP prefill, any decode) -> read the owner rank, tp0
These assume the proxy delivers a FAIR request stream (each prefill instance +
owner dp-rank an equal share) and verify the connector never re-introduces a
bottleneck. That proxy contract is checked separately against the real
``flat_interleaved_dp_route`` in the toy-proxy tests (PR #46115).
A node runs one of two modes (8 GPUs each):
* TP8 -> (dp_size, tp_size) = (1, 8); MLA latent KV REPLICATED on all ranks.
* DP8EP -> (dp_size, tp_size) = (8, 1); KV PARTITIONED, one owner rank/request.
"""
from collections import Counter
from dataclasses import dataclass
from unittest.mock import MagicMock
import pytest
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
ReqMeta,
get_port_offset,
)
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
MoRIIOConnectorWorker,
)
MODE_DIMS = {"TP8": (1, 8), "DP8EP": (8, 1)} # (dp_size, tp_size), 8 GPUs/node
@dataclass(frozen=True)
class PDConfig:
name: str
p_mode: str
d_mode: str
@property
def p_dp(self) -> int:
return MODE_DIMS[self.p_mode][0]
@property
def p_tp(self) -> int:
return MODE_DIMS[self.p_mode][1]
@property
def d_dp(self) -> int:
return MODE_DIMS[self.d_mode][0]
@property
def d_tp(self) -> int:
return MODE_DIMS[self.d_mode][1]
@property
def n_prefill_gpus(self) -> int:
return self.p_dp * self.p_tp
CONFIGS = [
PDConfig("1P_TP8:1D_TP8", "TP8", "TP8"),
PDConfig("2P_TP8:1D_DP8EP", "TP8", "DP8EP"),
PDConfig("2P_TP8:2D_TP8", "TP8", "TP8"),
PDConfig("2P_DP8EP:3D_DP8EP", "DP8EP", "DP8EP"),
PDConfig("2P_DP8EP:4D_TP8", "DP8EP", "TP8"),
]
CONFIG_IDS = [c.name for c in CONFIGS]
def make_decode_worker(*, world_size: int, tp_rank: int, dp_rank: int):
w = object.__new__(MoRIIOConnectorWorker)
w.world_size = world_size
w.tp_rank = tp_rank
w.dp_rank = dp_rank
w.use_mla = True
return w
def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0"):
return ReqMeta(
transfer_id="t",
local_block_ids=[1],
remote_block_ids=[2],
remote_host=host,
remote_port=1234,
remote_handshake_port=6301,
remote_notify_port=61005,
remote_engine_id=f"{host}:6301",
tp_size=p_tp,
remote_dp_size=p_dp,
remote_dp_rank=remote_dp_rank,
)
def build_decode_workers(cfg: PDConfig) -> list:
"""Decode workers that issue reads for one prefill instance. A TP decode
instance reads from every tp rank; a DP decode instance from each dp-rank
worker. Reused across requests so per-worker round-robin state advances."""
if cfg.d_tp > 1:
return [
make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0)
for r in range(cfg.d_tp)
]
return [
make_decode_worker(world_size=1, tp_rank=0, dp_rank=d) for d in range(cfg.d_dp)
]
def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter:
"""Drive the REAL _resolve_read_source over a fair input stream; tally which
prefill GPU each read targets. The only modelled assumption is the proxy
contract -- each owner dp-rank delivered equally (``for owner_dp in
range(p_dp)``); no proxy algorithm is reproduced."""
workers = build_decode_workers(cfg)
hits: Counter = Counter()
for _ in range(rounds):
for owner_dp in range(cfg.p_dp):
for worker in workers:
meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=owner_dp)
chosen_tp, _flexible = worker._resolve_read_source(meta)
target = get_port_offset(owner_dp, chosen_tp, cfg.p_tp)
assert 0 <= target < cfg.n_prefill_gpus
hits[target] += 1
return hits
@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS)
def test_kv_read_load_is_hardware_fair(cfg: PDConfig) -> None:
# rounds a multiple of p_tp so the round-robin closes on an exact cycle.
hits = prefill_target_multiset(cfg, rounds=16)
gpus = set(range(cfg.n_prefill_gpus))
assert set(hits) == gpus, f"unused prefill GPUs: {sorted(gpus - set(hits))}"
assert max(hits.values()) == min(hits.values()), dict(sorted(hits.items()))
@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS)
def test_flexible_gate_fires_only_for_tp_prefill_dp_decode(cfg: PDConfig) -> None:
flags = set()
for worker in build_decode_workers(cfg):
meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=0)
_chosen, flexible = worker._resolve_read_source(meta)
flags.add(flexible)
is_mirror = cfg.d_tp == 1 and cfg.p_dp == 1 and cfg.p_tp > 1
assert flags == {is_mirror}
def test_symmetric_tp_is_a_bijection() -> None:
targets = []
for tp_rank in range(8):
worker = make_decode_worker(world_size=8, tp_rank=tp_rank, dp_rank=0)
chosen_tp, flexible = worker._resolve_read_source(
make_meta(p_tp=8, p_dp=1, remote_dp_rank=0)
)
assert not flexible
targets.append(chosen_tp)
assert sorted(targets) == list(range(8))
def test_owner_dp_read_is_faithful_and_covers_every_rank() -> None:
worker = make_decode_worker(world_size=8, tp_rank=3, dp_rank=0)
targets = []
for owner_dp in range(8):
chosen_tp, flexible = worker._resolve_read_source(
make_meta(p_tp=1, p_dp=8, remote_dp_rank=owner_dp)
)
assert not flexible
assert chosen_tp == 0 # p_tp == 1, only tp0 exists
targets.append(get_port_offset(owner_dp, chosen_tp, 1))
assert sorted(targets) == list(range(8))
def test_flexible_round_robin_is_deterministic_uniform_and_staggered() -> None:
w0 = make_decode_worker(world_size=1, tp_rank=0, dp_rank=0)
seq = [w0._next_flex_tp_rank(8) for _ in range(64)]
assert seq[:8] == list(range(8))
assert Counter(seq) == Counter({t: 8 for t in range(8)})
first_pick = [
make_decode_worker(world_size=1, tp_rank=0, dp_rank=d)._next_flex_tp_rank(8)
for d in range(8)
]
assert sorted(first_pick) == list(range(8))
def test_read_blocks_for_req_threads_chosen_tp() -> None:
# The resolved (chosen_tp, flexible) must reach _read_blocks, which keys the
# session AND the notify port off that single value -- so a read and its
# completion notify address the same prefill rank.
worker = make_decode_worker(world_size=1, tp_rank=0, dp_rank=3)
worker._read_blocks = MagicMock()
worker._read_blocks_for_req("r", make_meta(p_tp=8, p_dp=1, remote_dp_rank=0))
kw = worker._read_blocks.call_args.kwargs
assert kw["flexible"] is True
assert kw["chosen_tp"] == 3 # first flexible pick = dp_rank seed
assert get_port_offset(0, kw["chosen_tp"], 8) == 3
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+5
View File
@@ -422,6 +422,10 @@ class ReqMeta:
remote_engine_id: str
tp_size: int
remote_dp_size: int
# Prefill DP rank that owns this request's KV (forwarded by the proxy). The
# read must target this rank's memory registration; the default 0 preserves
# the symmetric single-DP behaviour.
remote_dp_rank: int = 0
class MoRIIOConnectorMetadata(KVConnectorMetadata):
@@ -475,6 +479,7 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
remote_notify_port=int(remote_notify_port),
tp_size=kv_transfer_params.get("tp_size", 1),
remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0),
)
if write_mode:
self.reqs_to_save[request_id] = _req
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+252 -7
View File
@@ -1014,6 +1014,8 @@ class MoRIIOConnectorWorker:
self._handshake_futures: dict[EngineId, Future[set[str]]] = {}
# Protects _handshake_futures and _remote_agents.
self._handshake_lock = threading.RLock()
# Remote engines already covered by the eager pre-forward handshake.
self._eager_handshaked_engines: set[EngineId] = set()
self.block_size = vllm_config.cache_config.block_size
self.model_config = vllm_config.model_config
@@ -1261,8 +1263,16 @@ class MoRIIOConnectorWorker:
remote_tp_size: int,
expected_engine_id: str,
remote_dp_rank: int = 0,
remote_tp_rank: int | None = None,
) -> set[str]:
"""Do a MoRIIO handshake with a remote instance."""
"""Do a MoRIIO handshake with a remote instance.
remote_tp_rank: explicit remote TP index to dial. Flexible-read callers
pass the chosen prefill TP rank so the handshake, the (dp, tp) session
key and the notify port all address the SAME rank. None falls back to the
local-rank mapping _remote_tp_rank -- byte-identical for callers not yet
TP-aware.
"""
start_time = time.perf_counter()
@@ -1270,9 +1280,12 @@ class MoRIIOConnectorWorker:
# a hack to keep us moving. We will switch when moving to etcd
# or where we have a single ZMQ socket in the scheduler.
port_offset = get_port_offset(
remote_dp_rank, self._remote_tp_rank(remote_tp_size)
dial_tp_rank = (
self._remote_tp_rank(remote_tp_size)
if remote_tp_rank is None
else int(remote_tp_rank)
)
port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size)
path = make_zmq_path("tcp", host, port + port_offset)
logger.debug("handshake Querying metadata on path: %s", path)
@@ -1729,6 +1742,149 @@ class MoRIIOConnectorWorker:
def get_engine_name_with_dp(self, engine_name, dp_rank):
return f"{engine_name}_dp{dp_rank}"
def get_engine_name_with_dp_tp(self, engine_name, dp_rank, tp_rank):
# Per-(dp, tp) session key. The flexible mirror read keys sessions per
# (dp, tp) so one decode worker can hold a session to EACH prefill TP
# rank and spread reads across them; other configs keep the DP-only key.
return f"{engine_name}_dp{dp_rank}_tp{tp_rank}"
def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> None:
"""Handshake EVERY remote prefill DP rank BEFORE the decode forward pass,
identically across all local TP workers.
Why this exists (the deadlock it prevents): with heterogeneous DP prefill
a decode TP worker reads KV from whichever prefill DP rank owns the
request, so across requests every worker must reach several prefill DP
ranks. The decode forward issues per-layer TP collectives (e.g. an
all-gather) that all local TP workers must enter together. If the
handshakes are left to fire lazily on the read path, the workers diverge:
a worker whose target rank is already cached races ahead into the forward
collective while a peer is still blocked in a handshake recv(). The first
worker then waits inside the collective for the stuck peer -> 600s NCCL
timeout / hang. This was observed directly with mixed TP<->DP configs.
Fix: complete ALL prefill-DP-rank handshakes for every referenced remote
engine HERE, before any read enters the forward, so no worker is still
handshaking once its peers reach a collective. Fires ONCE per remote
engine (first contact), gated by _eager_handshaked_engines. The engine
set comes from scheduler-built metadata (identical on every TP worker),
so all workers run the same handshakes in the same order and reach the
all-reduce barrier below together.
Failure handling: handshake exceptions are caught, never raised before
the collective (raising early would hang the peers still waiting for it).
Every worker reaches the all-reduce(MIN) vote; if ANY worker failed, ALL
raise the same error AFTER the collective, so the step fails fast and
uniformly in ~seconds instead of one rank hanging the forward for 600s.
"""
import torch.distributed as dist
# Distinct remote engines referenced this step, in metadata (==
# scheduler) order so every TP worker iterates engines identically.
engines: dict[str, ReqMeta] = {}
for _req_id, meta in metadata.reqs_to_recv.items():
remote_engine_id = (
str(meta.remote_host) + ":" + str(meta.remote_handshake_port)
)
engines.setdefault(remote_engine_id, meta)
for remote_engine_id, meta in engines.items():
if remote_engine_id in self._eager_handshaked_engines:
continue
remote_dp_size = int(meta.remote_dp_size)
port = int(meta.remote_handshake_port)
tp_size = int(meta.tp_size)
# Flexible mirror (TP prefill + MLA, world_size==1 decode): the read
# round-robins over prefill TP ranks, so pre-warm a session to EVERY
# (dp, tp) rank. Other configs pre-warm per DP rank (tp resolved by
# the fixed local-rank mapping) -- byte-identical to before. The
# mirror's decode is DP+EP, whose forward all-to-all is the collective
# that the eager barrier keeps everyone in step for.
flexible = (
self.world_size == 1
and self.use_mla
and remote_dp_size == 1
and tp_size > 1
)
# (engine_id, dp_rank, tp_rank_or_None); tp_rank is None on the legacy
# path so _moriio_handshake falls back to its _remote_tp_rank mapping.
targets: list[tuple[Any, int, int | None]]
if flexible:
targets = [
(self.get_engine_name_with_dp_tp(remote_engine_id, dp, tp), dp, tp)
for dp in range(remote_dp_size)
for tp in range(max(1, tp_size))
]
else:
targets = [
(self.get_engine_name_with_dp(remote_engine_id, dp), dp, None)
for dp in range(remote_dp_size)
]
# Submit handshakes for every not-yet-known target UNDER the lock; do
# NOT hold it across the join or the collective (a stalled recv must
# not block another thread's lock acquisition). Gate on BOTH
# _remote_agents AND layer metadata: a rank with an agent entry but no
# layer metadata is half-handshaked and would KeyError at read time.
futures: list[tuple[str, Future[set[str]]]] = []
with self._handshake_lock:
for eid, cur_dp_rank, cur_tp_rank in targets:
if (
eid in self._remote_agents
and eid in self.layer_name_to_remote_kv_cache_metadata
):
continue
fut = self._handshake_initiation_executor.submit(
self._moriio_handshake,
meta.remote_host,
port,
tp_size,
eid,
cur_dp_rank,
cur_tp_rank,
)
futures.append((eid, fut))
# Join outside the lock. Bounded handshake errors are recorded here
# and reported after the all-reduce.
all_ok = True
results: dict[str, set[str]] = {}
for eid, fut in futures:
try:
results[eid] = fut.result()
except Exception:
logger.exception("Eager MoRIIO handshake failed for %s", eid)
all_ok = False
with self._handshake_lock:
for eid, agents in results.items():
self._remote_agents[eid] = agents
logger.info(
"Eager MoRIIO handshake: engine=%s dp_size=%d new_ranks=%d "
"ok=%s tp_rank=%d",
remote_engine_id,
remote_dp_size,
len(futures),
all_ok,
self.tp_rank,
)
# CPU all-reduce = TP-uniform success vote AND lockstep barrier: it
# blocks until every TP worker arrives, gives them the same verdict,
# and stays off the model compute stream.
vote = torch.tensor([1 if all_ok else 0], device="cpu", dtype=torch.int32)
dist.all_reduce(vote, group=self.tp_group.cpu_group, op=dist.ReduceOp.MIN)
if int(vote.item()) == 0:
raise HandshakeError(
f"Eager MoRIIO handshake failed for {remote_engine_id} on "
"at least one TP rank; failing this step fast to avoid a "
"TP collective hang"
)
self._eager_handshaked_engines.add(remote_engine_id)
def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
"""
Start loading by triggering non-blocking moriio_xfer.
@@ -1750,6 +1906,12 @@ class MoRIIOConnectorWorker:
if self.mode == MoRIIOMode.WRITE:
return
# Handshake every referenced remote prefill rank up front, before any
# read enters the forward pass. A lazy per-rank handshake on the read
# path lets TP workers diverge into a forward collective while a peer is
# still blocked handshaking -> NCCL hang (see below).
self._eager_handshake_all_dp_ranks(metadata)
wait_handshake_readd_req = False
remote_engine_id = None
@@ -1758,8 +1920,15 @@ class MoRIIOConnectorWorker:
str(meta.remote_host) + ":" + str(meta.remote_handshake_port)
)
meta.remote_engine_id = remote_engine_id
# The eager handshake above already covered every referenced engine
# (and keys the mirror per (dp, tp), which the DP-only dp0 probe below
# would miss). Only fall back to the lazy background handshake for an
# engine it did not cover.
dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0)
if dp0_remote_engine_id not in self._remote_agents:
if (
remote_engine_id not in self._eager_handshaked_engines
and dp0_remote_engine_id not in self._remote_agents
):
# Initiate handshake with remote engine to exchange metadata.
with self._handshake_lock:
if remote_engine_id not in self._remote_agents:
@@ -1809,12 +1978,53 @@ class MoRIIOConnectorWorker:
self.save_kv_layer(metadata, layer_name, kv_layer, None)
self._writer.seal_pending_transfers()
def _next_flex_tp_rank(self, remote_tp_size: int) -> int:
"""Deterministic round-robin over prefill tp0..N-1 for the flexible read.
Round-robin (not random): exactly uniform and testable, with the same
prefill-NIC balancing. Seeded from this decode rank's dp_rank so
concurrent decode DP ranks are phase-staggered -- at a given read index
distinct decode ranks target distinct prefill TP ranks.
"""
rr = getattr(self, "_flex_tp_rr", None)
if rr is None:
rr = int(getattr(self, "dp_rank", 0) or 0)
self._flex_tp_rr = rr + 1
return rr % remote_tp_size
def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]:
"""Resolve (chosen_tp, flexible) for reading this request's KV.
Flexible mirror (decode world_size==1 + MLA + pure-TP prefill): MLA
replicates the latent KV across the prefill TP ranks, so any is a valid
source; round-robin across them to spread RDMA/NIC load. Otherwise the
source TP rank is fixed by the local-rank mapping (_remote_tp_rank) --
forward DP8EP->TP8 -> tp0; symmetric TP -> tp_rank -- byte-identical to
prior behaviour. chosen_tp is the single value threaded into the (dp, tp)
session key, the handshake dial and the notify port, so all three address
the SAME prefill rank (drift -> read one rank but notify another -> the
read rank's prefill buffer is never freed).
"""
remote_tp_size = int(meta.tp_size)
flexible = (
self.world_size == 1
and self.use_mla
and int(meta.remote_dp_size) == 1
and remote_tp_size > 1
)
if flexible:
chosen_tp = self._next_flex_tp_rank(remote_tp_size)
else:
chosen_tp = self._remote_tp_rank(remote_tp_size)
return chosen_tp, flexible
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
logger.debug(
"Remote agent %s available, calling _read_blocks for req %s",
meta.remote_engine_id,
req_id,
)
chosen_tp, flexible = self._resolve_read_source(meta)
self._read_blocks(
request_id=req_id,
transfer_id=meta.transfer_id,
@@ -1824,6 +2034,9 @@ class MoRIIOConnectorWorker:
remote_host=meta.remote_host,
remote_notify_port=meta.remote_notify_port,
remote_tp_size=meta.tp_size,
remote_dp_rank=meta.remote_dp_rank,
chosen_tp=chosen_tp,
flexible=flexible,
)
def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
@@ -1976,12 +2189,37 @@ class MoRIIOConnectorWorker:
remote_host: str,
remote_notify_port: int,
remote_tp_size: int,
remote_dp_rank: int = 0,
chosen_tp: int | None = None,
flexible: bool = False,
) -> None:
if self.mode == MoRIIOMode.WRITE:
return
dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)
sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)
# Read from the prefill rank that actually computed this request's KV
# (forwarded by the proxy). Hardcoding DP0 reads from a different rank's
# memory registration; per-rank num_blocks differ, so high block ids can
# overrun the wrong rank's region.
#
# eff_tp = the remote TP rank this read targets. The flexible mirror
# reads from a round-robin-chosen prefill TP rank and keys the session
# per (dp, tp); other configs use the fixed local-rank mapping (eff_tp ==
# _remote_tp_rank), byte-identical to before. This key MUST match the one
# the eager handshake stored the session under.
eff_tp = (
int(chosen_tp)
if chosen_tp is not None
else self._remote_tp_rank(remote_tp_size)
)
if flexible:
remote_dp_engine_id = self.get_engine_name_with_dp_tp(
dst_engine_id, int(remote_dp_rank), eff_tp
)
else:
remote_dp_engine_id = self.get_engine_name_with_dp(
dst_engine_id, int(remote_dp_rank)
)
sessions, remote_moriio_meta = self._get_built_session(remote_dp_engine_id)
# SQ-full backpressure deadline, shared across this request's layers.
_sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout
@@ -2030,6 +2268,13 @@ class MoRIIOConnectorWorker:
self._recving_transfers[request_id].append(transfer_status)
self._recving_transfers_callback_addr[request_id] = (
remote_host,
str(remote_notify_port + self._remote_tp_rank(remote_tp_size)),
str(
remote_notify_port
+ get_port_offset(
int(remote_dp_rank),
eff_tp,
remote_tp_size,
)
),
transfer_id,
)