@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 90531daf6d

[Core][Frontend] Bind KV-event publishers at port 0 and expose the bound endpoints (#55844)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: mukeshbhandarkar <shell00035@gmail.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>

hangangqiang committed Sep 19, 2026 at 09:11 UTC 90531daf6dde1c961d326d43c493477feca510c7

11 files changed +282 -14

rust/src/server/src/routes.rs

@@ -6,6 +6,7 @@ mod cache;

6 mod collective_rpc;

7 mod health;

8 mod inference;

9 +mod kv_event_sources;

10 mod load;

11 mod lora;

12 mod metrics;

@@ -91,6 +92,7 @@ fn build_router_with_options(

92 .route("/metrics", get(metrics::scrape))

93 .route("/load", get(load::load))

94 .route("/version", get(version::version))

95 + .route("/kv_event_sources", get(kv_event_sources::kv_event_sources))

96 // OpenAI-compatible endpoints

97 .route("/v1/models", get(openai::list_models))

98 .route("/v1/completions", post(openai::completions))

rust/src/server/src/routes/kv_event_sources.rs new

+25

@@ -0,0 +1,25 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +use std::collections::BTreeMap;

5 +use std::sync::Arc;

6 +

7 +use axum::Json;

8 +use axum::extract::State;

9 +use vllm_engine_core_client::protocol::handshake::KvEventsConfig;

10 +

11 +use crate::state::AppState;

12 +

13 +/// KV-event publisher config of each data-parallel rank, keyed by rank.

14 +pub async fn kv_event_sources(

15 + State(state): State<Arc<AppState>>,

16 +) -> Json<BTreeMap<u32, KvEventsConfig>> {

17 + Json(

18 + state

19 + .engine_core_client()

20 + .ready_responses()

21 + .into_iter()

22 + .filter_map(|ready| Some((ready.data_parallel_rank, ready.kv_events_config.clone()?)))

23 + .collect(),

24 + )

25 +}

rust/src/server/src/routes/tests.rs

+50

@@ -2054,6 +2054,56 @@ async fn version_returns_engine_vllm_version() {

2054 );

2055 }

2056

2057 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

2058 +#[serial]

2059 +async fn kv_event_sources_reports_each_ready_engine() {

2060 + let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {

2061 + data_parallel_rank: 1,

2062 + kv_events_config: Some(

2063 + vllm_engine_core_client::protocol::handshake::KvEventsConfig {

2064 + enable_kv_cache_events: true,

2065 + publisher: "zmq".to_string(),

2066 + endpoint: "tcp://10.0.0.2:41233".to_string(),

2067 + replay_endpoint: None,

2068 + buffer_steps: 10_000,

2069 + hwm: 100_000,

2070 + max_queue_size: 100_000,

2071 + topic: "kv".to_string(),

2072 + },

2073 + ),

2074 + ..default_ready_response()

2075 + };

2076 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;

2077 + let response = app

2078 + .call(

2079 + Request::builder()

2080 + .uri("/kv_event_sources")

2081 + .body(Body::empty())

2082 + .expect("build request"),

2083 + )

2084 + .await

2085 + .expect("call app");

2086 +

2087 + assert_eq!(response.status(), StatusCode::OK);

2088 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

2089 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");

2090 + assert_eq!(

2091 + json,

2092 + json!({

2093 + "1": {

2094 + "enable_kv_cache_events": true,

2095 + "publisher": "zmq",

2096 + "endpoint": "tcp://10.0.0.2:41233",

2097 + "replay_endpoint": null,

2098 + "buffer_steps": 10000,

2099 + "hwm": 100000,

2100 + "max_queue_size": 100000,

2101 + "topic": "kv",

2102 + }

2103 + })

2104 + );

2105 +}

2106 +

2107 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

2108 #[serial]

2109 async fn server_info_endpoint_is_dev_mode_only() {

tests/distributed/conftest.py

@@ -91,8 +91,11 @@ class MockSubscriber:

91 if isinstance(replay_endpoints, str):

92 replay_endpoints = [replay_endpoints]

94 + ipv6 = any("[" in ep for ep in pub_endpoints + (replay_endpoints or []))

95 +

96 # Set up subscriber socket - connect to all endpoints

97 self.sub = self.ctx.socket(zmq.SUB)

98 + self.sub.setsockopt(zmq.IPV6, ipv6)

99 self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode("utf-8"))

100 for endpoint in pub_endpoints:

101 self.sub.connect(endpoint)

@@ -103,6 +106,7 @@ class MockSubscriber:

106 if replay_endpoints:

107 for replay_endpoint in replay_endpoints:

108 replay = self.ctx.socket(zmq.DEALER)

109 + replay.setsockopt(zmq.IPV6, ipv6)

110 replay.connect(replay_endpoint)

111 self.replay_sockets.append(replay)

112

tests/distributed/test_events.py

+74

@@ -332,3 +332,77 @@ def test_event_publisher_factory(random_port):

332 publisher = EventPublisherFactory.create(config, DP_RANK)

333 assert isinstance(publisher, ZmqEventPublisher)

334 publisher.shutdown()

335 +

336 +

337 +@pytest.mark.parametrize(

338 + ("endpoint", "expected"),

339 + [("tcp://*:5557", "tcp://*:5558"), ("tcp://*:0", "tcp://*:0")],

340 +)

341 +def test_offset_endpoint_port_keeps_port_zero(endpoint, expected):

342 + """Port 0 is not offset: every rank asks the OS for its own port."""

343 + from vllm.distributed.kv_events import ZmqEventPublisher

344 +

345 + assert ZmqEventPublisher.offset_endpoint_port(endpoint, 1) == expected

346 +

347 +

348 +@pytest.mark.parametrize("node_ip", ["127.0.0.1", "::1"])

349 +def test_ephemeral_publishers_report_bound_endpoints(monkeypatch, node_ip):

350 + """Two ranks binding tcp://*:0 get distinct OS-assigned ports on this

351 + node's IP and report endpoints a subscriber can use for events and replay."""

352 + from vllm.config.kv_events import KVEventsConfig

353 + from vllm.distributed import kv_events

354 + from vllm.utils.network_utils import split_zmq_path

355 +

356 + from .conftest import MockSubscriber

357 +

358 + monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip)

359 + config = KVEventsConfig(

360 + enable_kv_cache_events=True,

361 + publisher="zmq",

362 + endpoint="tcp://*:0",

363 + replay_endpoint="tcp://*:0",

364 + topic="kv",

365 + )

366 + pubs = [EventPublisherFactory.create(config, rank) for rank in range(2)]

367 + try:

368 + resolved = [pub.get_publisher_config() for pub in pubs]

369 + endpoints = [e for c in resolved for e in (c.endpoint, c.replay_endpoint)]

370 + hosts_ports = {split_zmq_path(e)[1:] for e in endpoints}

371 + assert {h for h, _ in hosts_ports} == {node_ip}

372 + assert len(hosts_ports) == 4 and "0" not in {p for _, p in hosts_ports}

373 + assert config.endpoint == "tcp://*:0"

374 +

375 + sub = MockSubscriber(endpoints[::2], endpoints[1::2], topic="kv")

376 + try:

377 + # PUB/SUB drops messages sent before the subscription lands, so

378 + # keep publishing until both ranks are heard.

379 + ranks: set[int] = set()

380 + deadline = time.time() + 10

381 + while ranks != {0, 1} and time.time() < deadline:

382 + for pub in pubs:

383 + pub.publish(create_test_events(1))

384 + while (result := sub.receive_one(timeout=100)) is not None:

385 + ranks.add(result[1].data_parallel_rank)

386 + assert ranks == {0, 1}

387 + sub.request_replay(0, socket_idx=1)

388 + assert sub.receive_replay(socket_idx=1)

389 + finally:

390 + sub.close()

391 + finally:

392 + for pub in pubs:

393 + pub.shutdown()

394 +

395 +

396 +@pytest.mark.parametrize("node_ip", ["0.0.0.0", "::"])

397 +def test_ephemeral_bind_needs_a_node_ip(monkeypatch, node_ip):

398 + """tcp://*:0 is advertised as this node's IP; with no usable IP the

399 + publisher refuses instead of reporting a wildcard."""

400 + from vllm.config.kv_events import KVEventsConfig

401 + from vllm.distributed import kv_events

402 +

403 + monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip)

404 + config = KVEventsConfig(

405 + enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://*:0"

406 + )

407 + with pytest.raises(ValueError, match="VLLM_HOST_IP"):

408 + EventPublisherFactory.create(config, DP_RANK)

tests/v1/engine/test_engine_core_client.py

+55

@@ -214,6 +214,7 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM

214 client.core_engines = [bytes([i, 0]) for i in range(num_engines)]

215 client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]

216 client.eng_start_index = 0

217 + client._kv_event_sources = {}

218 return client

219

220

@@ -1490,3 +1491,57 @@ def test_ready_response_queries_executor_with_weight_transfer():

1491

1492 executor.supports_draft_weight_updates.assert_called_once()

1493 assert response.supports_draft_weight_updates is True

1494 +

1495 +

1496 +def test_apply_ready_response_retains_kv_event_sources():

1497 + """Each engine's resolved publisher config is kept by DP rank; engines

1498 + without KV events contribute nothing."""

1499 + import msgspec

1500 +

1501 + from vllm.config.kv_events import KVEventsConfig

1502 +

1503 + client = object.__new__(MPClient)

1504 + client._effective_attention_block_sizes = set()

1505 + client._kv_event_sources = {}

1506 + client.vllm_config = SimpleNamespace(

1507 + cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),

1508 + model_config=SimpleNamespace(max_model_len=8192),

1509 + )

1510 +

1511 + def ready(rank: int, config: KVEventsConfig | None) -> bytes:

1512 + return msgspec.msgpack.encode(

1513 + EngineCoreReadyResponse(

1514 + max_model_len=8192,

1515 + num_gpu_blocks=100,

1516 + block_size=16,

1517 + dp_stats_address=None,

1518 + dtype="bfloat16",

1519 + vllm_version="test",

1520 + world_size=2,

1521 + data_parallel_size=2,

1522 + tensor_parallel_size=1,

1523 + pipeline_parallel_size=1,

1524 + decode_context_parallel_size=1,

1525 + data_parallel_rank=rank,

1526 + max_num_seqs=256,

1527 + max_num_batched_tokens=8192,

1528 + instance_id="test-instance",

1529 + supports_lora=False,

1530 + max_loras=0,

1531 + kv_events_config=config,

1532 + )

1533 + )

1534 +

1535 + rank1 = KVEventsConfig(

1536 + enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:41233"

1537 + )

1538 + client._apply_ready_response(ready(1, rank1))

1539 + client._apply_ready_response(ready(0, None))

1540 + assert client.get_kv_event_sources() == {1: rank1}

1541 +

1542 + # A rank that comes back (elastic re-add) reports its new port.

1543 + rank1_again = KVEventsConfig(

1544 + enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:47001"

1545 + )

1546 + client._apply_ready_response(ready(1, rank1_again))

1547 + assert client.get_kv_event_sources() == {1: rank1_again}

vllm/distributed/kv_events.py

+40 -14

@@ -17,6 +17,12 @@ import zmq

18 from vllm.config.kv_events import KVEventsConfig

19 from vllm.logger import init_logger

20 +from vllm.utils.network_utils import (

21 + get_ip,

22 + get_tcp_uri,

23 + is_valid_ipv6_address,

24 + split_zmq_path,

25 +)

26 from vllm.v1.core.kv_cache_utils import ExternalBlockHash

28 logger = init_logger(__name__)

@@ -349,6 +355,8 @@ class ZmqEventPublisher(EventPublisher):

355 replay_endpoint, self._dp_rank

356 )

357 assert self._endpoint is not None

358 + self._hwm = hwm

359 + self._socket_setup()

360 self._publisher_config = KVEventsConfig(

361 enable_kv_cache_events=True,

362 publisher="zmq",

@@ -359,8 +367,6 @@ class ZmqEventPublisher(EventPublisher):

367 max_queue_size=max_queue_size,

368 topic=topic,

369 )

362 - self._hwm = hwm

363 - self._socket_setup()

370

371 # Payload

372 self._seq_gen = count()

@@ -423,16 +429,10 @@ class ZmqEventPublisher(EventPublisher):

429 if self._pub is None:

430 self._pub = self._ctx.socket(zmq.PUB)

431 self._pub.set_hwm(self._hwm)

426 - # Heuristic: bind if wildcard / * present, else connect.

427 - # bind stable, connect volatile convention

428 - if self._endpoint is not None and (

429 - "*" in self._endpoint

430 - or "::" in self._endpoint

431 - or self._endpoint.startswith("ipc://")

432 - or self._endpoint.startswith("inproc://")

433 - ):

434 - self._pub.bind(self._endpoint)

435 - elif self._endpoint is not None:

432 + assert self._endpoint is not None

433 + if self._is_bind_endpoint(self._endpoint):

434 + self._endpoint = self._bind(self._pub, self._endpoint)

435 + else:

436 self._pub.connect(self._endpoint)

437

438 # Set up replay socket: use ROUTER

@@ -441,7 +441,31 @@ class ZmqEventPublisher(EventPublisher):

441 # 3) works in our non‑blocking poll loop alongside PUB

442 if self._replay_endpoint is not None:

443 self._replay = self._ctx.socket(zmq.ROUTER)

444 - self._replay.bind(self._replay_endpoint)

444 + self._replay_endpoint = self._bind(self._replay, self._replay_endpoint)

445 +

446 + @staticmethod

447 + def _is_bind_endpoint(endpoint: str) -> bool:

448 + """Bind wildcard and ipc/inproc endpoints; connect to the rest."""

449 + return (

450 + "*" in endpoint

451 + or "::" in endpoint

452 + or endpoint.startswith(("ipc://", "inproc://"))

453 + )

454 +

455 + @staticmethod

456 + def _bind(sock: zmq.Socket, endpoint: str) -> str:

457 + """Bind; tcp://*:0 binds this node's IP and returns the bound endpoint."""

458 + if endpoint != "tcp://*:0":

459 + sock.bind(endpoint)

460 + return endpoint

461 + ip = get_ip()

462 + if ip in ("0.0.0.0", "::"):

463 + raise ValueError("tcp://*:0 needs a node IP; set VLLM_HOST_IP")

464 + if is_valid_ipv6_address(ip):

465 + sock.setsockopt(zmq.IPV6, 1)

466 + sock.bind(get_tcp_uri(ip, 0))

467 + _, _, port = split_zmq_path(sock.getsockopt_string(zmq.LAST_ENDPOINT))

468 + return get_tcp_uri(ip, int(port))

469

470 def _publisher_thread(self) -> None:

471 """Background thread that processes the event queue."""

@@ -514,7 +538,7 @@ class ZmqEventPublisher(EventPublisher):

538

539 Returns:

540 The endpoint with the port offset by data_parallel_rank

517 - or suffix appended

541 + or suffix appended; a tcp port of 0 is left for the OS.

542

543 """

544 # Do nothing if input is None or data_parallel_rank is 0

@@ -529,6 +553,8 @@ class ZmqEventPublisher(EventPublisher):

553 last_colon_idx = endpoint.rfind(":")

554 base_addr = endpoint[:last_colon_idx]

555 base_port = int(endpoint[last_colon_idx + 1 :])

556 + if base_port == 0:

557 + return endpoint

558 new_port = base_port + data_parallel_rank

559 return f"{base_addr}:{new_port}"

560 return endpoint

vllm/engine/protocol.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass

7 from typing import TYPE_CHECKING, Any

9 from vllm.config import ModelConfig, VllmConfig

10 +from vllm.config.kv_events import KVEventsConfig

11 from vllm.distributed.weight_transfer.base import (

12 WeightTransferInitRequest,

13 WeightTransferUpdateRequest,

@@ -62,6 +63,10 @@ class EngineClient(ABC):

63 @abstractmethod

64 def dead_error(self) -> BaseException: ...

66 + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:

67 + """KV-event publisher config of each engine, keyed by DP rank."""

68 + return {}

69 +

70 def check_admission( # noqa: B027

71 self, n: int = 1, request_id: str | None = None

72 ) -> None:

vllm/entrypoints/serve/instrumentator/basic.py

@@ -54,3 +54,10 @@ async def get_server_load_metrics(request: Request):

54 async def show_version():

55 ver = {"version": VLLM_VERSION}

56 return JSONResponse(content=ver)

57 +

58 +

59 +@router.get("/kv_event_sources")

60 +async def kv_event_sources(raw_request: Request):

61 + """KV-event publisher config per DP rank; empty on render-only servers."""

62 + client: EngineClient | None = raw_request.app.state.engine_client

63 + return client.get_kv_event_sources() if client is not None else {}

vllm/v1/engine/async_llm.py

@@ -12,6 +12,7 @@ from typing import Any

12 import vllm.envs as envs

13 from vllm import TokensPrompt

14 from vllm.config import VllmConfig

15 +from vllm.config.kv_events import KVEventsConfig

16 from vllm.distributed.weight_transfer.base import (

17 WeightTransferInitRequest,

18 WeightTransferUpdateRequest,

@@ -1282,3 +1283,6 @@ class AsyncLLM(EngineClient):

1283 async def get_weight_version(self) -> str:

1284 """Return the latest committed weight version."""

1285 return await self.engine_core.get_weight_version_async()

1286 +

1287 + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:

1288 + return self.engine_core.get_kv_event_sources()

vllm/v1/engine/core_client.py

+16

@@ -23,6 +23,7 @@ import zmq.asyncio

24 from vllm import envs

25 from vllm.config import VllmConfig

26 +from vllm.config.kv_events import KVEventsConfig

27 from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S

28 from vllm.logger import init_logger

29 from vllm.lora.request import LoRARequest

@@ -86,6 +87,10 @@ class EngineCoreClient(ABC):

87 * AsyncMPClient: ZMQ + background proc EngineCore w/ asyncio (for AsyncLLM)

88 """

90 + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:

91 + """KV-event publisher config of each ready engine, keyed by DP rank."""

92 + return {}

93 +

94 @staticmethod

95 def make_client(

96 multiprocess_mode: bool,

@@ -572,6 +577,7 @@ class MPClient(EngineCoreClient):

577 self.vllm_config = vllm_config

578 self._renderer: BaseRenderer | None = renderer

579 self._effective_attention_block_sizes: set[int | None] = set()

580 + self._kv_event_sources: dict[int, KVEventsConfig] = {}

581

582 # ZMQ setup.

583 sync_ctx = zmq.Context(io_threads=2)

@@ -860,6 +866,14 @@ class MPClient(EngineCoreClient):

866 else:

867 assert response.dp_stats_address == self.stats_update_address

868

869 + if response.kv_events_config is not None:

870 + self._kv_event_sources[response.data_parallel_rank] = (

871 + response.kv_events_config

872 + )

873 +

874 + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:

875 + return dict(self._kv_event_sources)

876 +

877

878 def _process_utility_output(

879 output: UtilityOutput, utility_results: dict[int, AnyFuture]

@@ -1944,6 +1958,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1958 self.lb_engines = self.lb_engines[:new_data_parallel_size]

1959 # Pending coordinator snapshots must use the surviving ranks too.

1960 self.engine_ranks_managed = self.engine_ranks_managed[:new_data_parallel_size]

1961 + for rank in range(new_data_parallel_size, cur_data_parallel_size):

1962 + self._kv_event_sources.pop(rank, None)

1963 removed_dp_size = cur_data_parallel_size - new_data_parallel_size

1964 pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size

1965 pause_futures = [