feat(kv_index): the chain index, a lossless KV-event relay with state snapshots, and liveness-aware cache routing

.gitignore

@@ -116,6 +116,8 @@ work_dirs/
 
 # Rust target directory
 target/
+# Older builds of the kv-index churn bench wrote this into the crate directory under cargo test
+churn_result.json
 
 # VSCode
 .vscode

.pre-commit-config.yaml

@@ -7,6 +7,8 @@ repos:
       - id: check-symlinks
       - id: destroyed-symlinks
       - id: trailing-whitespace
+        # Patch files keep their whitespace: a blank context line is a single space.
+        exclude: \.patch$
       - id: end-of-file-fixer
       - id: check-yaml
         args: [--allow-multiple-documents, --unsafe]

bindings/python/src/lib.rs

@@ -549,6 +549,17 @@ struct Router {
     prefill_max_inflight_requests_per_worker: i32,
     prefill_queue_size: Option<usize>,
     prefill_queue_timeout_secs: Option<u64>,
+    worker_overload_shed: bool,
+    kv_index: String,
+    worker_stall_secs: u64,
+    worker_wedge_secs: u64,
+    worker_warmup_secs: u64,
+    worker_warmup_share: f32,
+    worker_warmup_blocks: usize,
+    worker_warmup_thin_ratio: f32,
+    worker_warmup_divert_every: u64,
+    selection_policy: String,
+    selection_accounting_ttl_ms: u64,
     /// The keyword-only `discovery` mapping, read by the same rules as
     /// `RouterConfig.discovery`.
     discovery: Option<config::DiscoveryConfig>,
@@ -663,6 +674,14 @@ impl Router {
             })
             .transpose()?;
 
+        let kv_index = config::KvIndexKind::parse(&self.kv_index).ok_or_else(|| {
+            config::ConfigError::InvalidValue {
+                field: "kv_index".to_string(),
+                value: self.kv_index.clone(),
+                reason: "expected 'positional' or 'chain'".to_string(),
+            }
+        })?;
+
         let convert_policy = |policy: &PolicyType| -> config::ConfigResult<ConfigPolicyConfig> {
             Ok(match policy {
                 PolicyType::Random => ConfigPolicyConfig::Random,
@@ -682,6 +701,9 @@ impl Router {
                     cache_index: self.parse_cache_index()?,
                     cache_ttl_secs: self.cache_ttl_secs,
                     cache_boundaries: self.cache_boundaries.clone(),
+                    selection_policy: (self.selection_policy != policies::cost::DEFAULT_POLICY)
+                        .then(|| self.selection_policy.clone()),
+                    selection_accounting_ttl_ms: self.selection_accounting_ttl_ms,
                 },
                 PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
                     load_check_interval_secs: self.load_monitor_interval,
@@ -899,6 +921,17 @@ impl Router {
             .worker_overload_waiting_requests(self.worker_overload_waiting_requests)
             .worker_overload_token_usage(self.worker_overload_token_usage)
             .worker_overload_protection(self.worker_overload_protection)
+            .worker_overload_shed(self.worker_overload_shed)
+            .worker_stall_secs(self.worker_stall_secs)
+            .worker_wedge_secs(self.worker_wedge_secs)
+            .worker_warmup(
+                self.worker_warmup_secs,
+                self.worker_warmup_share,
+                self.worker_warmup_blocks,
+                self.worker_warmup_thin_ratio,
+                self.worker_warmup_divert_every,
+            )
+            .kv_index(kv_index)
             .disable_load_monitoring(self.disable_load_monitoring)
             .load_monitor_interval_secs(self.load_monitor_interval)
             .pd_admission_wait_secs(self.pd_admission_wait_secs)
@@ -1159,9 +1192,9 @@ impl Router {
         cache_ttl_secs = 180,
         job_queue_capacity = 1000,
         job_queue_concurrency = 200,
-        worker_overload_waiting_requests = None,
-        worker_overload_token_usage = None,
-        worker_overload_protection = false,
+        worker_overload_waiting_requests = Some(8),
+        worker_overload_token_usage = Some(0.8),
+        worker_overload_protection = true,
         disable_load_monitoring = false,
         max_buffered_request_bytes = 1_048_576,
         kv_connector_annotation = String::from("smg.ai/kv-connector"),
@@ -1184,6 +1217,17 @@ impl Router {
         prefill_max_inflight_requests_per_worker = -1,
         prefill_queue_size = None,
         prefill_queue_timeout_secs = None,
+        worker_overload_shed = false,
+        kv_index = String::from("positional"),
+        worker_stall_secs = 2,
+        worker_wedge_secs = 3,
+        worker_warmup_secs = 60,
+        worker_warmup_share = 0.25,
+        worker_warmup_blocks = 1024,
+        worker_warmup_thin_ratio = 0.5,
+        worker_warmup_divert_every = 8,
+        selection_policy = String::from("cache-aware-default"),
+        selection_accounting_ttl_ms = 0,
         // Keyword-only, so it never takes a positional slot.
         *,
         discovery = None,
@@ -1352,6 +1396,17 @@ impl Router {
         prefill_max_inflight_requests_per_worker: i32,
         prefill_queue_size: Option<usize>,
         prefill_queue_timeout_secs: Option<u64>,
+        worker_overload_shed: bool,
+        kv_index: String,
+        worker_stall_secs: u64,
+        worker_wedge_secs: u64,
+        worker_warmup_secs: u64,
+        worker_warmup_share: f32,
+        worker_warmup_blocks: usize,
+        worker_warmup_thin_ratio: f32,
+        worker_warmup_divert_every: u64,
+        selection_policy: String,
+        selection_accounting_ttl_ms: u64,
         discovery: Option<Bound<'_, PyAny>>,
     ) -> PyResult<Self> {
         // Two spellings of one choice: refuse both rather than pick one.
@@ -1546,6 +1601,17 @@ impl Router {
             prefill_max_inflight_requests_per_worker,
             prefill_queue_size,
             prefill_queue_timeout_secs,
+            worker_overload_shed,
+            kv_index,
+            worker_stall_secs,
+            worker_wedge_secs,
+            worker_warmup_secs,
+            worker_warmup_share,
+            worker_warmup_blocks,
+            worker_warmup_thin_ratio,
+            worker_warmup_divert_every,
+            selection_policy,
+            selection_accounting_ttl_ms,
             discovery,
         })
     }

bindings/python/src/servicer.rs

@@ -787,6 +787,7 @@ impl PyTokenSpeedGrpcServer {
         max_running_requests = 0,
         data_parallel_size = 1,
         kv_events_endpoint = String::new(),
+        kv_events_replay_endpoint = String::new(),
         kv_events_topic = String::new(),
         engine_startup_timeout_secs = None,
     ))]
@@ -822,6 +823,7 @@ impl PyTokenSpeedGrpcServer {
         max_running_requests: i32,
         data_parallel_size: i32,
         kv_events_endpoint: String,
+        kv_events_replay_endpoint: String,
         kv_events_topic: String,
         engine_startup_timeout_secs: Option<f64>,
     ) -> PyResult<Self> {
@@ -850,6 +852,7 @@ impl PyTokenSpeedGrpcServer {
             max_running_requests,
             data_parallel_size,
             kv_events_endpoint,
+            kv_events_replay_endpoint,
             kv_events_topic,
         };
         let config = TokenSpeedServicerConfig {
@@ -952,6 +955,9 @@ impl PySglangGrpcServer {
         sglang_version = String::new(),
         max_running_requests = 0,
         data_parallel_size = 1,
+        kv_events_endpoint = String::new(),
+        kv_events_replay_endpoint = String::new(),
+        kv_events_topic = String::new(),
         engine_startup_timeout_secs = None,
     ))]
     #[expect(clippy::too_many_arguments)]
@@ -985,6 +991,9 @@ impl PySglangGrpcServer {
         sglang_version: String,
         max_running_requests: i32,
         data_parallel_size: i32,
+        kv_events_endpoint: String,
+        kv_events_replay_endpoint: String,
+        kv_events_topic: String,
         engine_startup_timeout_secs: Option<f64>,
     ) -> PyResult<Self> {
         let model = SglangModelInfo {
@@ -1011,6 +1020,9 @@ impl PySglangGrpcServer {
             sglang_version,
             max_running_requests,
             data_parallel_size,
+            kv_events_endpoint,
+            kv_events_replay_endpoint,
+            kv_events_topic,
         };
         let config = SglangServicerConfig {
             bind_address,