feat(kv-router): index Mooncake KV events

components/src/dynamo/sglang/register.py

@@ -8,7 +8,6 @@
 from typing import Any, List, Optional
 
 import sglang as sgl
-from sglang.srt.environ import envs
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
@@ -232,33 +231,6 @@ def _get_mooncake_runtime_data(server_args: ServerArgs) -> Optional[dict[str, An
         getattr(server_args, "hicache_storage_backend_extra_config", None)
     )
 
-    try:
-        from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import (
-            MooncakeStoreConfig,
-        )
-    except ImportError as e:
-        logging.warning(f"MooncakeStoreConfig import unavailable: {e}")
-        return None
-
-    # Graceful degradation: Mooncake runtime metadata is optional. If config
-    # resolution fails for any reason (file not found, malformed env vars,
-    # upstream API change), skip publishing the metadata rather than crashing
-    # the worker -- the worker still serves requests, just without HiCache
-    # router hints. Broad catch is intentional per python-guidelines.md.
-    try:
-        if extra_config and (
-            extra_config.get("master_server_address") is not None
-            or extra_config.get("client_server_address") is not None
-        ):
-            mooncake_config = MooncakeStoreConfig.load_from_extra_config(extra_config)
-        elif envs.SGLANG_HICACHE_MOONCAKE_CONFIG_PATH.is_set():
-            mooncake_config = MooncakeStoreConfig.from_file()
-        else:
-            mooncake_config = MooncakeStoreConfig.load_from_env()
-    except Exception as e:
-        logging.warning(f"Failed to resolve Mooncake config for runtime metadata: {e}")
-        return None
-
     tp_size = int(getattr(server_args, "tp_size", 1) or 1)
     pp_size = int(getattr(server_args, "pp_size", 1) or 1)
 
@@ -296,10 +268,6 @@ def _get_mooncake_runtime_data(server_args: ServerArgs) -> Optional[dict[str, An
     if not isinstance(extra_backend_tag, str) or not extra_backend_tag:
         extra_backend_tag = None
 
-    master_server_address = getattr(mooncake_config, "master_server_address", None)
-    if not isinstance(master_server_address, str) or not master_server_address:
-        master_server_address = None
-
     return {
         "backend": "mooncake",
         "page_size": int(getattr(server_args, "page_size", 1) or 1),
@@ -310,10 +278,7 @@ def _get_mooncake_runtime_data(server_args: ServerArgs) -> Optional[dict[str, An
         "tp_lcm_size": tp_lcm_size,
         "should_split_heads": should_split_heads,
         "extra_backend_tag": extra_backend_tag,
-        "master_server_address": master_server_address,
-        "master_metrics_port": int(
-            getattr(mooncake_config, "master_metrics_port", 9003)
-        ),
+        "kv_events_endpoint": os.getenv("DYN_MOONCAKE_KV_EVENTS_ENDPOINT") or None,
     }
 
 

docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/hicache.md

@@ -17,7 +17,7 @@ SGLang HiCache extends RadixAttention with a multi-tier KV cache that transparen
 What Dynamo adds on top of HiCache:
 
 - **Tier-aware routing.** The KV router tracks which cache tier each block lives on (GPU / Host / External) and uses that when scoring candidate workers — not just device overlap.
-- **Shared-pool awareness.** When an external backend such as Mooncake is configured, the router queries the shared pool in parallel with its own indexer so it can discount prefill cost for blocks any worker can fetch, not just blocks the candidate holds locally.
+- **Shared-pool awareness.** When an external backend such as Mooncake is configured, the router tracks Mooncake object events so it can discount prefill cost for blocks any worker can fetch, not just blocks the candidate holds locally.
 
 If you are running a single worker with HiCache and no shared pool, no Dynamo-side configuration is required — the worker reports KV events to the router as usual.
 
@@ -84,15 +84,17 @@ flowchart LR
 
     Worker -- "KV events (store/remove + medium)" --> Router
     Worker -- "writes pages" --> Mooncake
-    Router -- "batch_query on each request" --> Mooncake
+    Mooncake -- "object events (stored/removed)" --> Router
 ```
 
-On every request the router runs two lookups in parallel:
+The router maintains two indexes:
 
 - Its own radix tree, built from worker KV events (per-tier).
-- A batch query to the Mooncake master for blocks reachable from the shared pool.
+- A set of Mooncake object keys and verified logical groups, built from the Mooncake master's event stream.
 
-If the shared-pool query fails, the router falls back to indexer-only scoring and logs a warning. The request still succeeds.
+Request routing checks both indexes locally. When an event includes SGLang's `group_id`, the first lookup verifies every physical object in the group and caches the result. Later requests use one group lookup per page; any event for that group invalidates the cached result. Events without `group_id` use exact object-key checks.
+
+If the event stream starts after Mooncake already contains objects, the shared-pool index starts empty and learns about subsequent events. A sequence gap clears the shared-pool index to avoid stale hits.
 
 ### Scoring
 
@@ -143,7 +145,8 @@ Earlier SGLang versions do not emit `medium=CPU_PINNED` for Host-tier residency,
 You also need:
 
 - Dynamo router started with `--shared-cache-type hicache` (see [Configuration](#configuration)).
-- A Mooncake master reachable from the Dynamo frontend host. Worker-side Mooncake config (master address, page size, TP/PP layout, split-head layout) is published automatically via each worker's registration metadata when the worker is started with `--hicache-storage-backend mooncake`.
+- A Mooncake master with KV events enabled. This requires the event publisher introduced by [kvcache-ai/Mooncake#2214](https://github.com/kvcache-ai/Mooncake/pull/2214) until that change is available in a Mooncake release.
+- A Mooncake KV event endpoint reachable from the Dynamo frontend host. Worker-side Mooncake config (event endpoint, page size, TP/PP layout, and split-head layout) is published automatically through each worker's registration metadata.
 
 ## Setup
 
@@ -160,15 +163,16 @@ python -m dynamo.sglang \
   --hicache-ratio 2 \
   --hicache-write-policy write_through \
   --hicache-storage-backend mooncake \
-  --hicache-storage-backend-extra-config '{"master_server_address": "mooncake-master.internal:50051"}' \
+  --hicache-storage-backend-extra-config '{"master_server_address": "mooncake-master.internal:50051", "enable_group_semantics": true}' \
   --skip-tokenizer-init
 ```
 
 Launch additional workers on other GPUs / hosts with the same Mooncake config so they back to the same cluster.
 
-**Dynamo frontend** — enable tier-aware routing:
+**Dynamo frontend** — configure the Mooncake event endpoint and enable tier-aware routing:
 
 ```bash
+DYN_MOONCAKE_KV_EVENTS_ENDPOINT=tcp://mooncake-master.internal:5557 \
 python -m dynamo.frontend \
   --http-port 8000 \
   --router-mode kv \
@@ -180,12 +184,15 @@ python -m dynamo.frontend \
 
 | Flag                        | Env var                       | Default | Description                                                                                                                                                        |
 | --------------------------- | ----------------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `--shared-cache-type`       | `DYN_SHARED_CACHE_TYPE`       | `none`  | `none` disables shared-pool lookups; `hicache` enables Mooncake queries.                                                                                           |
-| `--shared-cache-multiplier` | `DYN_SHARED_CACHE_MULTIPLIER` | `0.5`   | Discount factor for shared-pool hits. `0.0` queries but ignores them; `0.5` treats a shared hit as half a device hit; `1.0` treats shared and device hits equally. |
+| `--shared-cache-type`       | `DYN_SHARED_CACHE_TYPE`       | `none`  | `none` disables shared-pool tracking; `hicache` enables the Mooncake event-backed index.                                                                           |
+| `--shared-cache-multiplier` | `DYN_SHARED_CACHE_MULTIPLIER` | `0.5`   | Discount factor for shared-pool hits. `0.0` ignores them; `0.5` treats a shared hit as half a device hit; `1.0` treats shared and device hits equally.             |
+| —                           | `DYN_MOONCAKE_KV_EVENTS_ENDPOINT` | unset | Mooncake PUB endpoint consumed by the frontend. When unset, Dynamo uses one consistently advertised worker endpoint. |
 
 Per-request overrides are available via `RouterConfigOverride.shared_cache_multiplier` for A/B experimentation without restarting the router.
 
-No extra flags are required on the worker. When `--hicache-storage-backend mooncake` is set, Dynamo publishes the required metadata (page size, TP/PP layout, master address) via the worker's `ModelRuntimeConfig.engine_specific` blob under the key `sglang_hicache_mooncake`.
+Set `DYN_MOONCAKE_KV_EVENTS_ENDPOINT` on the frontend to the Mooncake PUB endpoint, such as `tcp://mooncake-master.internal:5557`. The endpoint must be reachable from the frontend. Workers can advertise the same variable as a fallback, but a missing worker value does not disable shared-cache routing.
+
+Set `enable_group_semantics` to `true` in `--hicache-storage-backend-extra-config` to include SGLang logical group IDs in Mooncake metadata. Dynamo falls back to exact physical-key checks when group metadata is unavailable.
 
 ## Verification
 
@@ -199,12 +206,13 @@ python -m dynamo.sglang ... --log-level debug 2>&1 | grep -E 'BlockStored|BlockR
 
 If `medium` is missing or Host-tier transitions never report `CPU_PINNED`, confirm that the worker runs SGLang 0.5.11 or later (or a custom build that includes PR #22894).
 
-**Router sees the shared pool.** Two new histograms are exposed on the frontend's Prometheus endpoint:
+**Router sees the shared pool.** Shared-cache metrics are exposed on the frontend's Prometheus endpoint:
 
-| Metric                              | Meaning                                                                  |
-| ----------------------------------- | ------------------------------------------------------------------------ |
-| `router_shared_cache_hit_rate`      | Fraction of request blocks found in the shared pool (0.0–1.0).           |
-| `router_shared_cache_beyond_blocks` | Blocks in the shared pool _beyond_ the selected worker's device overlap. |
+| Metric                                    | Meaning                                                                  |
+| ----------------------------------------- | ------------------------------------------------------------------------ |
+| `router_shared_cache_hit_rate`            | Fraction of request blocks found in the shared pool (0.0–1.0).           |
+| `router_shared_cache_beyond_blocks`       | Blocks in the shared pool _beyond_ the selected worker's device overlap. |
+| `dynamo_router_shared_cache_errors_total` | Shared-cache query and Mooncake subscriber failures.                     |
 
 ```bash
 curl -s localhost:8000/metrics | grep shared_cache
@@ -214,10 +222,10 @@ curl -s localhost:8000/metrics | grep shared_cache
 
 | Symptom                                                  | Likely cause                                                                 | Fix                                                                                           |
 | -------------------------------------------------------- | ---------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------- |
-| `shared_cache_hit_rate` is always 0                      | Mooncake master unreachable from the router host                             | Check network path; the router logs `Shared cache query failed` when it can't reach Mooncake. |
-| Events only ever carry `medium=GPU`                      | SGLang older than 0.5.11 or a custom build missing [PR #22894](https://github.com/sgl-project/sglang/pull/22894) | Upgrade to SGLang 0.5.11 or later.                                                      |
-| Workers registered but router never queries shared cache | `--shared-cache-type` left at default `none`                                 | Set `--shared-cache-type hicache` on the frontend.                                            |
-| Queries issued but winning worker rarely changes         | `--shared-cache-multiplier 0.0`                                              | Raise the multiplier — typical starting range is `0.3`–`0.7`.                                 |
+| `shared_cache_hit_rate` is always 0                     | Event endpoint missing, unreachable, or connected after objects were stored | Set `DYN_MOONCAKE_KV_EVENTS_ENDPOINT` on the frontend and inspect `dynamo_router_shared_cache_errors_total`. |
+| Events only ever carry `medium=GPU`                     | SGLang older than 0.5.11 or a custom build missing [PR #22894](https://github.com/sgl-project/sglang/pull/22894) | Upgrade to SGLang 0.5.11 or later. |
+| Workers registered but router never tracks shared cache | `--shared-cache-type` left at default `none`                                | Set `--shared-cache-type hicache` on the frontend.                                        |
+| Shared hits do not affect worker selection              | `--shared-cache-multiplier 0.0`                                             | Raise the multiplier — typical starting range is `0.3`–`0.7`.                             |
 | Page-size mismatch warnings                              | Router `--page-size` doesn't match worker `--page-size`                      | They must agree; the router hashes pages using the worker's page size.                        |
 | Router logs "no workers have HiCache enabled"            | No worker published `sglang_hicache_mooncake` metadata                       | Confirm workers started with `--hicache-storage-backend mooncake`.                            |
 

lib/llm/src/discovery/model_manager.rs

@@ -148,6 +148,9 @@ pub struct ModelManager {
     /// Per-endpoint runtime config watchers. Keyed by EndpointId (includes namespace).
     runtime_configs: DashMap<EndpointId, RuntimeConfigWatch>,
 
+    /// Per-endpoint HiCache state and its one Mooncake event subscriber.
+    hicache_caches: DashMap<EndpointId, HicacheSharedKvCache>,
+
     /// Shared KV-source membership coordinators, scoped by exact serving endpoint.
     /// Weak ownership lets the discovery loop stop when its last consumer goes away.
     kv_source_memberships: DashMap<EndpointId, Weak<KvSourceMembershipCoordinator>>,
@@ -188,6 +191,7 @@ impl ModelManager {
             prefill_router_activators: DashMap::new(),
             encoder_router_activators: DashMap::new(),
             runtime_configs: DashMap::new(),
+            hicache_caches: DashMap::new(),
             kv_source_memberships: DashMap::new(),
             lora_domains: DashMap::new(),
             lora_enabled: crate::lora::lora_serving_enabled(),
@@ -1050,6 +1054,45 @@ impl ModelManager {
         lora_enabled && worker_type == crate::protocols::common::timing::WORKER_TYPE_DECODE
     }
 
+    fn hicache_cache_for(
+        &self,
+        endpoint: &Endpoint,
+        runtime_configs: RuntimeConfigWatch,
+    ) -> HicacheSharedKvCache {
+        self.hicache_caches
+            .entry(endpoint.id())
+            .or_insert_with(|| {
+                let frontend_kv_events_endpoint = std::env::var("DYN_MOONCAKE_KV_EVENTS_ENDPOINT")
+                    .ok()
+                    .filter(|endpoint| !endpoint.is_empty());
+                let cache = HicacheSharedKvCache::new_with_cancellation_and_endpoint(
+                    runtime_configs,
+                    endpoint.component().drt().child_token(),
+                    frontend_kv_events_endpoint,
+                );
+                cache.start_subscriber();
+                cache
+            })
+            .clone()
+    }
+
+    pub fn remove_hicache_caches(&self, namespace: &str, component: &str) {
+        let endpoint_ids = self
+            .hicache_caches
+            .iter()
+            .filter(|entry| {
+                entry.key().namespace == namespace && entry.key().component == component
+            })
+            .map(|entry| entry.key().clone())
+            .collect::<Vec<_>>();
+
+        for endpoint_id in endpoint_ids {
+            if let Some((_, cache)) = self.hicache_caches.remove(&endpoint_id) {
+                cache.shutdown();
+            }
+        }
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub async fn kv_chooser_for(
         &self,
@@ -1127,9 +1170,9 @@ impl ModelManager {
                     worker_component = worker_component_name,
                     "Using HiCache shared KV cache"
                 );
-                Some(Box::new(HicacheSharedKvCache::new(
-                    workers_with_configs.clone(),
-                )))
+                Some(Box::new(
+                    self.hicache_cache_for(endpoint, workers_with_configs.clone()),
+                ))
             }
         };
 
@@ -2378,6 +2421,38 @@ mod tests {
         assert!(mm.get_model("llama").is_some());
     }
 
+    #[test]
+    fn remove_hicache_caches_cancels_only_the_removed_component() {
+        let manager = ModelManager::new();
+        let (_tx, runtime_configs) =
+            tokio::sync::watch::channel(HashMap::<WorkerId, ModelRuntimeConfig>::new());
+        let cancelled = tokio_util::sync::CancellationToken::new();
+        let retained = tokio_util::sync::CancellationToken::new();
+        manager.hicache_caches.insert(
+            EndpointId::from("ns.worker.generate"),
+            HicacheSharedKvCache::new_with_cancellation(runtime_configs.clone(), cancelled.clone()),
+        );
+        manager.hicache_caches.insert(
+            EndpointId::from("ns.other.generate"),
+            HicacheSharedKvCache::new_with_cancellation(runtime_configs, retained.clone()),
+        );
+
+        manager.remove_hicache_caches("ns", "worker");
+
+        assert!(cancelled.is_cancelled());
+        assert!(!retained.is_cancelled());
+        assert!(
+            !manager
+                .hicache_caches
+                .contains_key(&EndpointId::from("ns.worker.generate"))
+        );
+        assert!(
+            manager
+                .hicache_caches
+                .contains_key(&EndpointId::from("ns.other.generate"))
+        );
+    }
+
     #[test]
     fn test_alias_resolution_maps_to_primary() {
         let mm = ModelManager::new();