48fc2e2797 · tduncan/Nvidia-Optimization-Stack-VLLM

rust/proto/control.proto

+21

@@ -8,6 +8,7 @@ service Control {

8 rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {}

9 rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {}

10 rpc Abort (AbortRequest) returns (AbortResponse) {}

11 + rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {}

12 }

13

14 message GetServerInfoRequest {}

@@ -51,3 +52,23 @@ message AbortRequest {

52 }

53

54 message AbortResponse {}

55 +

56 +// ======================================================================================

57 +// KV discovery

58 +// ======================================================================================

59 +

60 +message GetKvEventSourcesRequest {}

61 +message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; }

62 +

63 +message KvEventSource {

64 + string transport = 1;

65 + string endpoint = 2;

66 + string topic = 3;

67 + string replay_endpoint = 4;

68 + optional uint32 data_parallel_rank = 5;

69 + string encoding = 6;

70 + uint32 schema_version = 7;

71 + uint32 buffer_steps = 8;

72 + uint32 hwm = 9;

73 + uint32 max_queue_size = 10;

74 +}

rust/src/engine-core-client/src/mock_engine.rs

+1

@@ -66,6 +66,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {

66 instance_id: "test-instance".to_string(),

67 kv_cache_size_tokens: None,

68 kv_cache_max_concurrency: None,

69 + kv_events_config: None,

70 }

71 }

72

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

+16

@@ -24,6 +24,19 @@ pub struct ReadyMessage {

24 pub parallel_config_hash: Option<String>,

25 }

26

27 +/// KV-event publisher configuration reported by EngineCore.

28 +#[derive(Debug, Clone, Serialize, Deserialize)]

29 +pub struct KvEventsConfig {

30 + pub enable_kv_cache_events: bool,

31 + pub publisher: String,

32 + pub endpoint: String,

33 + pub replay_endpoint: Option<String>,

34 + pub buffer_steps: u32,

35 + pub hwm: u32,

36 + pub max_queue_size: u32,

37 + pub topic: String,

38 +}

39 +

40 /// Post-initialization configuration sent from each engine on the input socket

41 /// registration message, after the handshake completes.

42 ///

@@ -71,6 +84,9 @@ pub struct EngineCoreReadyResponse {

84 pub kv_cache_size_tokens: Option<u64>,

85 /// Maximum achievable request concurrency given the KV cache, if reported.

86 pub kv_cache_max_concurrency: Option<f64>,

87 + /// KV-event publisher configuration, if configured.

88 + #[serde(default)]

89 + pub kv_events_config: Option<KvEventsConfig>,

90 }

91

92 /// Frontend-owned ZMQ addresses that are sent to the engine during startup

rust/src/engine-core-client/src/tests/client.rs

+16 -1

@@ -19,7 +19,7 @@ use zeromq::prelude::{Socket, SocketRecv, SocketSend};

19 use zeromq::util::PeerIdentity;

20 use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage};

21

22 -use crate::protocol::handshake::{HandshakeInitMessage, ReadyMessage};

22 +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};

23 use crate::protocol::logprobs::MaybeWireLogprobs;

24 use crate::protocol::multimodal::{

25 MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange,

@@ -2677,6 +2677,21 @@ fn python_msgpack_fixtures_match_rust_encoding() {

2677 rust_ready_keys, python_ready_keys,

2678 "EngineCoreReadyResponse drifted from the Python dataclass",

2679 );

2680 +

2681 + let ready_response: EngineCoreReadyResponse =

2682 + rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap();

2683 + let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode");

2684 + assert!(kv_events_config.enable_kv_cache_events);

2685 + assert_eq!(kv_events_config.publisher, "zmq");

2686 + assert_eq!(kv_events_config.endpoint, "tcp://127.0.0.1:5557");

2687 + assert_eq!(

2688 + kv_events_config.replay_endpoint.as_deref(),

2689 + Some("tcp://127.0.0.1:5558")

2690 + );

2691 + assert_eq!(kv_events_config.topic, "kv");

2692 + assert_eq!(kv_events_config.buffer_steps, 10_000);

2693 + assert_eq!(kv_events_config.hwm, 100_000);

2694 + assert_eq!(kv_events_config.max_queue_size, 100_000);

2695 }

2696

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

rust/src/engine-core-client/src/tests/python_compat.py

+23

@@ -353,6 +353,18 @@ multipart_prompt_logprobs = engine_outputs_wire(

353 )

354

355

356 +@dataclass

357 +class KVEventsConfig:

358 + enable_kv_cache_events: bool

359 + publisher: str

360 + endpoint: str

361 + replay_endpoint: str | None

362 + buffer_steps: int

363 + hwm: int

364 + max_queue_size: int

365 + topic: str

366 +

367 +

368 @dataclass

369 class EngineCoreReadyResponse:

370 max_model_len: int

@@ -372,6 +384,7 @@ class EngineCoreReadyResponse:

384 instance_id: str

385 kv_cache_size_tokens: int | None = None

386 kv_cache_max_concurrency: float | None = None

387 + kv_events_config: KVEventsConfig | None = None

388

389

390 ready_response = EngineCoreReadyResponse(

@@ -390,6 +403,16 @@ ready_response = EngineCoreReadyResponse(

403 max_num_seqs=256,

404 max_num_batched_tokens=8192,

405 instance_id="test-instance",

406 + kv_events_config=KVEventsConfig(

407 + enable_kv_cache_events=True,

408 + publisher="zmq",

409 + endpoint="tcp://127.0.0.1:5557",

410 + replay_endpoint="tcp://127.0.0.1:5558",

411 + buffer_steps=10_000,

412 + hwm=100_000,

413 + max_queue_size=100_000,

414 + topic="kv",

415 + ),

416 )

417

418 print(msgspec.msgpack.encode(request).hex())

rust/src/server/src/grpc/control.rs

+29

@@ -103,4 +103,33 @@ impl pb::control_server::Control for ControlServiceImpl {

103 .map_err(|error| Status::internal(error.to_report_string()))?;

104 Ok(Response::new(pb::AbortResponse {}))

105 }

106 +

107 + async fn get_kv_event_sources(

108 + &self,

109 + _request: Request<pb::GetKvEventSourcesRequest>,

110 + ) -> Result<Response<pb::GetKvEventSourcesResponse>, Status> {

111 + let client = self.state.engine_core_client();

112 + let sources = client.ready_responses().into_iter().filter_map(kv_event_source).collect();

113 + Ok(Response::new(pb::GetKvEventSourcesResponse { sources }))

114 + }

115 +}

116 +

117 +pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option<pb::KvEventSource> {

118 + let config = response.kv_events_config.as_ref()?;

119 + if !config.enable_kv_cache_events || config.publisher != "zmq" {

120 + return None;

121 + }

122 +

123 + Some(pb::KvEventSource {

124 + transport: "zmq".to_string(),

125 + endpoint: config.endpoint.clone(),

126 + topic: config.topic.clone(),

127 + replay_endpoint: config.replay_endpoint.clone().unwrap_or_default(),

128 + data_parallel_rank: Some(response.data_parallel_rank),

129 + encoding: "msgpack".to_string(),

130 + schema_version: 1,

131 + buffer_steps: config.buffer_steps,

132 + hwm: config.hwm,

133 + max_queue_size: config.max_queue_size,

134 + })

135 }

rust/src/server/src/grpc/tests.rs

+36

@@ -29,6 +29,7 @@ use vllm_engine_core_client::mock_engine::{

29 DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS,

30 default_ready_response,

31 };

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

33 use vllm_engine_core_client::protocol::output::{

34 EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,

35 };

@@ -44,6 +45,7 @@ use vllm_tokenizer::test_utils::TestTokenizer;

45 use zeromq::prelude::{SocketRecv, SocketSend};

46 use zeromq::{DealerSocket, PushSocket, ZmqMessage};

47

48 +use super::control::kv_event_source;

49 use super::pb::control_client::ControlClient;

50 use super::pb::inference_client::InferenceClient;

51 use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb};

@@ -1319,6 +1321,40 @@ async fn control_aggregates_multi_engine_capacity() {

1321

1322 drop(engine_tasks);

1323 }

1324 +

1325 +#[test]

1326 +fn kv_event_source_filters_and_exposes_zmq_publisher() {

1327 + let mut ready = default_ready_response();

1328 + ready.data_parallel_rank = 2;

1329 + ready.kv_events_config = Some(KvEventsConfig {

1330 + enable_kv_cache_events: false,

1331 + publisher: "null".to_string(),

1332 + endpoint: "tcp://*:5559".to_string(),

1333 + replay_endpoint: Some("tcp://*:5560".to_string()),

1334 + buffer_steps: 10_000,

1335 + hwm: 100_000,

1336 + max_queue_size: 100_000,

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

1338 + });

1339 +

1340 + assert!(kv_event_source(&ready).is_none());

1341 +

1342 + let config = ready.kv_events_config.as_mut().unwrap();

1343 + config.enable_kv_cache_events = true;

1344 + config.publisher = "zmq".to_string();

1345 + let source = kv_event_source(&ready).expect("configured ZMQ event source");

1346 + assert_eq!(source.transport, "zmq");

1347 + assert_eq!(source.endpoint, "tcp://*:5559");

1348 + assert_eq!(source.topic, "kv");

1349 + assert_eq!(source.replay_endpoint, "tcp://*:5560");

1350 + assert_eq!(source.data_parallel_rank, Some(2));

1351 + assert_eq!(source.encoding, "msgpack");

1352 + assert_eq!(source.schema_version, 1);

1353 + assert_eq!(source.buffer_steps, 10_000);

1354 + assert_eq!(source.hwm, 100_000);

1355 + assert_eq!(source.max_queue_size, 100_000);

1356 +}

1357 +

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

1359 #[serial]

1360 async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {

tests/distributed/test_events.py

+8 -3

@@ -284,7 +284,7 @@ def test_data_parallel_rank_tagging(publisher_config):

284 sub_1.close()

285

286

287 -def test_event_publisher_factory():

287 +def test_event_publisher_factory(random_port):

288 """Test event publisher factory creation behavior under different configurations"""

289 from vllm.config.kv_events import KVEventsConfig

290 from vllm.distributed.kv_events import ZmqEventPublisher

@@ -308,10 +308,15 @@ def test_event_publisher_factory():

308 config = KVEventsConfig(

309 enable_kv_cache_events=True,

310 publisher="zmq",

311 - endpoint="inproc://test-factory-true",

311 + endpoint=f"tcp://*:{random_port}",

312 + replay_endpoint=f"tcp://*:{random_port + 100}",

313 )

313 - publisher = EventPublisherFactory.create(config, DP_RANK)

314 + publisher = EventPublisherFactory.create(config, DP_RANK + 1)

315 assert isinstance(publisher, ZmqEventPublisher)

316 + resolved_config = publisher.get_publisher_config()

317 + assert resolved_config.endpoint == f"tcp://*:{random_port + 1}"

318 + assert resolved_config.replay_endpoint == f"tcp://*:{random_port + 101}"

319 + assert config.endpoint == f"tcp://*:{random_port}"

320 publisher.shutdown()

321

322 # test unknown publisher

vllm/distributed/kv_events.py

+18

@@ -272,6 +272,10 @@ class EventPublisher(ABC):

272 def shutdown(self) -> None:

273 """Shutdown the publisher."""

274

275 + def get_publisher_config(self) -> KVEventsConfig | None:

276 + """Return the publisher's resolved runtime configuration."""

277 + return None

278 +

279

280 class NullEventPublisher(EventPublisher):

281 """No-op implementation (default when disabled)."""

@@ -335,6 +339,17 @@ class ZmqEventPublisher(EventPublisher):

339 self._replay_endpoint = self.offset_endpoint_port(

340 replay_endpoint, self._dp_rank

341 )

342 + assert self._endpoint is not None

343 + self._publisher_config = KVEventsConfig(

344 + enable_kv_cache_events=True,

345 + publisher="zmq",

346 + endpoint=self._endpoint,

347 + replay_endpoint=self._replay_endpoint,

348 + buffer_steps=buffer_steps,

349 + hwm=hwm,

350 + max_queue_size=max_queue_size,

351 + topic=topic,

352 + )

353 self._hwm = hwm

354 self._socket_setup()

355

@@ -351,6 +366,9 @@ class ZmqEventPublisher(EventPublisher):

366 )

367 self._thread.start()

368

369 + def get_publisher_config(self) -> KVEventsConfig:

370 + return self._publisher_config

371 +

372 def publish(self, events: EventBatch) -> None:

373 if not self._running:

374 raise RuntimeError("Publisher is closed")

vllm/v1/core/sched/interface.py

+4

@@ -9,6 +9,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

9

10 if TYPE_CHECKING:

11 from vllm.config import VllmConfig

12 + from vllm.config.kv_events import KVEventsConfig

13 from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase

14 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1

15 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

@@ -256,3 +257,6 @@ class SchedulerInterface(ABC):

257

258 def get_ec_connector(self) -> "ECConnectorBase | None":

259 return None

260 +

261 + def get_kv_event_publisher_config(self) -> "KVEventsConfig | None":

262 + return None

vllm/v1/core/sched/scheduler.py

+4 -1

@@ -8,7 +8,7 @@ from dataclasses import replace

8 from typing import Any

9

10 from vllm.compilation.cuda_graph import CUDAGraphStat

11 -from vllm.config import VllmConfig

11 +from vllm.config import KVEventsConfig, VllmConfig

12 from vllm.distributed.ec_transfer.ec_connector.base import (

13 ECConnectorBase,

14 ECConnectorMetadata,

@@ -2554,6 +2554,9 @@ class Scheduler(SchedulerInterface):

2554 def get_ec_connector(self) -> ECConnectorBase | None:

2555 return self.ec_connector

2556

2557 + def get_kv_event_publisher_config(self) -> KVEventsConfig | None:

2558 + return self.kv_event_publisher.get_publisher_config()

2559 +

2560 def _connector_finished(

2561 self, request: Request

2562 ) -> tuple[bool, dict[str, Any] | None]:

vllm/v1/engine/__init__.py

+2

@@ -11,6 +11,7 @@ import msgspec

11 import numpy as np

12 import torch

13

14 +from vllm.config.kv_events import KVEventsConfig

15 from vllm.lora.request import LoRARequest

16 from vllm.multimodal.inputs import MultiModalFeatureSpec

17 from vllm.pooling_params import PoolingParams

@@ -90,6 +91,7 @@ class EngineCoreReadyResponse:

91 # KV cache capacity (None for encoder-only/attention-free models).

92 kv_cache_size_tokens: int | None = None

93 kv_cache_max_concurrency: float | None = None

94 + kv_events_config: KVEventsConfig | None = None

95

96

97 class EngineCoreRequest(

vllm/v1/engine/core.py

+1

@@ -1639,6 +1639,7 @@ class EngineCoreProc(EngineCore):

1639 max_num_seqs=scheduler_config.max_num_seqs,

1640 max_num_batched_tokens=scheduler_config.max_num_batched_tokens,

1641 instance_id=self.vllm_config.instance_id,

1642 + kv_events_config=self.scheduler.get_kv_event_publisher_config(),

1643 )

1644

1645 def process_input_sockets(

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