From 48fc2e2797ec3f35605c3ff399bbaab7938b0459 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 29 Jul 2026 13:46:40 -0700 Subject: [PATCH] feat(grpc): add KV event source discovery (#50033) Signed-off-by: Connor Carpenter Co-authored-by: Nick Hill --- rust/proto/control.proto | 21 +++++++++++ .../src/engine-core-client/src/mock_engine.rs | 1 + .../src/protocol/handshake.rs | 16 +++++++++ .../engine-core-client/src/tests/client.rs | 17 ++++++++- .../src/tests/python_compat.py | 23 ++++++++++++ rust/src/server/src/grpc/control.rs | 29 +++++++++++++++ rust/src/server/src/grpc/tests.rs | 36 +++++++++++++++++++ tests/distributed/test_events.py | 11 ++++-- vllm/distributed/kv_events.py | 18 ++++++++++ vllm/v1/core/sched/interface.py | 4 +++ vllm/v1/core/sched/scheduler.py | 5 ++- vllm/v1/engine/__init__.py | 2 ++ vllm/v1/engine/core.py | 1 + 13 files changed, 179 insertions(+), 5 deletions(-) diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 7b858cb3e2..0e25aea264 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -8,6 +8,7 @@ service Control { rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} rpc Abort (AbortRequest) returns (AbortResponse) {} + rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} } message GetServerInfoRequest {} @@ -51,3 +52,23 @@ message AbortRequest { } message AbortResponse {} + +// ====================================================================================== +// KV discovery +// ====================================================================================== + +message GetKvEventSourcesRequest {} +message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } + +message KvEventSource { + string transport = 1; + string endpoint = 2; + string topic = 3; + string replay_endpoint = 4; + optional uint32 data_parallel_rank = 5; + string encoding = 6; + uint32 schema_version = 7; + uint32 buffer_steps = 8; + uint32 hwm = 9; + uint32 max_queue_size = 10; +} diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 5540ae8bc1..ee066635ff 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -66,6 +66,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { instance_id: "test-instance".to_string(), kv_cache_size_tokens: None, kv_cache_max_concurrency: None, + kv_events_config: None, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index 125a82b39d..a4e1cebde6 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -24,6 +24,19 @@ pub struct ReadyMessage { pub parallel_config_hash: Option, } +/// KV-event publisher configuration reported by EngineCore. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KvEventsConfig { + pub enable_kv_cache_events: bool, + pub publisher: String, + pub endpoint: String, + pub replay_endpoint: Option, + pub buffer_steps: u32, + pub hwm: u32, + pub max_queue_size: u32, + pub topic: String, +} + /// Post-initialization configuration sent from each engine on the input socket /// registration message, after the handshake completes. /// @@ -71,6 +84,9 @@ pub struct EngineCoreReadyResponse { pub kv_cache_size_tokens: Option, /// Maximum achievable request concurrency given the KV cache, if reported. pub kv_cache_max_concurrency: Option, + /// KV-event publisher configuration, if configured. + #[serde(default)] + pub kv_events_config: Option, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 7b3e688d86..11d08db646 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -19,7 +19,7 @@ use zeromq::prelude::{Socket, SocketRecv, SocketSend}; use zeromq::util::PeerIdentity; use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage}; -use crate::protocol::handshake::{HandshakeInitMessage, ReadyMessage}; +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; use crate::protocol::logprobs::MaybeWireLogprobs; use crate::protocol::multimodal::{ MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange, @@ -2677,6 +2677,21 @@ fn python_msgpack_fixtures_match_rust_encoding() { rust_ready_keys, python_ready_keys, "EngineCoreReadyResponse drifted from the Python dataclass", ); + + let ready_response: EngineCoreReadyResponse = + rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode"); + assert!(kv_events_config.enable_kv_cache_events); + assert_eq!(kv_events_config.publisher, "zmq"); + assert_eq!(kv_events_config.endpoint, "tcp://127.0.0.1:5557"); + assert_eq!( + kv_events_config.replay_endpoint.as_deref(), + Some("tcp://127.0.0.1:5558") + ); + assert_eq!(kv_events_config.topic, "kv"); + assert_eq!(kv_events_config.buffer_steps, 10_000); + assert_eq!(kv_events_config.hwm, 100_000); + assert_eq!(kv_events_config.max_queue_size, 100_000); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index ed16f2bd8b..c68e2b6b47 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -353,6 +353,18 @@ multipart_prompt_logprobs = engine_outputs_wire( ) +@dataclass +class KVEventsConfig: + enable_kv_cache_events: bool + publisher: str + endpoint: str + replay_endpoint: str | None + buffer_steps: int + hwm: int + max_queue_size: int + topic: str + + @dataclass class EngineCoreReadyResponse: max_model_len: int @@ -372,6 +384,7 @@ class EngineCoreReadyResponse: instance_id: str kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None + kv_events_config: KVEventsConfig | None = None ready_response = EngineCoreReadyResponse( @@ -390,6 +403,16 @@ ready_response = EngineCoreReadyResponse( max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + kv_events_config=KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint="tcp://127.0.0.1:5557", + replay_endpoint="tcp://127.0.0.1:5558", + buffer_steps=10_000, + hwm=100_000, + max_queue_size=100_000, + topic="kv", + ), ) print(msgspec.msgpack.encode(request).hex()) diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index e989e6abae..8972098222 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -103,4 +103,33 @@ impl pb::control_server::Control for ControlServiceImpl { .map_err(|error| Status::internal(error.to_report_string()))?; Ok(Response::new(pb::AbortResponse {})) } + + async fn get_kv_event_sources( + &self, + _request: Request, + ) -> Result, Status> { + let client = self.state.engine_core_client(); + let sources = client.ready_responses().into_iter().filter_map(kv_event_source).collect(); + Ok(Response::new(pb::GetKvEventSourcesResponse { sources })) + } +} + +pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { + let config = response.kv_events_config.as_ref()?; + if !config.enable_kv_cache_events || config.publisher != "zmq" { + return None; + } + + Some(pb::KvEventSource { + transport: "zmq".to_string(), + endpoint: config.endpoint.clone(), + topic: config.topic.clone(), + replay_endpoint: config.replay_endpoint.clone().unwrap_or_default(), + data_parallel_rank: Some(response.data_parallel_rank), + encoding: "msgpack".to_string(), + schema_version: 1, + buffer_steps: config.buffer_steps, + hwm: config.hwm, + max_queue_size: config.max_queue_size, + }) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 808f342275..8265fefa05 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -29,6 +29,7 @@ use vllm_engine_core_client::mock_engine::{ DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, default_ready_response, }; +use vllm_engine_core_client::protocol::handshake::KvEventsConfig; use vllm_engine_core_client::protocol::output::{ EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, }; @@ -44,6 +45,7 @@ use vllm_tokenizer::test_utils::TestTokenizer; use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; +use super::control::kv_event_source; use super::pb::control_client::ControlClient; use super::pb::inference_client::InferenceClient; use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb}; @@ -1319,6 +1321,40 @@ async fn control_aggregates_multi_engine_capacity() { drop(engine_tasks); } + +#[test] +fn kv_event_source_filters_and_exposes_zmq_publisher() { + let mut ready = default_ready_response(); + ready.data_parallel_rank = 2; + ready.kv_events_config = Some(KvEventsConfig { + enable_kv_cache_events: false, + publisher: "null".to_string(), + endpoint: "tcp://*:5559".to_string(), + replay_endpoint: Some("tcp://*:5560".to_string()), + buffer_steps: 10_000, + hwm: 100_000, + max_queue_size: 100_000, + topic: "kv".to_string(), + }); + + assert!(kv_event_source(&ready).is_none()); + + let config = ready.kv_events_config.as_mut().unwrap(); + config.enable_kv_cache_events = true; + config.publisher = "zmq".to_string(); + let source = kv_event_source(&ready).expect("configured ZMQ event source"); + assert_eq!(source.transport, "zmq"); + assert_eq!(source.endpoint, "tcp://*:5559"); + assert_eq!(source.topic, "kv"); + assert_eq!(source.replay_endpoint, "tcp://*:5560"); + assert_eq!(source.data_parallel_rank, Some(2)); + assert_eq!(source.encoding, "msgpack"); + assert_eq!(source.schema_version, 1); + assert_eq!(source.buffer_steps, 10_000); + assert_eq!(source.hwm, 100_000); + assert_eq!(source.max_queue_size, 100_000); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() { diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index 9b5601ad1d..894f3f7f8b 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -284,7 +284,7 @@ def test_data_parallel_rank_tagging(publisher_config): sub_1.close() -def test_event_publisher_factory(): +def test_event_publisher_factory(random_port): """Test event publisher factory creation behavior under different configurations""" from vllm.config.kv_events import KVEventsConfig from vllm.distributed.kv_events import ZmqEventPublisher @@ -308,10 +308,15 @@ def test_event_publisher_factory(): config = KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", - endpoint="inproc://test-factory-true", + endpoint=f"tcp://*:{random_port}", + replay_endpoint=f"tcp://*:{random_port + 100}", ) - publisher = EventPublisherFactory.create(config, DP_RANK) + publisher = EventPublisherFactory.create(config, DP_RANK + 1) assert isinstance(publisher, ZmqEventPublisher) + resolved_config = publisher.get_publisher_config() + assert resolved_config.endpoint == f"tcp://*:{random_port + 1}" + assert resolved_config.replay_endpoint == f"tcp://*:{random_port + 101}" + assert config.endpoint == f"tcp://*:{random_port}" publisher.shutdown() # test unknown publisher diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index 88cb3541f2..cd222567a3 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -272,6 +272,10 @@ class EventPublisher(ABC): def shutdown(self) -> None: """Shutdown the publisher.""" + def get_publisher_config(self) -> KVEventsConfig | None: + """Return the publisher's resolved runtime configuration.""" + return None + class NullEventPublisher(EventPublisher): """No-op implementation (default when disabled).""" @@ -335,6 +339,17 @@ class ZmqEventPublisher(EventPublisher): self._replay_endpoint = self.offset_endpoint_port( replay_endpoint, self._dp_rank ) + assert self._endpoint is not None + self._publisher_config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=self._endpoint, + replay_endpoint=self._replay_endpoint, + buffer_steps=buffer_steps, + hwm=hwm, + max_queue_size=max_queue_size, + topic=topic, + ) self._hwm = hwm self._socket_setup() @@ -351,6 +366,9 @@ class ZmqEventPublisher(EventPublisher): ) self._thread.start() + def get_publisher_config(self) -> KVEventsConfig: + return self._publisher_config + def publish(self, events: EventBatch) -> None: if not self._running: raise RuntimeError("Publisher is closed") diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index 562946fb72..c15296dd05 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -9,6 +9,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.config.kv_events import KVEventsConfig from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -256,3 +257,6 @@ class SchedulerInterface(ABC): def get_ec_connector(self) -> "ECConnectorBase | None": return None + + def get_kv_event_publisher_config(self) -> "KVEventsConfig | None": + return None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 1802c48111..7ba41e05e6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -8,7 +8,7 @@ from dataclasses import replace from typing import Any from vllm.compilation.cuda_graph import CUDAGraphStat -from vllm.config import VllmConfig +from vllm.config import KVEventsConfig, VllmConfig from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorBase, ECConnectorMetadata, @@ -2554,6 +2554,9 @@ class Scheduler(SchedulerInterface): def get_ec_connector(self) -> ECConnectorBase | None: return self.ec_connector + def get_kv_event_publisher_config(self) -> KVEventsConfig | None: + return self.kv_event_publisher.get_publisher_config() + def _connector_finished( self, request: Request ) -> tuple[bool, dict[str, Any] | None]: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 83033ecf81..5142820802 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -11,6 +11,7 @@ import msgspec import numpy as np import torch +from vllm.config.kv_events import KVEventsConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.pooling_params import PoolingParams @@ -90,6 +91,7 @@ class EngineCoreReadyResponse: # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None + kv_events_config: KVEventsConfig | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 6613527300..11c78c7e74 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1639,6 +1639,7 @@ class EngineCoreProc(EngineCore): max_num_seqs=scheduler_config.max_num_seqs, max_num_batched_tokens=scheduler_config.max_num_batched_tokens, instance_id=self.vllm_config.instance_id, + kv_events_config=self.scheduler.get_kv_event_publisher_config(), ) def process_input_sockets(