refactor(kv_index)!: the chain index is the only KV index, the positional indexer goes

.github/workflows/benchmark-radix-tree.yml

@@ -65,7 +65,6 @@ jobs:
           cargo bench -p kv-index --bench throughput_bench -- \
             --num-sessions 200000 \
             --num-workers 256 \
-            --jump-size 8 \
             --sweep-min-ms 10 \
             --sweep-steps 20 2>&1 | tee throughput_output.txt
 
@@ -90,7 +89,6 @@ jobs:
               --trace-path "$GITHUB_WORKSPACE/traces/${trace}.jsonl" \
               --trace-duplication-factor 20 \
               --num-workers 256 \
-              --jump-size 8 \
               --sweep-min-ms 10 \
               --sweep-steps 20 2>&1 | tee -a mooncake_output.txt
             echo "" >> mooncake_output.txt

Cargo.toml

@@ -14,7 +14,7 @@ llm-tokenizer = { version = "1.8.0", path = "crates/tokenizer" }
 smg-auth = { version = "1.2.3", path = "crates/auth" }
 smg-mcp = { version = "2.3.4", path = "crates/mcp" }
 smg-external-router = { version = "0.1.0", path = "crates/external_router" }
-kv-index = { version = "1.5.0", path = "crates/kv_index" }
+kv-index = { version = "2.0.0", path = "crates/kv_index" }
 smg-radix-tree = { version = "0.1.0", path = "crates/radix_tree" }
 smg-data-connector = { version = "2.3.4", path = "crates/data_connector", package = "data-connector" }
 llm-multimodal = { version = "1.12.0", path = "crates/multimodal" }

bindings/python/src/lib.rs

@@ -550,7 +550,6 @@ struct Router {
     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,
@@ -674,14 +673,6 @@ 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,
@@ -931,7 +922,6 @@ impl Router {
                 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)
@@ -1218,7 +1208,6 @@ impl Router {
         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,
@@ -1397,7 +1386,6 @@ impl Router {
         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,
@@ -1602,7 +1590,6 @@ impl Router {
             prefill_queue_size,
             prefill_queue_timeout_secs,
             worker_overload_shed,
-            kv_index,
             worker_stall_secs,
             worker_wedge_secs,
             worker_warmup_secs,

bindings/python/src/smg/router_args.py

@@ -296,9 +296,6 @@ class RouterArgs:
     # Refuse with a 503 when every worker a request could use is overloaded,
     # instead of routing it to the least-loaded one
     worker_overload_shed: bool = False
-    # The event-driven KV index behind cache-aware routing: "positional"
-    # (one entry per block position) or "chain" (chains as runs)
-    kv_index: str = "positional"
     # Liveness: seconds without contact after a transport failure before a
     # worker is vetoed; seconds without progress before a loaded one is wedged
     worker_stall_secs: int = 2
@@ -688,16 +685,6 @@ def add_cli_args(
                 " and dispatch. Off by default."
             ),
         )
-        routing_group.add_argument(
-            f"--{prefix}kv-index",
-            type=str,
-            default=RouterArgs.kv_index,
-            help=(
-                "The event-driven KV index behind cache-aware routing: 'positional'"
-                " (one entry per block position, the default) or 'chain' (chains as"
-                " runs with per-run worker coverage; lock-free, store-free lookups)."
-            ),
-        )
         routing_group.add_argument(
             f"--{prefix}worker-stall-secs",
             type=int,

bindings/python/tests/test_arg_parser.py

@@ -705,14 +705,12 @@ def test_parse_overload_protection_and_monitoring_flags(self):
         disabled = parse_router_args(["--disable-worker-overload-protection"])
         assert disabled.worker_overload_protection is False
 
-    def test_parse_overload_shed_liveness_warmup_index_and_selection_flags(self):
+    def test_parse_overload_shed_liveness_warmup_and_selection_flags(self):
         """Every RouterConfig field the Rust CLI exposes reaches RouterArgs,
         with the CLI's defaults when the flags are absent."""
         router_args = parse_router_args(
             [
                 "--worker-overload-shed",
-                "--kv-index",
-                "chain",
                 "--worker-stall-secs",
                 "5",
                 "--worker-wedge-secs",
@@ -734,7 +732,6 @@ def test_parse_overload_shed_liveness_warmup_index_and_selection_flags(self):
             ]
         )
         assert router_args.worker_overload_shed is True
-        assert router_args.kv_index == "chain"
         assert router_args.worker_stall_secs == 5
         assert router_args.worker_wedge_secs == 7
         assert router_args.worker_warmup_secs == 30
@@ -747,7 +744,6 @@ def test_parse_overload_shed_liveness_warmup_index_and_selection_flags(self):
 
         defaults = parse_router_args([])
         assert defaults.worker_overload_shed is False
-        assert defaults.kv_index == "positional"
         assert defaults.worker_stall_secs == 2
         assert defaults.worker_wedge_secs == 3
         assert defaults.worker_warmup_secs == 60
@@ -1646,7 +1642,6 @@ class TestRouterArgsFieldOrder:
         "prefill_queue_timeout_secs",
         "discovery_provider",
         "worker_overload_shed",
-        "kv_index",
         "worker_stall_secs",
         "worker_wedge_secs",
         "worker_warmup_secs",
@@ -1695,7 +1690,6 @@ def test_new_fields_appended_after_positional_reserve(self):
             "rl_control_timeout_secs",
             "rl_fanout_concurrency",
             "worker_overload_shed",
-            "kv_index",
             "worker_stall_secs",
             "worker_wedge_secs",
             "worker_warmup_secs",

bindings/python/tests/test_router_config.py

@@ -491,7 +491,6 @@ def test_appended_parameters_follow_prefill_queue_timeout(self):
         assert tail == [
             "prefill_queue_timeout_secs",
             "worker_overload_shed",
-            "kv_index",
             "worker_stall_secs",
             "worker_wedge_secs",
             "worker_warmup_secs",

crates/kv_index/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "kv-index"
-version = "1.5.0"
+version = "2.0.0"
 edition = "2021"
 description = "Radix tree implementations for prefix matching and cache-aware routing"
 license = "Apache-2.0"

crates/kv_index/README.md

@@ -12,14 +12,14 @@ changes around it. Benchmarks and how to run them: `benches/README.md`.
 
 ## The problem with one entry per block
 
-The positional indexer (`src/event_tree.rs`, the default) keys one entry per (position, content
-hash) with the set of holding workers, and a lookup probes one entry per request block. It is
-exact (its jump search was not: the first commits of this series fixed it against the reference
-indexer and made it verify every position), but exactness costs it one dependent memory access per
-block per lookup, the per-entry worker set grows with every holder, and every store or removal is
-a hash-map write per block under a shard lock. A fleet of engines sharing long prefixes makes all
-three worse at once: the same chain is stored by many workers, requests are hundreds of blocks
-long, and the engines evict and re-store at the rate they serve.
+The positional indexer this crate shipped until 1.5 (`PositionalIndexer`, removed in 2.0) keyed
+one entry per (position, content hash) with the set of holding workers, and a lookup probed one
+entry per request block. It was exact (its jump search was not: the first commits of this series
+fixed it against the reference indexer and made it verify every position), but exactness cost it
+one dependent memory access per block per lookup, the per-entry worker set grew with every holder,
+and every store or removal was a hash-map write per block under a shard lock. A fleet of engines
+sharing long prefixes makes all three worse at once: the same chain is stored by many workers,
+requests are hundreds of blocks long, and the engines evict and re-store at the rate they serve.
 
 ## The chain index
 
@@ -160,11 +160,10 @@ The servicer relays the engine's publisher into `SubscribeKvEvents` through one
 
 ## Routing and recovery in the gateway
 
-- **Two backends, one type.** `KvIndex` (`model_gateway/src/worker/kv_index_backend.rs`) holds the
-  positional indexer or the chain index, selected by `--kv-index {positional,chain}`; the monitor
-  and the policy see the operations they used before. An in-gateway exactness test replays the
-  recorded engine streams and a synthetic hole corpus through the monitor's own apply path into
-  both backends and the reference indexer.
+- **One index, one type.** `KvIndex` (`model_gateway/src/worker/kv_index_backend.rs`) holds the
+  chain index; the monitor and the policy see the operations they used before. An in-gateway
+  exactness test replays generated engine streams and a synthetic hole corpus through the
+  monitor's own apply path into the index and the reference indexer.
 - **Per-rank admission.** One cursor per (worker, dp_rank): contiguous batches apply, duplicates
   skip, a gap gets one replay request and is otherwise settled (a small one keeps the blocks and
   marks the rank degraded, a large one clears the worker), a publisher restart clears the worker's
@@ -194,8 +193,8 @@ The servicer relays the engine's publisher into `SubscribeKvEvents` through one
 
 ## Flags and settings
 
-Gateway: `--kv-index {positional,chain}` (default `positional`; `run` is accepted as a deprecated
-alias of `chain`); `--kv-indexer-ttl-secs`, `--kv-indexer-max-entries` (positional only);
+Gateway: `--kv-indexer-ttl-secs`, `--kv-indexer-max-entries` (deprecated, accepted and ignored:
+the chain index has no prune);
 `--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-overload-protection` (on), `--disable-worker-overload-protection`,

crates/kv_index/benches/README.md

@@ -46,8 +46,8 @@ the last completion, drain included; latency is the lookup service time measured
   `--payload-home issuer` builds those payloads on the issuing thread instead of the main thread,
   so on a multi-socket machine the payloads live on the issuer's socket.
 
-Backends (`--backend`): `positional` (`PositionalIndexer`), `chain` (`ChainIndex`, or
-`ShardedChainIndex` with `--shards N`; `run` is its deprecated spelling), `reference` (the
+Backends (`--backend`): `chain` (`ChainIndex`, or `ShardedChainIndex` with `--shards N`; the
+default), `reference` (the
 single-threaded exactness reference; small corpora only) and `null` (no indexer: the harness's own
 ceiling on a layout). A new index plugs in by implementing `ReplayBackend`.
 
@@ -84,7 +84,7 @@ export is deterministic: two exports of the same arguments hash identically.
 
 | Flag | Meaning |
 | --- | --- |
-| `--backend`, `--jump-size`, `--max-workers`, `--shards`, `--lane-memory inherit\|local` | The index under test; `--shards` places each pinned event lane's workers on the shard of the lane's NUMA node; `--lane-memory local` makes a pinned lane prefer its own node for what it allocates. |
+| `--backend`, `--max-workers`, `--shards`, `--lane-memory inherit\|local` | The index under test; `--shards` places each pinned event lane's workers on the shard of the lane's NUMA node; `--lane-memory local` makes a pinned lane prefer its own node for what it allocates. |
 | `--benchmark-duration-ms` or `--offered-block-ops-per-sec` | The window, or the offered rate that sets it. |
 | `--query-lanes`, `--event-lanes`, `--issuer-threads`, `--query-issuer-threads` | Lane and issuer counts. |
 | `--issuer-cpus`, `--query-issuer-cpus`, `--backend-cpus`, `--pin-event-lanes`, `--issuer-by-lane` | Placement. Issuer CPUs inside the lane set are refused. |
@@ -116,10 +116,10 @@ Every indexer answers the same question as the single-threaded `ReferenceIndexer
 (`src/reference.rs`): after any replay, the set of (worker, position, block) and every lookup
 score equal the reference's. That is checked three ways:
 
-- `tests/exactness_positional.rs` and `tests/exactness_chain.rs` replay seeded corpora (new
-  conversations, extensions, siblings diverging at any position, tail and middle evictions,
-  clears, worker removal and arrival) into the positional and the chain index beside the
-  reference, compare every lookup kind after every 256 events and the full block set at the end;
+- `tests/exactness_chain.rs` replays seeded corpora (new conversations, extensions, siblings
+  diverging at any position, tail and middle evictions, clears, worker removal and arrival) into
+  the chain index beside the reference, compares every lookup kind after every 256 events and
+  the full block set at the end;
   `KV_INDEX_EXACTNESS_EVENTS`, `KV_INDEX_EXACTNESS_SEED` and `KV_INDEX_EXACTNESS_SHARDS` scale,
   reseed and shard them;
 - `tests/concurrency_chain.rs` runs 16 event lanes and 4 readers with worker replacement and

crates/kv_index/benches/match_insert.rs

@@ -24,8 +24,7 @@ use std::{
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
 use kv_index::{
     compute_content_hash, compute_request_content_hashes, request_prefix_hashes, ChainBlockMap,
-    ChainIndex, ContentHash, PositionalIndexer, SequenceHash, ShardedChainIndex, StoredBlock,
-    TokenTree, Tree, WorkerBlockMap,
+    ChainIndex, ContentHash, SequenceHash, ShardedChainIndex, StoredBlock, TokenTree, Tree,
 };
 
 const TENANTS: usize = 64;
@@ -260,10 +259,10 @@ fn bench_event_path(c: &mut Criterion) {
 
     // An indexer with 64 workers each holding 32 blocks: the size check the
     // policy runs per request before taking the event-driven path.
-    let indexer = PositionalIndexer::new(64);
+    let indexer = ChainIndex::new();
     for w in 0..TENANTS {
         let worker_id = indexer.intern_worker(&tenant(w)).expect("worker id");
-        let mut worker_blocks = WorkerBlockMap::default();
+        let mut worker_blocks = ChainBlockMap::default();
         let blocks: Vec<StoredBlock> = (0..32usize)
             .map(|p| StoredBlock {
                 seq_hash: SequenceHash::from(mix((w as u64) << 8 | p as u64)),