feat(mocker): publish native KV events from the vLLM gRPC server

lib/mocker/servers/vllm/Cargo.toml

@@ -17,7 +17,7 @@ name = "dynamo-vllm-mocker-server"
 path = "src/main.rs"
 
 [dependencies]
-dynamo-mocker = { workspace = true }
+dynamo-mocker = { workspace = true, features = ["zmq-events"] }
 dynamo-vllm-sidecar = { path = "../../../sidecar/vllm" }
 
 anyhow = { workspace = true }
@@ -34,5 +34,8 @@ tracing-subscriber = { workspace = true }
 uuid = { workspace = true }
 
 [dev-dependencies]
+dynamo-llm = { path = "../../../llm", default-features = false }
+dynamo-kv-router = { workspace = true }
+dynamo-runtime = { workspace = true }
 dynamo-backend-common = { workspace = true, features = ["testing"] }
 tokio-stream = { workspace = true, features = ["net"] }

lib/mocker/servers/vllm/README.md

@@ -38,6 +38,29 @@ Synthetic output plans are limited to 32,768 tokens. LiveEngine uses a small,
 fixed response buffer for each request and cancels slow consumers rather than
 turning declared output length into a second admission-control policy.
 
+## KV events
+
+Like regular mock workers, the server publishes KV cache events when prefix
+caching is enabled, except in decode mode. It uses the existing Mocker ZMQ
+publisher and reports the endpoint through native engine discovery. The
+sidecar forwards these events to Dynamo's router.
+
+The event publisher binds a free port by default. For a fixed port and an
+optional replay port, add these settings to the server command above:
+
+```bash
+--extra-engine-args '{"block_size":4,"zmq_kv_events_port":5557,"zmq_replay_port":5558}'
+```
+
+Run the sidecar as shown above and use a frontend with `--router-mode kv`.
+Sidecar discovers the event endpoint without an additional setting. When
+running across containers or hosts, expose the event port; a fixed port can
+make this easier.
+
+Set `"enable_prefix_caching":false` to disable both prefix caching and KV
+events. Decode servers do not publish events. As with regular mock workers,
+publisher setup failures are logged and serving continues without KV events.
+
 ## Disaggregated wire-flow
 
 Run separate endpoints for the two emulated vLLM roles:

lib/mocker/servers/vllm/src/main.rs

@@ -73,7 +73,8 @@ async fn main() -> anyhow::Result<()> {
             max_concurrent_requests: args.max_concurrent_requests,
         },
         engine_args,
-    )?;
+    )
+    .await?;
 
     tracing::info!(
         listen = %args.listen,

lib/mocker/servers/vllm/src/server.rs

@@ -6,9 +6,12 @@ use std::pin::Pin;
 use std::sync::Arc;
 
 use clap::ValueEnum;
-use dynamo_mocker::common::protocols::{EngineType, MockEngineArgs, OutputSignal, WorkerType};
-use dynamo_mocker::live::{LiveEngine, LiveRequest, stable_request_uuid};
+use dynamo_mocker::common::protocols::{
+    EngineType, KvEventPublishers, MockEngineArgs, OutputSignal, WorkerType,
+};
+use dynamo_mocker::live::{LiveEngine, LiveEngineConfig, LiveRequest, stable_request_uuid};
 use dynamo_mocker::scheduler::MockerMetrics;
+use dynamo_mocker::services::zmq_events::ZmqKvEventSink;
 use dynamo_vllm_sidecar::proto as pb;
 use futures::Stream;
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
@@ -67,11 +70,16 @@ pub struct VllmMockerService {
     model_info: Arc<pb::ModelInfo>,
     server_info: Arc<pb::ServerInfo>,
     engine: LiveEngine,
+    kv_event_sources: Arc<Vec<pb::KvEventSource>>,
     request_permits: Arc<Semaphore>,
 }
 
 impl VllmMockerService {
-    pub fn new(config: MockerServerConfig, engine_args: MockEngineArgs) -> anyhow::Result<Self> {
+    pub async fn new(
+        config: MockerServerConfig,
+        engine_args: MockEngineArgs,
+    ) -> anyhow::Result<Self> {
+        let engine_args = engine_args.normalized()?;
         anyhow::ensure!(
             engine_args.engine_type == EngineType::Vllm,
             "Mocker engine_type must be vllm"
@@ -134,11 +142,57 @@ impl VllmMockerService {
                 .unwrap_or_default(),
             rl_capabilities: None,
         };
+        // The wire role is separate from the aggregated scheduler used to
+        // emulate disaggregated requests.
+        let sink = if engine_args.needs_kv_publisher() && config.mode != ServerMode::Decode {
+            match ZmqKvEventSink::bind(
+                engine_args.zmq_kv_events_port,
+                engine_args.zmq_replay_port,
+                DP_RANK,
+                u32::try_from(engine_args.block_size)?,
+            )
+            .await
+            {
+                Ok(sink) => Some(sink),
+                Err(error) => {
+                    tracing::error!(dp_rank = DP_RANK, %error, "Failed to create ZMQ KV event sink");
+                    None
+                }
+            }
+        } else {
+            None
+        };
+        let kv_event_sources = sink
+            .as_ref()
+            .map(|sink| pb::KvEventSource {
+                transport: "zmq".to_string(),
+                endpoint: sink.endpoint().to_string(),
+                topic: String::new(),
+                data_parallel_rank: Some(DP_RANK),
+                replay_endpoint: sink.replay_endpoint().unwrap_or_default().to_string(),
+                encoding: "msgpack".to_string(),
+                schema_version: 1,
+                ..Default::default()
+            })
+            .into_iter()
+            .collect();
+        let engine = LiveEngine::start_with_config(
+            engine_args,
+            DP_RANK,
+            LiveEngineConfig {
+                kv_event_publishers: KvEventPublishers::new(
+                    None,
+                    sink.map(|sink| Arc::new(sink) as _),
+                ),
+                ..Default::default()
+            },
+        )?;
         Ok(Self {
             config: Arc::new(config),
             model_info: Arc::new(model_info),
             server_info: Arc::new(server_info),
-            engine: LiveEngine::start(engine_args, DP_RANK)?,
+            engine,
+            kv_event_sources: Arc::new(kv_event_sources),
             request_permits: Arc::new(Semaphore::new(max_concurrent_requests)),
         })
     }
@@ -319,7 +373,7 @@ impl pb::control_server::Control for VllmMockerService {
         _request: Request<pb::GetKvEventSourcesRequest>,
     ) -> Result<Response<pb::GetKvEventSourcesResponse>, Status> {
         Ok(Response::new(pb::GetKvEventSourcesResponse {
-            sources: Vec::new(),
+            sources: (*self.kv_event_sources).clone(),
         }))
     }
 

lib/mocker/servers/vllm/src/server_tests.rs

@@ -155,14 +155,15 @@ fn text_prompts_fail_with_an_actionable_status() {
     assert!(error.message().contains("token_ids"));
 }
 
-#[test]
-fn service_rejects_non_vllm_or_multi_rank_engines() {
+#[tokio::test]
+async fn service_rejects_non_vllm_or_multi_rank_engines() {
     let sglang = MockEngineArgs::builder()
         .engine_type(EngineType::Sglang)
         .build()
         .unwrap();
     assert!(
         VllmMockerService::new(MockerServerConfig::default(), sglang)
+            .await
             .err()
             .unwrap()
             .to_string()
@@ -172,6 +173,7 @@ fn service_rejects_non_vllm_or_multi_rank_engines() {
     let multi_rank = MockEngineArgs::builder().dp_size(2).build().unwrap();
     assert!(
         VllmMockerService::new(MockerServerConfig::default(), multi_rank)
+            .await
             .err()
             .unwrap()
             .to_string()
@@ -184,6 +186,7 @@ fn service_rejects_non_vllm_or_multi_rank_engines() {
         .unwrap();
     assert!(
         VllmMockerService::new(MockerServerConfig::default(), disaggregated)
+            .await
             .err()
             .unwrap()
             .to_string()
@@ -196,6 +199,7 @@ fn service_rejects_non_vllm_or_multi_rank_engines() {
     };
     assert!(
         VllmMockerService::new(disabled, MockEngineArgs::default())
+            .await
             .err()
             .unwrap()
             .to_string()
@@ -208,8 +212,9 @@ fn service_rejects_non_vllm_or_multi_rank_engines() {
 /// capability absence; this test catches it at the Control RPC boundary.
 #[tokio::test]
 async fn unsupported_rl_control_reports_unimplemented() {
-    let service =
-        VllmMockerService::new(MockerServerConfig::default(), MockEngineArgs::default()).unwrap();
+    let service = VllmMockerService::new(MockerServerConfig::default(), MockEngineArgs::default())
+        .await
+        .unwrap();
     let server_info = pb::control_server::Control::get_server_info(
         &service,
         Request::new(pb::GetServerInfoRequest {}),
@@ -241,7 +246,9 @@ async fn unary_generate_maps_capacity_rejection_to_resource_exhausted() {
         .speedup_ratio(0.0)
         .build()
         .unwrap();
-    let service = VllmMockerService::new(MockerServerConfig::default(), args).unwrap();
+    let service = VllmMockerService::new(MockerServerConfig::default(), args)
+        .await
+        .unwrap();
     let mut oversized = request("oversized");
     oversized.prompt = Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds {
         ids: vec![1, 2, 3, 4, 5],
@@ -269,6 +276,7 @@ async fn concurrent_request_limit_rejects_a_stalled_stream() {
         },
         args,
     )
+    .await
     .unwrap();
     let mut first_request = request("stalled");
     first_request.stopping.as_mut().unwrap().max_new_tokens = 100;
@@ -337,7 +345,9 @@ fn decode_rejects_a_handoff_missing_the_opacity_sentinel() {
 
 #[tokio::test]
 async fn unary_generate_accumulates_output_and_terminal_metadata() {
-    let service = VllmMockerService::new(MockerServerConfig::default(), admitting_args()).unwrap();
+    let service = VllmMockerService::new(MockerServerConfig::default(), admitting_args())
+        .await
+        .unwrap();
 
     let mut routed_request = Request::new(request("unary"));
     routed_request.metadata_mut().insert(
@@ -390,7 +400,9 @@ async fn streaming_generate_maps_capacity_rejection_to_resource_exhausted() {
         .speedup_ratio(0.0)
         .build()
         .unwrap();
-    let service = VllmMockerService::new(MockerServerConfig::default(), args).unwrap();
+    let service = VllmMockerService::new(MockerServerConfig::default(), args)
+        .await
+        .unwrap();
     let mut oversized = request("oversized-stream");
     oversized.prompt = Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds {
         ids: vec![1, 2, 3, 4, 5],
@@ -422,7 +434,9 @@ async fn streaming_survives_a_producer_that_outruns_a_stalled_consumer() {
     // The producer runs instantly while the consumer stalls, so the request
     // races far past LiveEngine's fixed per-request buffer. The server's pump
     // must absorb the burst instead of shedding the stream into an INTERNAL.
-    let service = VllmMockerService::new(MockerServerConfig::default(), admitting_args()).unwrap();
+    let service = VllmMockerService::new(MockerServerConfig::default(), admitting_args())
+        .await
+        .unwrap();
     let mut bursty = request("bursty");
     bursty.stopping.as_mut().unwrap().max_new_tokens = 50;
 
@@ -459,3 +473,59 @@ async fn streaming_survives_a_producer_that_outruns_a_stalled_consumer() {
     assert_eq!(finish.num_output_tokens, 50);
     assert_eq!(service.active_request_count(), 0);
 }
+
+#[tokio::test]
+async fn kv_event_discovery_follows_regular_mocker_rules() {
+    for (mode, prefix_caching, expected_sources) in [
+        (ServerMode::Aggregated, true, 1),
+        (ServerMode::Prefill, true, 1),
+        (ServerMode::Decode, true, 0),
+        (ServerMode::Aggregated, false, 0),
+    ] {
+        let mut args = admitting_args();
+        args.enable_prefix_caching = prefix_caching;
+        let service = VllmMockerService::new(
+            MockerServerConfig {
+                mode,
+                ..Default::default()
+            },
+            args,
+        )
+        .await
+        .unwrap();
+        assert_eq!(
+            pb::control_server::Control::get_kv_event_sources(
+                &service,
+                Request::new(pb::GetKvEventSourcesRequest {})
+            )
+            .await
+            .unwrap()
+            .into_inner()
+            .sources
+            .len(),
+            expected_sources
+        );
+    }
+}
+
+#[tokio::test]
+async fn failed_kv_publisher_is_not_advertised() {
+    let occupied = tokio::net::TcpListener::bind("0.0.0.0:0").await.unwrap();
+    let mut args = admitting_args();
+    args.zmq_kv_events_port = Some(occupied.local_addr().unwrap().port());
+    let service = VllmMockerService::new(MockerServerConfig::default(), args)
+        .await
+        .unwrap();
+    assert_eq!(
+        pb::control_server::Control::get_kv_event_sources(
+            &service,
+            Request::new(pb::GetKvEventSourcesRequest {})
+        )
+        .await
+        .unwrap()
+        .into_inner()
+        .sources
+        .len(),
+        0
+    );
+}