feat(grpc): add KV event source discovery (#50033)

Signed-off-by: Connor Carpenter <connorc@nvidia.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

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;
+}

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

rust/src/engine-core-client/src/protocol/handshake.rs

@@ -24,6 +24,19 @@ pub struct ReadyMessage {
     pub parallel_config_hash: Option<String>,
 }
 
+/// 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<String>,
+    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<u64>,
     /// Maximum achievable request concurrency given the KV cache, if reported.
     pub kv_cache_max_concurrency: Option<f64>,
+    /// KV-event publisher configuration, if configured.
+    #[serde(default)]
+    pub kv_events_config: Option<KvEventsConfig>,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup

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

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

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<pb::GetKvEventSourcesRequest>,
+    ) -> Result<Response<pb::GetKvEventSourcesResponse>, 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<pb::KvEventSource> {
+    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,
+    })
 }

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() {

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

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

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

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]:

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(

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(