[KVOffload][P2P] Well-known default host/port env vars and per-DP-rank control port (#47636)
Signed-off-by: Liran Schour <lirans@il.ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
liranschour committed
Jul 15, 2026 at 15:22 UTC
615834ee584fd6395f1ac4d2257cb29c47895295
13 files changed
+413
-61
docs/features/kv_offloading_usage.md
+11
-2
@@ -152,16 +152,25 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti
152
153
The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required.
154
155
+PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes.
156
+
157
| Key | Required | Default | Notes |
158
| --- | --- | --- | --- |
159
| `type` | yes | — | Must be `p2p`. |
158
-| `host` | no | `0.0.0.0` | Address the control socket binds to. |
159
-| `port` | no | `7777` | Port for the control socket. Must be reachable from peers. |
160
+| `host` | no | `$VLLM_P2P_SIDE_CHANNEL_HOST` (`localhost`) | Address the control socket binds to, used verbatim as the identity peers dial back. When omitted, resolves from the env var below. The `localhost` default binds loopback only — for cross-host P2P you **must** set it to the node's routable IP (see below). |
161
+| `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. |
162
| `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. |
163
| `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |
164
165
The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process.
166
167
+#### Environment Variables
168
+
169
+Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence.
170
+
171
+- `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide.
172
+- `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision.
173
+
174
## Tuning Tips
175
176
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
tests/test_envs.py
+12
@@ -36,6 +36,18 @@ def test_nixl_side_channel_host_is_not_compile_factor(
36
assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in envs.compile_factors()
37
38
39
+def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch):
40
+ monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False)
41
+ monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False)
42
+ assert envs.VLLM_P2P_SIDE_CHANNEL_HOST == "localhost"
43
+ assert envs.VLLM_P2P_SIDE_CHANNEL_PORT == 5710
44
+
45
+ monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.0.0.20")
46
+ monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799")
47
+ assert envs.VLLM_P2P_SIDE_CHANNEL_HOST == "10.0.0.20"
48
+ assert envs.VLLM_P2P_SIDE_CHANNEL_PORT == 5799
49
+
50
+
51
def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
52
monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1")
53
monkeypatch.setenv("VLLM_PORT", "1234")
tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py
+91
-18
@@ -14,7 +14,7 @@ Usage:
14
--port 8192 \
15
--prefiller-host 127.0.0.1 --prefiller-port 8100 \
16
--decoder-host 127.0.0.1 --decoder-port 8200 \
17
- --p2p-connector-host 127.0.0.1 --p2p-connector-port 7777
17
+ --p2p-connector-host 127.0.0.1 --p2p-connector-port 5710
18
"""
19
20
import argparse
@@ -73,14 +73,23 @@ async def lifespan(app: FastAPI):
73
app.state.prefill_iterator = itertools.cycle(range(len(app.state.prefill_clients)))
74
app.state.decode_iterator = itertools.cycle(range(len(app.state.decode_clients)))
75
76
+ # Round-robin over each role's data-parallel replicas (independent of the
77
+ # instance iterators above). For dp_size=1 these yield only rank 0.
78
+ app.state.prefill_dp_iterator = itertools.cycle(
79
+ range(global_args.prefiller_dp_size)
80
+ )
81
+ app.state.decode_dp_iterator = itertools.cycle(range(global_args.decoder_dp_size))
82
+
83
mode = "decoder-first" if global_args.decoder_first else "prefiller-first"
84
pd_host = global_args.p2p_connector_host
85
pd_port = global_args.p2p_connector_port
86
+ n_pref = len(app.state.prefill_clients)
87
+ n_dec = len(app.state.decode_clients)
88
print(
89
f"Proxy ready [{mode}]: "
81
- f"{len(app.state.prefill_clients)} prefiller(s), "
82
- f"{len(app.state.decode_clients)} decoder(s). "
83
- f"P2PConnector at {pd_host}:{pd_port}"
90
+ f"{n_pref} prefiller(s) x dp={global_args.prefiller_dp_size}, "
91
+ f"{n_dec} decoder(s) x dp={global_args.decoder_dp_size}. "
92
+ f"P2PConnector base at {pd_host}:{pd_port}"
93
)
94
yield
95
@@ -111,8 +120,9 @@ def parse_args():
120
p.add_argument(
121
"--p2p-connector-port",
122
type=int,
114
- default=7777,
115
- help="Port of the prefiller's P2PConnector ZMQ socket",
123
+ default=int(os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710")),
124
+ help="Port of the prefiller's P2PConnector ZMQ socket "
125
+ "(default: $VLLM_P2P_SIDE_CHANNEL_PORT or 5710)",
126
)
127
# P2PConnector coordinates of the decoder — injected into prefill requests
128
# so the prefiller's submit_store can resolve the peer to push KV to.
@@ -125,8 +135,9 @@ def parse_args():
135
p.add_argument(
136
"--decoder-p2p-connector-port",
137
type=int,
128
- default=7778,
129
- help="Port of the decoder's P2PConnector ZMQ socket",
138
+ default=int(os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710")) + 1,
139
+ help="Port of the decoder's P2PConnector ZMQ socket "
140
+ "(default: $VLLM_P2P_SIDE_CHANNEL_PORT + 1 or 5711)",
141
)
142
p.add_argument(
143
"--decoder-first",
@@ -134,6 +145,23 @@ def parse_args():
145
help="Send decode request before prefill so decoder is already "
146
"waiting when KV blocks arrive (decoder-first mode)",
147
)
148
+ p.add_argument(
149
+ "--prefiller-dp-size",
150
+ type=int,
151
+ default=1,
152
+ help="Data-parallel replica count of the prefiller. When >1 the proxy "
153
+ "round-robins prefill across ranks via the X-data-parallel-rank header "
154
+ "and injects remote_port = p2p-connector-port + rank into the decode "
155
+ "request so the decoder pulls KV from that replica's control socket. "
156
+ "Assumes a single prefiller instance (one host/port fronting N replicas).",
157
+ )
158
+ p.add_argument(
159
+ "--decoder-dp-size",
160
+ type=int,
161
+ default=1,
162
+ help="Data-parallel replica count of the decoder. When >1 the proxy "
163
+ "round-robins decode across ranks via the X-data-parallel-rank header.",
164
+ )
165
args = p.parse_args()
166
if len(args.prefiller_hosts) != len(args.prefiller_ports):
167
raise ValueError("Prefiller host/port count mismatch")
@@ -150,6 +178,23 @@ def _get_next(app, service: str):
178
return app.state.decode_clients[next(app.state.decode_iterator)]
179
180
181
+def _next_dp_rank(app, service: str):
182
+ """Advance the round-robin DP cursor for a role.
183
+
184
+ Returns (rank, header_rank). ``rank`` (0..dp_size-1) is always used for the
185
+ P2P remote_port arithmetic; ``header_rank`` is the same value when the role
186
+ has dp_size>1 and None otherwise (so dp=1 sends no header, matching the
187
+ single-replica behavior).
188
+ """
189
+ if service == "prefill":
190
+ rank = next(app.state.prefill_dp_iterator)
191
+ dp_size = global_args.prefiller_dp_size
192
+ else:
193
+ rank = next(app.state.decode_dp_iterator)
194
+ dp_size = global_args.decoder_dp_size
195
+ return rank, (rank if dp_size > 1 else None)
196
+
197
+
198
def _auth_headers(request_id: str) -> dict:
199
headers: dict = {"X-Request-Id": request_id}
200
api_key = os.environ.get("OPENAI_API_KEY", "")
@@ -158,7 +203,7 @@ def _auth_headers(request_id: str) -> dict:
203
return headers
204
205
161
-async def _prefill(client_info, endpoint, req_data, request_id):
206
+async def _prefill(client_info, endpoint, req_data, request_id, dp_rank=None):
207
"""Send a prefill-only request (max_tokens=1) to the prefiller."""
208
data = req_data.copy()
209
data["kv_transfer_params"] = {
@@ -174,14 +219,21 @@ async def _prefill(client_info, endpoint, req_data, request_id):
219
data.pop("min_completion_tokens", None)
220
221
headers = _auth_headers(request_id)
222
+ if dp_rank is not None:
223
+ # Pin this prefill to a specific DP replica so its P2PConnector (bound
224
+ # at base_port + data_parallel_index) holds the produced KV blocks.
225
+ headers["X-data-parallel-rank"] = str(dp_rank)
226
resp = await client_info["client"].post(endpoint, json=data, headers=headers)
227
resp.raise_for_status()
228
await resp.aread()
229
return resp
230
231
183
-async def _stream_decode(client_info, endpoint, req_data, request_id):
232
+async def _stream_decode(client_info, endpoint, req_data, request_id, dp_rank=None):
233
headers = _auth_headers(request_id)
234
+ if dp_rank is not None:
235
+ # Round-robin this decode onto a specific DP replica of the decoder.
236
+ headers["X-data-parallel-rank"] = str(dp_rank)
237
async with client_info["client"].stream(
238
"POST", endpoint, json=req_data, headers=headers
239
) as resp:
@@ -195,24 +247,39 @@ async def _handle_completions(api: str, request: Request):
247
req_data = await request.json()
248
request_id = str(uuid.uuid4())
249
250
+ # Pick DP replicas for both roles up front (before any await) so the
251
+ # two round-robin advances are atomic together. Header rank is None for
252
+ # a role with dp_size==1 (no header → single replica, unchanged).
253
+ prefill_rank, prefill_hdr = _next_dp_rank(request.app, "prefill")
254
+ decode_rank, decode_hdr = _next_dp_rank(request.app, "decode")
255
+
256
prefill_client = _get_next(request.app, "prefill")
199
- await _prefill(prefill_client, api, req_data, request_id)
257
+ await _prefill(prefill_client, api, req_data, request_id, dp_rank=prefill_hdr)
258
259
# Inject the prefiller's P2PConnector address so the decoder can pull
202
- # KV blocks from it via the P2PConnector transport.
260
+ # KV blocks from it. remote_port = base + prefill_rank targets the
261
+ # replica that produced the KV (base+0 == base when dp=1).
262
req_data["kv_transfer_params"] = {
263
"prefill": {
264
"kv_request_id": request_id,
265
"remote_host": global_args.p2p_connector_host,
207
- "remote_port": global_args.p2p_connector_port,
266
+ "remote_port": global_args.p2p_connector_port + prefill_rank,
267
},
268
}
269
270
decode_client = _get_next(request.app, "decode")
212
- logger.debug("prefill=%s decode=%s", prefill_client, decode_client)
271
+ logger.debug(
272
+ "prefill=%s dp=%s decode=%s dp=%s",
273
+ prefill_client,
274
+ prefill_rank,
275
+ decode_client,
276
+ decode_rank,
277
+ )
278
279
async def generate():
215
- async for chunk in _stream_decode(decode_client, api, req_data, request_id):
280
+ async for chunk in _stream_decode(
281
+ decode_client, api, req_data, request_id, dp_rank=decode_hdr
282
+ ):
283
yield chunk
284
285
return StreamingResponse(generate(), media_type="application/json")
@@ -237,6 +304,10 @@ async def _handle_completions_decoder_first(api: str, request: Request):
304
req_data = await request.json()
305
request_id = str(uuid.uuid4())
306
307
+ # Pick DP replicas up front (see _handle_completions for rationale).
308
+ prefill_rank, prefill_hdr = _next_dp_rank(request.app, "prefill")
309
+ decode_rank, decode_hdr = _next_dp_rank(request.app, "decode")
310
+
311
prefill_client = _get_next(request.app, "prefill")
312
decode_client = _get_next(request.app, "decode")
313
@@ -245,7 +316,7 @@ async def _handle_completions_decoder_first(api: str, request: Request):
316
"prefill": {
317
"kv_request_id": request_id,
318
"remote_host": global_args.p2p_connector_host,
248
- "remote_port": global_args.p2p_connector_port,
319
+ "remote_port": global_args.p2p_connector_port + prefill_rank,
320
},
321
}
322
@@ -255,7 +326,7 @@ async def _handle_completions_decoder_first(api: str, request: Request):
326
async def _run_decode():
327
try:
328
async for chunk in _stream_decode(
258
- decode_client, api, decode_data, request_id
329
+ decode_client, api, decode_data, request_id, dp_rank=decode_hdr
330
):
331
await queue.put(("data", chunk))
332
except Exception as exc:
@@ -268,7 +339,9 @@ async def _handle_completions_decoder_first(api: str, request: Request):
339
340
# 2. Send prefill — blocks are computed and pushed to the decoder
341
try:
271
- await _prefill(prefill_client, api, req_data, request_id)
342
+ await _prefill(
343
+ prefill_client, api, req_data, request_id, dp_rank=prefill_hdr
344
+ )
345
except Exception as exc:
346
logger.warning("decoder-first: prefill failed: %s", exc)
347
tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh
+34
-3
@@ -16,6 +16,12 @@
16
# NUM_DECODE_INSTANCES default 1
17
# PREFILLER_TP_SIZE default 1
18
# DECODER_TP_SIZE default 1
19
+# DP_EP when set, deploy the decoder in DP-EP attention
20
+# mode (dp=DECODER_TP_SIZE, tp=1,
21
+# --enable-expert-parallel) instead of tensor
22
+# parallel; the proxy round-robins decode across
23
+# the DP replicas. Mirrors
24
+# nixl_integration/run_accuracy_test.sh.
25
# GPU_MEMORY_UTILIZATION default 0.45
26
# MAX_MODEL_LEN default 512
27
# PREFILL_BLOCK_SIZE default 128
@@ -29,6 +35,8 @@
35
# NUM_PREFILL_INSTANCES=2 NUM_DECODE_INSTANCES=2 \
36
# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh
37
# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh --decoder-first
38
+# DP_EP=1 DECODER_TP_SIZE=2 \
39
+# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh
40
41
set -xe
42
@@ -71,6 +79,14 @@ NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1}
79
NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1}
80
PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
81
DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
82
+# DP-EP attention mode (see header). When DP_EP is set the decoder uses
83
+# data-parallel replicas (dp=DECODER_TP_SIZE) that the proxy round-robins over.
84
+DP_EP=${DP_EP:-}
85
+if [[ -n "$DP_EP" ]]; then
86
+ DECODER_DP_SIZE=${DECODER_TP_SIZE}
87
+else
88
+ DECODER_DP_SIZE=1
89
+fi
90
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.45}
91
MAX_MODEL_LEN=${MAX_MODEL_LEN:-512}
92
PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
@@ -162,7 +178,11 @@ run_tests_for_model() {
178
echo "================================"
179
echo "Testing model: $model_name"
180
echo " prefillers=${NUM_PREFILL_INSTANCES} (tp=${PREFILLER_TP_SIZE})"
165
- echo " decoders=${NUM_DECODE_INSTANCES} (tp=${DECODER_TP_SIZE})"
181
+ if [[ -n "$DP_EP" ]]; then
182
+ echo " decoders=${NUM_DECODE_INSTANCES} (dp-ep=${DECODER_TP_SIZE}, tp=1)"
183
+ else
184
+ echo " decoders=${NUM_DECODE_INSTANCES} (tp=${DECODER_TP_SIZE})"
185
+ fi
186
echo " decoder_first=${DECODER_FIRST}"
187
echo "================================"
188
@@ -241,9 +261,19 @@ run_tests_for_model() {
261
--block-size ${DECODE_BLOCK_SIZE} \
262
--gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \
263
--max-model-len ${MAX_MODEL_LEN} \
244
- --tensor-parallel-size ${DECODER_TP_SIZE} \
264
--kv-transfer-config '${kv_cfg}'"
265
266
+ # DP-EP attention mode: data-parallel + expert-parallel decode replicas
267
+ # (dp=DECODER_TP_SIZE, tp=1) instead of tensor parallel. Mirrors
268
+ # nixl_integration/run_accuracy_test.sh.
269
+ if [[ -z "$DP_EP" ]]; then
270
+ BASE_CMD="${BASE_CMD} --tensor-parallel-size ${DECODER_TP_SIZE}"
271
+ else
272
+ echo "DP-EP Attention enabled, deploying decoder with dp=${DECODER_TP_SIZE} and tp=1"
273
+ BASE_CMD="${BASE_CMD} --data-parallel-size ${DECODER_TP_SIZE} \
274
+ --tensor-parallel-size 1 --enable-expert-parallel"
275
+ fi
276
+
277
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
278
IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
279
for arg in "${extra_args[@]}"; do
@@ -282,7 +312,8 @@ run_tests_for_model() {
312
--p2p-connector-host ${P2P_HOST} \
313
--p2p-connector-port ${PREFILL_PD_PORTS[0]} \
314
--decoder-p2p-connector-host ${P2P_HOST} \
285
- --decoder-p2p-connector-port ${DECODE_PD_PORTS[0]}"
315
+ --decoder-p2p-connector-port ${DECODE_PD_PORTS[0]} \
316
+ --decoder-dp-size ${DECODER_DP_SIZE}"
317
318
if [[ "${DECODER_FIRST}" == "true" ]]; then
319
PROXY_CMD="${PROXY_CMD} --decoder-first"
tests/v1/kv_offload/tiering/p2p/test_data_transport.py
+29
@@ -183,6 +183,35 @@ class TestNixlTransportWithMockedAgent:
183
assert result.done == ()
184
assert result.failed == ()
185
186
+ def test_poll_peer_id_scopes_to_peer(self):
187
+ """poll(peer_id) drains only that peer's transfers.
188
+
189
+ Regression: the transport is shared across peer sessions (e.g. a
190
+ single prefiller serving a DP>1 decoder). An unscoped poll by one
191
+ session used to consume and discard sibling sessions' completions,
192
+ starving them until timeout. poll(peer_id) must leave other peers'
193
+ transfers inflight.
194
+ """
195
+ transport = self._make_transport()
196
+ transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
197
+ transport.add_remote_peer("peer:2", b"meta", 0x2000, 8, 1024)
198
+
199
+ tid1 = transport.write_blocks("peer:1", [0], [1])
200
+ tid2 = transport.write_blocks("peer:2", [2], [3])
201
+ transport._agent.check_xfer_state.return_value = "DONE"
202
+
203
+ # Polling peer:1 must not consume peer:2's completed transfer.
204
+ result = transport.poll(peer_id="peer:1")
205
+ assert tid1 in result.done
206
+ assert tid2 not in result.done
207
+ assert tid1 not in transport._inflight
208
+ assert tid2 in transport._inflight
209
+
210
+ # peer:2 sees its own completion when it polls.
211
+ result2 = transport.poll(peer_id="peer:2")
212
+ assert tid2 in result2.done
213
+ assert tid2 not in transport._inflight
214
+
215
def test_cancel_removes_inflight(self):
216
"""cancel removes transfers and releases handles."""
217
transport = self._make_transport()
tests/v1/kv_offload/tiering/p2p/test_manager.py
+116
-2
@@ -9,9 +9,11 @@ using fake transport and session objects.
9
from __future__ import annotations
10
11
import time
12
+import uuid
13
from types import SimpleNamespace
14
15
import numpy as np
16
+import pytest
17
18
from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext
19
from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
@@ -598,7 +600,7 @@ class _ShutdownFakeData:
600
return list(self._still_queue[0])
601
return []
602
601
- def poll(self):
603
+ def poll(self, peer_id=None):
604
self.poll_calls += 1
605
606
class _Empty:
@@ -831,7 +833,7 @@ class _FakeData:
833
self._inflight_done.append(tid)
834
return tid
835
834
- def poll(self):
836
+ def poll(self, peer_id=None):
837
from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
838
839
done = self._inflight_done[:]
@@ -1373,3 +1375,115 @@ class TestConnectionDeathMidTransfer:
1375
assert (900, False) not in finishes
1376
assert (900, True) not in finishes
1377
assert "req-store" in mgr_a._unbound_stores
1378
+
1379
+
1380
+# ---------------------------------------------------------------------------
1381
+# Tests for host/port resolution in __init__ (env-var defaults)
1382
+# ---------------------------------------------------------------------------
1383
+
1384
+
1385
+class TestBindHostPortDefaults:
1386
+ """host/port fall back to VLLM_P2P_SIDE_CHANNEL_* when not in config."""
1387
+
1388
+ @staticmethod
1389
+ def _construct(monkeypatch, dp_index=0, **kwargs) -> P2PSecondaryTierManager:
1390
+ """Build a manager with the transports/file-mapper stubbed out.
1391
+
1392
+ The host is used verbatim (no resolution). The transport constructor
1393
+ args are recorded on ``mgr._test_calls`` so tests can assert the ZMQ
1394
+ identity (``host:port``) stays decoupled from the NIXL agent name
1395
+ (a uuid).
1396
+ """
1397
+ monkeypatch.setattr(
1398
+ manager_module,
1399
+ "FileMapper",
1400
+ SimpleNamespace(
1401
+ from_offloading_spec=lambda **_: SimpleNamespace(
1402
+ get_run_config=lambda: {}
1403
+ )
1404
+ ),
1405
+ )
1406
+ calls: dict = {}
1407
+ monkeypatch.setattr(
1408
+ manager_module,
1409
+ "NixlTransport",
1410
+ lambda agent_name, *a, **k: calls.update(nixl_name=agent_name)
1411
+ or SimpleNamespace(),
1412
+ )
1413
+ monkeypatch.setattr(
1414
+ manager_module,
1415
+ "ZmqTransport",
1416
+ lambda local_id, host, port, *a, **k: calls.update(
1417
+ zmq_id=local_id, zmq_host=host, zmq_port=port
1418
+ )
1419
+ or SimpleNamespace(),
1420
+ )
1421
+ spec = SimpleNamespace(
1422
+ block_size_factor=1,
1423
+ vllm_config=SimpleNamespace(
1424
+ parallel_config=SimpleNamespace(data_parallel_index=dp_index)
1425
+ ),
1426
+ )
1427
+ mgr = P2PSecondaryTierManager(spec, memoryview(b""), **kwargs)
1428
+ mgr._test_calls = calls
1429
+ return mgr
1430
+
1431
+ def test_defaults_from_env_unset(self, monkeypatch):
1432
+ monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False)
1433
+ monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False)
1434
+ mgr = self._construct(monkeypatch)
1435
+ # localhost default is used verbatim as the dial-back identity.
1436
+ assert mgr._local_id == "localhost:5710"
1437
+
1438
+ def test_env_override(self, monkeypatch):
1439
+ monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.1.2.3")
1440
+ monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799")
1441
+ mgr = self._construct(monkeypatch)
1442
+ assert mgr._local_id == "10.1.2.3:5799"
1443
+
1444
+ def test_explicit_config_wins(self, monkeypatch):
1445
+ monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.1.2.3")
1446
+ monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799")
1447
+ mgr = self._construct(monkeypatch, host="192.0.2.5", port=6001)
1448
+ assert mgr._local_id == "192.0.2.5:6001"
1449
+
1450
+ def test_dp_index_offsets_default_port(self, monkeypatch):
1451
+ monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False)
1452
+ monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False)
1453
+ mgr = self._construct(monkeypatch, dp_index=2)
1454
+ assert mgr._local_id == "localhost:5712"
1455
+
1456
+ def test_dp_index_offsets_explicit_port(self, monkeypatch):
1457
+ mgr = self._construct(monkeypatch, dp_index=1, host="192.0.2.5", port=6000)
1458
+ assert mgr._local_id == "192.0.2.5:6001"
1459
+
1460
+ @pytest.mark.parametrize(
1461
+ "bind_host", ["localhost", "127.0.0.1", "::1", "0.0.0.0", "::", "192.0.2.5"]
1462
+ )
1463
+ def test_host_used_verbatim(self, monkeypatch, bind_host):
1464
+ # The host is never rewritten — loopbacks and wildcards alike become
1465
+ # the dial-back identity verbatim (mirrors the NIXL connector).
1466
+ mgr = self._construct(monkeypatch, host=bind_host, port=5710)
1467
+ assert mgr._local_id == f"{bind_host}:5710"
1468
+
1469
+ def test_nixl_name_decoupled_from_identity(self, monkeypatch):
1470
+ # The ZMQ identity is the verbatim host:port; the NIXL agent name is a
1471
+ # uuid, distinct from the identity and never used as an address.
1472
+ mgr = self._construct(monkeypatch, host="127.0.0.1", port=5710)
1473
+ assert mgr._test_calls["zmq_host"] == "127.0.0.1"
1474
+ assert mgr._test_calls["zmq_port"] == 5710
1475
+ assert mgr._test_calls["zmq_id"] == "127.0.0.1:5710"
1476
+ assert mgr._local_id == "127.0.0.1:5710"
1477
+ # nixl_name is a valid uuid4 and not the host:port identity.
1478
+ nixl_name = mgr._test_calls["nixl_name"]
1479
+ assert nixl_name != mgr._local_id
1480
+ assert uuid.UUID(nixl_name).version == 4
1481
+
1482
+ def test_same_host_port_gets_distinct_nixl_names(self, monkeypatch):
1483
+ # The original collision: two peers sharing a host:port must still get
1484
+ # distinct NIXL agent names so add_remote_agent doesn't reject a remote
1485
+ # whose name equals the local. The per-process uuid guarantees this.
1486
+ mgr_a = self._construct(monkeypatch, host="localhost", port=5710)
1487
+ mgr_b = self._construct(monkeypatch, host="localhost", port=5710)
1488
+ assert mgr_a._local_id == mgr_b._local_id == "localhost:5710"
1489
+ assert mgr_a._nixl_agent_name != mgr_b._nixl_agent_name
tests/v1/kv_offload/tiering/p2p/test_sessions.py
+1
-1
@@ -110,7 +110,7 @@ class FakeDataTransport:
110
self._transfers[tid] = (peer_id, local_idxs, remote_idxs)
111
return tid
112
113
- def poll(self):
113
+ def poll(self, peer_id=None):
114
from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
115
116
result = PollResult(done=list(self._poll_done), failed=list(self._poll_failed))
vllm/envs.py
+12
@@ -209,6 +209,8 @@ if TYPE_CHECKING:
209
VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False
210
VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
211
VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
212
+ VLLM_P2P_SIDE_CHANNEL_HOST: str = "localhost"
213
+ VLLM_P2P_SIDE_CHANNEL_PORT: int = 5710
214
VLLM_EC_SIDE_CHANNEL_HOST: str = "localhost"
215
VLLM_EC_SIDE_CHANNEL_PORT: int = 5601
216
VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
@@ -1586,6 +1588,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
1588
"VLLM_NIXL_SIDE_CHANNEL_PORT": lambda: int(
1589
os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5600")
1590
),
1591
+ # Address the P2P KV-offload control socket binds to. Defaults to
1592
+ # ``localhost`` (loopback only); must be set to the node IP for
1593
+ # cross-host P2P so remote peers can reach the socket.
1594
+ "VLLM_P2P_SIDE_CHANNEL_HOST": lambda: os.getenv(
1595
+ "VLLM_P2P_SIDE_CHANNEL_HOST", "localhost"
1596
+ ),
1597
+ # Port the P2P KV-offload control socket binds to.
1598
+ "VLLM_P2P_SIDE_CHANNEL_PORT": lambda: int(
1599
+ os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710")
1600
+ ),
1601
# IP address used for the EC connector's ZMQ side channel
1602
# (producer ROUTER bind, consumer DEALER dial).
1603
"VLLM_EC_SIDE_CHANNEL_HOST": lambda: os.getenv(
vllm/v1/kv_offload/tiering/p2p/data/base.py
+11
-2
@@ -207,8 +207,17 @@ class DataTransport(ABC):
207
...
208
209
@abstractmethod
210
- def poll(self) -> PollResult:
211
- """Poll all inflight transfers for completion.
210
+ def poll(self, peer_id: str | None = None) -> PollResult:
211
+ """Poll inflight transfers for completion.
212
+
213
+ Args:
214
+ peer_id: If given, only poll (and drain) transfers submitted for
215
+ this peer_id — the value passed to ``write_blocks``. This is
216
+ required when a single transport is shared across multiple
217
+ peer sessions: ``poll()`` pops completed handles, so an
218
+ unscoped poll by one session would consume and discard the
219
+ completions of its siblings, starving them. ``None`` polls
220
+ every peer's transfers (used only for the shutdown drain).
221
222
Returns:
223
PollResult with lists of completed and failed transfer_ids.
vllm/v1/kv_offload/tiering/p2p/data/nixl.py
+46
-25
@@ -8,7 +8,7 @@ from __future__ import annotations
8
9
import itertools
10
from collections.abc import Iterable
11
-from typing import Any
11
+from typing import Any, NamedTuple
12
13
from vllm.distributed.nixl_utils import NixlWrapper as _NixlAgent
14
from vllm.distributed.nixl_utils import nixl_agent_config as _NixlAgentConfig
@@ -27,6 +27,17 @@ logger = init_logger(__name__)
27
_EMPTY_POLL_RESULT: PollResult = PollResult(done=(), failed=())
28
29
30
+class _Inflight(NamedTuple):
31
+ """A submitted-but-not-yet-drained transfer.
32
+
33
+ ``peer_id`` lets poll() scope to a single owning session, since the
34
+ transport is shared across all peer sessions of the engine.
35
+ """
36
+
37
+ peer_id: str
38
+ handle: object
39
+
40
+
41
class NixlTransport(DataTransport):
42
"""Manages a NIXL agent, memory registration, and block transfers.
43
@@ -37,14 +48,14 @@ class NixlTransport(DataTransport):
48
49
def __init__(
50
self,
40
- local_id: str,
51
+ agent_name: str,
52
view: memoryview,
53
config_fields: dict | None = None,
54
backends: list[str] | None = None,
55
num_threads: int = 4,
56
) -> None:
57
super().__init__(view, config_fields=config_fields)
47
- self._local_id = local_id
58
+ self._agent_name = agent_name
59
self._backends = list(backends) if backends else ["UCX"]
60
self._num_threads = num_threads
61
self._agent: Any = None
@@ -52,7 +63,8 @@ class NixlTransport(DataTransport):
63
self._local_dlist: Any = None
64
self._remote_dlists: dict[str, object] = {}
65
self._peer_nixl_names: dict[str, str] = {}
55
- self._inflight: dict[int, object] = {} # transfer_id → handle
66
+ # transfer_id → _Inflight(peer_id, handle).
67
+ self._inflight: dict[int, _Inflight] = {}
68
self._next_id = itertools.count()
69
70
self._init(view)
@@ -70,7 +82,7 @@ class NixlTransport(DataTransport):
82
cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True)
83
logger.info(
84
"NixlTransport %s: NIXL backends=%s",
73
- self._local_id,
85
+ self._agent_name,
86
self._backends,
87
)
88
else:
@@ -79,10 +91,10 @@ class NixlTransport(DataTransport):
91
)
92
logger.info(
93
"NixlTransport %s: NIXL backends=[UCX] num_threads=%d",
82
- self._local_id,
94
+ self._agent_name,
95
self._num_threads,
96
)
85
- self._agent = _NixlAgent(self._local_id, cfg)
97
+ self._agent = _NixlAgent(self._agent_name, cfg)
98
99
total_size = self._num_blocks * self._block_len
100
reg_descs = [(self._base_addr, total_size, 0, "")]
@@ -95,7 +107,7 @@ class NixlTransport(DataTransport):
107
xfer_dlist = self._agent.get_xfer_descs(block_tuples, mem_type="DRAM")
108
self._local_dlist = self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", xfer_dlist)
109
logger.info(
98
- "NixlTransport %s: registered %d blocks", self._local_id, self._num_blocks
110
+ "NixlTransport %s: registered %d blocks", self._agent_name, self._num_blocks
111
)
112
113
def get_agent_metadata(self) -> bytes:
@@ -152,14 +164,14 @@ class NixlTransport(DataTransport):
164
logger.warning(
165
"NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s "
166
"(known peers=%s)",
155
- self._local_id,
167
+ self._agent_name,
168
peer_id,
169
list(self._remote_dlists.keys()),
170
)
171
return None
172
logger.debug(
173
"NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d",
162
- self._local_id,
174
+ self._agent_name,
175
peer_id,
176
len(local_idxs),
177
)
@@ -172,11 +184,16 @@ class NixlTransport(DataTransport):
184
)
185
self._agent.transfer(handle)
186
transfer_id = next(self._next_id)
175
- self._inflight[transfer_id] = handle
187
+ self._inflight[transfer_id] = _Inflight(peer_id, handle)
188
return transfer_id
189
178
- def poll(self) -> PollResult:
179
- """Poll all inflight transfers.
190
+ def poll(self, peer_id: str | None = None) -> PollResult:
191
+ """Poll inflight transfers.
192
+
193
+ When *peer_id* is given, only transfers submitted for that peer_id are
194
+ checked and drained — the transport is shared across peer sessions, so
195
+ an unscoped poll by one session would consume and discard siblings'
196
+ completions. *peer_id* None polls every peer (shutdown drain only).
197
198
Returns PollResult(done=..., failed=...) with transfer IDs.
199
Completed handles are released automatically.
@@ -187,13 +204,15 @@ class NixlTransport(DataTransport):
204
done_ids: list[int] | None = None
205
failed_ids: list[int] | None = None
206
190
- for transfer_id, handle in self._inflight.items():
207
+ for transfer_id, entry in self._inflight.items():
208
+ if peer_id is not None and entry.peer_id != peer_id:
209
+ continue
210
try:
192
- state = self._agent.check_xfer_state(handle)
211
+ state = self._agent.check_xfer_state(entry.handle)
212
except Exception as exc:
213
logger.warning(
214
"NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s",
196
- self._local_id,
215
+ self._agent_name,
216
transfer_id,
217
exc,
218
)
@@ -212,9 +231,9 @@ class NixlTransport(DataTransport):
231
232
handles_to_release = []
233
for tid in done_ids or ():
215
- handles_to_release.append(self._inflight.pop(tid))
234
+ handles_to_release.append(self._inflight.pop(tid).handle)
235
for tid in failed_ids or ():
217
- handles_to_release.append(self._inflight.pop(tid))
236
+ handles_to_release.append(self._inflight.pop(tid).handle)
237
self._release_handles(handles_to_release)
238
239
return PollResult(
@@ -236,22 +255,24 @@ class NixlTransport(DataTransport):
255
"""
256
if mode == "immediate":
257
handles = [
239
- self._inflight.pop(tid) for tid in transfer_ids if tid in self._inflight
258
+ self._inflight.pop(tid).handle
259
+ for tid in transfer_ids
260
+ if tid in self._inflight
261
]
262
self._release_handles(handles)
263
return []
264
265
still_inflight: list[int] = []
266
for tid in transfer_ids:
246
- handle = self._inflight.get(tid)
247
- if handle is None:
267
+ entry = self._inflight.get(tid)
268
+ if entry is None:
269
continue
270
try:
250
- self._agent.release_xfer_handle(handle)
271
+ self._agent.release_xfer_handle(entry.handle)
272
except Exception as exc:
273
logger.debug(
274
"NixlTransport %s: cancel pending for transfer_id=%d: %s",
254
- self._local_id,
275
+ self._agent_name,
276
tid,
277
exc,
278
)
@@ -267,7 +288,7 @@ class NixlTransport(DataTransport):
288
def close(self) -> None:
289
if self._agent is None:
290
return
270
- self._release_handles(list(self._inflight.values()))
291
+ self._release_handles([entry.handle for entry in self._inflight.values()])
292
self._inflight.clear()
293
for peer_id in list(self._remote_dlists):
294
self.remove_remote_peer(peer_id)
@@ -292,6 +313,6 @@ class NixlTransport(DataTransport):
313
except Exception as exc:
314
logger.warning(
315
"NixlTransport %s: release_xfer_handle failed: %s",
295
- self._local_id,
316
+ self._agent_name,
317
exc,
318
)
vllm/v1/kv_offload/tiering/p2p/manager.py
+37
-7
@@ -9,12 +9,14 @@ Owns transports and a single bidirectional P2PSession per remote peer.
9
from __future__ import annotations
10
11
import time
12
+import uuid
13
from collections.abc import Iterable, Sequence
14
from dataclasses import dataclass, field
15
from typing import TYPE_CHECKING, Any
16
17
from typing_extensions import override
18
19
+import vllm.envs as envs
20
from vllm.logger import init_logger
21
from vllm.v1.kv_offload.base import (
22
LookupResult,
@@ -111,8 +113,8 @@ class P2PSecondaryTierManager(SecondaryTierManager):
113
offloading_spec: OffloadingSpec,
114
primary_kv_view: memoryview,
115
tier_type: str = "p2p",
114
- host: str = "0.0.0.0",
115
- port: int = 7777,
116
+ host: str | None = None,
117
+ port: int | None = None,
118
backends: list[str] | None = None,
119
num_threads: int = 4,
120
**kwargs: Any,
@@ -130,9 +132,19 @@ class P2PSecondaryTierManager(SecondaryTierManager):
132
primary_kv_view: Memoryview over the CPU primary tier; the
133
NIXL agent registers this region for RDMA transfers.
134
tier_type: Tier identifier (defaults to ``"p2p"``).
133
- host: Address the ZMQ control socket binds to.
134
- port: Port for the ZMQ control socket. Must be reachable
135
- from peers.
135
+ host: Address the ZMQ control socket binds to, used verbatim
136
+ as both the bind address and the identity peers dial back
137
+ (mirrors the NIXL connector's ``VLLM_NIXL_SIDE_CHANNEL_HOST``;
138
+ no auto-detection). Defaults to
139
+ ``VLLM_P2P_SIDE_CHANNEL_HOST`` (``localhost``) when not set;
140
+ must be set to the node's routable IP for cross-host P2P so
141
+ remote peers can reach the socket.
142
+ port: Base port for the ZMQ control socket. Must be
143
+ reachable from peers. Defaults to
144
+ ``VLLM_P2P_SIDE_CHANNEL_PORT`` (``5710``) when not set.
145
+ The bound port is ``base + data_parallel_index`` so each
146
+ DP replica gets a distinct port (one socket per replica,
147
+ like NIXL); for DP=1 the offset is 0.
148
backends: NIXL transport backends (e.g. ``["UCX"]``,
149
``["MOONCAKE"]``, ``["LIBFABRIC"]``). Defaults to
150
``["UCX"]``. When any non-UCX backend is requested, the
@@ -145,8 +157,26 @@ class P2PSecondaryTierManager(SecondaryTierManager):
157
**kwargs: Reserved for future tier-specific options.
158
"""
159
super().__init__(offloading_spec, primary_kv_view, tier_type)
148
- port = int(port)
160
+ if host is None:
161
+ host = envs.VLLM_P2P_SIDE_CHANNEL_HOST
162
+ if port is None:
163
+ port = envs.VLLM_P2P_SIDE_CHANNEL_PORT
164
+ # One control socket per DP replica: offset the base by the global
165
+ # data-parallel index so replicas on a host don't collide (mirrors
166
+ # NIXL). For DP=1 the index is 0, leaving the base port unchanged.
167
+ dp_index = offloading_spec.vllm_config.parallel_config.data_parallel_index
168
+ port = int(port) + dp_index
169
+ # Two decoupled identities:
170
+ # _local_id (``host:port``): the ZMQ control identity that peers
171
+ # dial back, used verbatim (the socket binds this host/port and
172
+ # the address is parsed back into host:port by the remote).
173
+ # _nixl_agent_name (uuid4): the NIXL agent name. It is never dialed
174
+ # — it travels opaquely inside the agent metadata blob — so it
175
+ # only needs to be globally unique. A per-process uuid guarantees
176
+ # that even for peers sharing a host:port (mirrors the NIXL
177
+ # connector; avoids the "remote agent name equals local" reject).
178
self._local_id = f"{host}:{port}"
179
+ self._nixl_agent_name = str(uuid.uuid4())
180
181
config_fields = FileMapper.from_offloading_spec(
182
root_dir="",
@@ -155,7 +185,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
185
parallel_agnostic=True,
186
).get_run_config()
187
self._data: DataTransport = NixlTransport(
158
- self._local_id,
188
+ self._nixl_agent_name,
189
primary_kv_view,
190
config_fields=config_fields,
191
backends=backends,
vllm/v1/kv_offload/tiering/p2p/session/client.py
+8
@@ -139,6 +139,14 @@ class ClientRole:
139
"""Handle an AbortAckMsg from the peer."""
140
req = self._inbound.pop(kv_request_id, None)
141
if req is not None:
142
+ logger.warning(
143
+ "P2PSession %s: load request %s (job_id=%d) timed out; "
144
+ "load job completed with failure. If this recurs, ensure "
145
+ "PYTHONHASHSEED is set to the same value on all nodes.",
146
+ self._peer_id,
147
+ kv_request_id,
148
+ req.job_id,
149
+ )
150
self._completed_loads.append(
151
LoadResult(
152
job_id=req.job_id,
vllm/v1/kv_offload/tiering/p2p/session/server.py
+5
-1
@@ -282,7 +282,11 @@ class ServerRole:
282
results.extend(self._pending_store_results)
283
self._pending_store_results.clear()
284
285
- poll_result = self._transport.poll()
285
+ # Scope the poll to this peer: the transport is shared across all peer
286
+ # sessions of the engine, and poll() drains completed handles. An
287
+ # unscoped poll here would consume sibling sessions' completions and
288
+ # report them as "unknown transfer_id", starving those sessions.
289
+ poll_result = self._transport.poll(self._peer_id)
290
291
for tid in poll_result.done:
292
xfer = self._inflight_pop(tid)