[Router] Shard the cache-aware KV tree by chain root (#39167)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>

experimental/sgl-router/Cargo.lock

@@ -3304,6 +3304,7 @@ dependencies = [
  "reqwest",
  "rmp",
  "rmp-serde",
+ "rustc-hash 2.1.3",
  "serde",
  "serde_json",
  "sgl-kv-indexer",

experimental/sgl-router/Cargo.toml

@@ -83,6 +83,12 @@ uuid = { version = "1", features = ["v4"] }
 # `python/sglang/srt/disaggregation/kv_events.py`.
 parking_lot = "0.12"
 rmp-serde = "1"
+# FxHash for the KV-event radix tree's integer-keyed maps (i64 block
+# hashes, u64 node ids), off the routing hot path's SipHash cost. The
+# keys are client-derivable, not trusted; what bounds a collision attack
+# is that planting a key costs a real request — see the DoS note in
+# src/policies/kv_events/tree.rs.
+rustc-hash = "2"
 sha2 = "0.10"
 url = "2"
 zeromq = { version = "0.6", default-features = false, features = ["tokio-runtime", "tcp-transport"] }

experimental/sgl-router/benches/tree_lookup.rs

@@ -9,6 +9,17 @@
 //!
 //!   * `insert` — populate one worker's prefix.
 //!   * `match_prefix` — score an incoming request against the tree.
+//!   * `insert_continuation` — insert with `parent_hash = Some(..)`, which
+//!     is what the pump emits for every block after a sequence's first.
+//!     `HashTree::route_insert` resolves that parent across ALL shards
+//!     before writing one, so it costs strictly more than a `None` insert;
+//!     the paired `parent_none` case is the same block count without the
+//!     scan.
+//!   * `contended_match` — reader `match_prefix` throughput WHILE a
+//!     background writer hammers `insert` / `remove`. The case sharding
+//!     targets: under one process-wide lock every event write blocks every
+//!     routing read. Run against both writer shapes so the headline ratio
+//!     is not read off the cheapest possible writer.
 //!
 //! Output is `criterion`'s default (target/criterion/...). To run:
 //!
@@ -17,6 +28,10 @@
 //!
 //! See `BENCHMARKS.md` for the SMG↔sgl-router comparison table.
 
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::thread;
+
 use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
@@ -56,6 +71,42 @@ fn bench_insert(c: &mut Criterion) {
     group.finish();
 }
 
+/// Cost of an insert that carries a parent, against one that does not.
+///
+/// `route_insert` early-returns on `parent_hash = None` and writes a single
+/// shard; a `Some(p)` takes a read lock on every shard to find which one
+/// already holds `p`. The pump passes `Some` for every block after a
+/// sequence's first, so that is the steady-state write path, and without
+/// this case the suite would time only the one shape the pump rarely
+/// sends. The scan is real but small next to the descent it precedes
+/// (~36ns against ~800ns measured), and the unsharded tree shows no gap
+/// between the two cases at all, having no shards to scan.
+///
+/// Both cases re-insert blocks the tree already holds, which is idempotent
+/// — the tree neither grows nor needs teardown between iterations, so the
+/// two are timed on one prebuilt tree at equal block counts.
+fn bench_insert_continuation(c: &mut Criterion) {
+    let mut group = c.benchmark_group("hashtree_insert_continuation");
+    let tree = build_tree(64, 64, 0xDEADBEEF);
+    // Worker 0 and its chain, re-derived from `build_tree`'s seed.
+    let worker = KvWorkerId::new("http://w0:30000".to_string(), 0);
+    let chain: Vec<i64> = {
+        let mut rng = StdRng::seed_from_u64(0xDEADBEEF);
+        (0..64).map(|_| rng.gen::<i64>()).collect()
+    };
+    let (head, tail) = chain.split_at(32);
+    let parent_hash = head[head.len() - 1];
+
+    group.throughput(Throughput::Elements(tail.len() as u64));
+    group.bench_function("parent_none", |b| {
+        b.iter(|| tree.insert(&worker, None, black_box(head)))
+    });
+    group.bench_function("parent_some", |b| {
+        b.iter(|| tree.insert(&worker, Some(black_box(parent_hash)), black_box(tail)))
+    });
+    group.finish();
+}
+
 fn bench_match_prefix(c: &mut Criterion) {
     let mut group = c.benchmark_group("hashtree_match_prefix");
     // (workers, blocks_per_worker, query_len) cases that span the
@@ -71,9 +122,10 @@ fn bench_match_prefix(c: &mut Criterion) {
         let label = format!("w{workers}_bpw{bpw}_q{query_len}");
         group.throughput(Throughput::Elements(query_len as u64));
         let tree = build_tree(workers, bpw, 0xDEADBEEF);
-        // Pull one real worker's prefix so the query has a non-trivial
+        // Re-derive worker 0's prefix (the first `bpw` i64s `build_tree`
-        // partial match — closer to the production hot path.
+        // drew from this seed) so the bench times a real descent; fresh
-        let mut rng = StdRng::seed_from_u64(0x12345);
+        // randoms would miss at the root and time a single lookup.
+        let mut rng = StdRng::seed_from_u64(0xDEADBEEF);
         let probe: Vec<i64> = (0..query_len).map(|_| rng.gen::<i64>()).collect();
         group.bench_function(label, |b| {
             b.iter(|| {
@@ -85,5 +137,116 @@ fn bench_match_prefix(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_insert, bench_match_prefix);
+/// How the background writer shapes its inserts. A `Rooted` writer touches
+/// exactly one shard; a `Continuation` writer first resolves its parent
+/// across all of them, which is the shape the pump emits for every block
+/// after a sequence's first.
+///
+/// Both are measured because the cross-shard resolution looks like it
+/// should erase the sharding win and does not: `route_insert` scans under
+/// READ locks, which readers share, so only its single `write()` blocks
+/// anyone. Measured, the continuation shape costs the reader ~7% over the
+/// rooted one, against ~450x for removing the global lock. This case
+/// exists to keep that true — a future `route_insert` that took write
+/// locks to scan, or serialised the scan behind the reader path, would
+/// show up here and nowhere else.
+#[derive(Clone, Copy)]
+enum WriterShape {
+    Rooted,
+    Continuation,
+}
+
+impl WriterShape {
+    fn label(self) -> &'static str {
+        match self {
+            Self::Rooted => "reader_under_write_pressure",
+            Self::Continuation => "reader_under_continuation_write_pressure",
+        }
+    }
+}
+
+/// Reader `match_prefix` throughput while a background writer hammers
+/// `insert` / `remove` — the read-vs-write contention sharding is built
+/// for. A single global lock serialises the two paths and the reader rate
+/// collapses; sharded, the writer's churn leaves reads on other roots
+/// uncontended, to the extent the writer stays off their shards.
+fn bench_contended_match(c: &mut Criterion) {
+    for shape in [WriterShape::Rooted, WriterShape::Continuation] {
+        bench_contended_match_with(c, shape);
+    }
+}
+
+fn bench_contended_match_with(c: &mut Criterion, shape: WriterShape) {
+    let mut group = c.benchmark_group("hashtree_contended_match");
+    let tree = Arc::new(build_tree(64, 64, 0xDEADBEEF));
+    // Worker 0's chain, re-derived from the same seed, so the reader gets a
+    // non-trivial full match.
+    let warm_chain: Vec<i64> = {
+        let mut rng = StdRng::seed_from_u64(0xDEADBEEF);
+        (0..64).map(|_| rng.gen::<i64>()).collect()
+    };
+
+    // One background writer, insert + remove of a fresh 4-block chain per
+    // round, so it keeps taking write locks with the tree size bounded.
+    let stop = Arc::new(AtomicBool::new(false));
+    let writer = {
+        let tree = tree.clone();
+        let stop = stop.clone();
+        thread::spawn(move || {
+            let scratch = KvWorkerId::new("http://scratch:30000".to_string(), 0);
+            let mut round = 0i64;
+            while !stop.load(Ordering::Relaxed) {
+                // Cycled, so a long run cannot drift the scratch roots into
+                // the warm chain's space or overflow the multiply.
+                let base = 1_000_000 + (round % 100_000) * 7;
+                let chain = [base, base + 1, base + 2, base + 3];
+                match shape {
+                    WriterShape::Rooted => tree.insert(&scratch, None, &chain),
+                    // Root the chain, then extend it the way the pump does:
+                    // only a sequence's first block carries no parent.
+                    WriterShape::Continuation => {
+                        tree.insert(&scratch, None, &chain[..1]);
+                        tree.insert(&scratch, Some(base), &chain[1..]);
+                    }
+                }
+                tree.remove(&scratch, &chain);
+                round = round.wrapping_add(1);
+            }
+        })
+    };
+    // Signals the writer on the way out however we leave — a panic in the
+    // bench body unwinds past any explicit store and would otherwise leave
+    // the thread spinning for the rest of the process.
+    let _stop_writer = StopOnDrop(stop);
+
+    group.throughput(Throughput::Elements(warm_chain.len() as u64));
+    group.bench_function(shape.label(), |b| {
+        b.iter(|| {
+            let m = tree.match_prefix(None, black_box(&warm_chain));
+            black_box(m.matched_blocks)
+        });
+    });
+
+    drop(_stop_writer);
+    writer.join().expect("bench writer thread panicked");
+    group.finish();
+}
+
+/// Sets its flag on drop, so a background thread parked on it is stopped by
+/// an unwind as reliably as by the normal path.
+struct StopOnDrop(Arc<AtomicBool>);
+
+impl Drop for StopOnDrop {
+    fn drop(&mut self) {
+        self.0.store(true, Ordering::Relaxed);
+    }
+}
+
+criterion_group!(
+    benches,
+    bench_insert,
+    bench_insert_continuation,
+    bench_match_prefix,
+    bench_contended_match
+);
 criterion_main!(benches);

experimental/sgl-router/src/policies/kv_events/index.rs

@@ -679,7 +679,7 @@ mod tests {
 
         let m = tree.match_prefix(None, &[10, 20, 30]);
         assert_eq!(m.matched_blocks, 3);
-        assert!(m.workers().contains(&id), "tree must hold the worker");
+        assert!(m.holds(&id), "tree must hold the worker");
     }
 
     /// The pump must carry each event's `medium` into the tree. The engine's
@@ -725,7 +725,7 @@ mod tests {
 
         let m = tree.match_prefix(None, &[10, 20]);
         assert_eq!(m.matched_blocks, 2, "host copy keeps the block routable");
-        assert!(m.workers().contains(&id));
+        assert!(m.holds(&id));
         assert!(!m.device_workers().contains(&id), "device copy is gone");
     }
 
@@ -1161,4 +1161,57 @@ mod tests {
         assert_eq!(index.engine_load().expected_count(), 0);
         index.shutdown().await;
     }
+
+    /// `remove_worker` is the tree's SECOND writer: it clears every rank's
+    /// state from whatever task service discovery calls it on
+    /// (`workers/manager.rs`) while the pump is live and writing the same
+    /// tree. This pins the wiring that makes the collision
+    /// `HashTree::concurrent_writers_never_orphan_a_chain` covers reachable
+    /// at all — a refactor moving the clear onto the pump would retire it.
+    #[tokio::test]
+    async fn remove_worker_clears_the_tree_off_the_pump_task() {
+        let index = KvEventIndex::new();
+        let url = "http://127.0.0.1:59124";
+        let cfg = EventConfig {
+            host: "127.0.0.1".into(),
+            port_base: 59124,
+            topic: String::new(),
+            load_port_base: None,
+            load_topic: None,
+            block_size: 64,
+            dp_size: 2,
+            is_bigram: false,
+        };
+        index.add_worker(url, Some(cfg)).await;
+
+        // Stand in for events the pump already applied for both ranks.
+        let tree = index.tree();
+        let (r0, r1) = (worker_id(url, 0), worker_id(url, 1));
+        tree.insert(&r0, None, &[10, 20]);
+        tree.insert(&r1, None, &[30, 40]);
+        assert!(tree.match_prefix(None, &[10, 20]).holds(&r0));
+        assert!(tree.match_prefix(None, &[30, 40]).holds(&r1));
+
+        // Parked on its channel, not finished, so the clear below genuinely
+        // runs on a different task than the one applying events.
+        assert!(
+            index.pump.lock().as_ref().is_some_and(|h| !h.is_finished()),
+            "pump task must still be live, or this proves nothing about a second writer",
+        );
+
+        index.remove_worker(url).await;
+
+        assert_eq!(
+            tree.match_prefix(None, &[10, 20]).matched_blocks,
+            0,
+            "remove_worker must clear the removed worker's tree state",
+        );
+        assert_eq!(
+            tree.match_prefix(None, &[30, 40]).matched_blocks,
+            0,
+            "every dp rank of the removed worker must be cleared",
+        );
+        assert_eq!(tree.node_count(), 0, "cleared chains must prune");
+        index.shutdown().await;
+    }
 }

experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs

@@ -3,7 +3,7 @@
 
 //! Concurrent-mutation stress test for `HashTree`.
 //!
-//! The 19 inline tests in `policies::kv_events::tree` are all
+//! The inline tests in `policies::kv_events::tree` are all
 //! single-threaded.  Under production load, multiple worker subscribers
 //! drive `insert` / `remove` / `clear_worker` against the same tree from
 //! tokio worker threads while the chat handler simultaneously calls
@@ -155,7 +155,7 @@ fn match_prefix_is_consistent_with_concurrent_clear() {
         // consistent.
         if m.matched_blocks == chain.len() {
             assert!(
-                m.workers().contains(&w),
+                m.holds(&w),
                 "full match must include worker; got {:?}",
                 m.workers(),
             );
@@ -165,3 +165,102 @@ fn match_prefix_is_consistent_with_concurrent_clear() {
     stop.store(true, std::sync::atomic::Ordering::Relaxed);
     mutator.join().unwrap();
 }
+
+/// Reader storm concurrent with a writer hammering insert/remove on
+/// DISTINCT chain roots — the pattern sharding targets. Asserts
+/// CORRECTNESS under that contention: warm chains are pre-inserted and
+/// never removed, so a reader must always get a full match with the warm
+/// worker present, and the writer's scratch chains are fully removed each
+/// round, so after join only the warm chains remain.
+///
+/// No sleeps, no wall-clock — the readers run a fixed number of bounded
+/// iterations and the writer churns until they are done, so nothing here
+/// can flake on timing.
+#[test]
+fn readers_unaffected_by_writer_on_distinct_roots() {
+    let tree = Arc::new(HashTree::new());
+
+    // Fixed chains the readers query and the writer never touches, on
+    // distinct roots (1_000 apart) so they spread across shards.
+    let warm = KvWorkerId {
+        url: "http://warm:30000".into(),
+        dp_rank: 0,
+    };
+    const WARM_CHAINS: i64 = 16;
+    let warm_chain = |c: i64| -> Vec<i64> { vec![c * 1_000, c * 1_000 + 1, c * 1_000 + 2] };
+    for c in 0..WARM_CHAINS {
+        tree.insert(&warm, None, &warm_chain(c));
+    }
+
+    let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
+
+    // Scratch chains on distinct roots far from the warm space, insert then
+    // immediate remove. Runs until the readers are done rather than for a
+    // fixed count the writer could burn through early, leaving the test with
+    // no contention at all. Every round is a complete pair, so stopping at
+    // any point leaves no residue.
+    let writer = {
+        let tree = tree.clone();
+        let stop = stop.clone();
+        thread::spawn(move || {
+            let scratch = KvWorkerId {
+                url: "http://scratch:30000".into(),
+                dp_rank: 0,
+            };
+            let mut round = 0i64;
+            while !stop.load(std::sync::atomic::Ordering::Relaxed) {
+                // Cycled so the arithmetic cannot drift into the warm space.
+                let base = 1_000_000 + (round % 100_000) * 7;
+                let chain = [base, base + 1, base + 2, base + 3];
+                tree.insert(&scratch, None, &chain);
+                tree.remove(&scratch, &chain);
+                round = round.wrapping_add(1);
+            }
+        })
+    };
+
+    // Each reader asserts the warm worker is present at full depth on every
+    // warm chain, regardless of writer churn.
+    let mut readers = Vec::new();
+    for _ in 0..4 {
+        let tree = tree.clone();
+        let warm = warm.clone();
+        readers.push(thread::spawn(move || {
+            for _ in 0..2_000 {
+                for c in 0..WARM_CHAINS {
+                    let chain = warm_chain(c);
+                    let m = tree.match_prefix(None, &chain);
+                    assert_eq!(
+                        m.matched_blocks,
+                        chain.len(),
+                        "warm chain {c} must always fully match despite writer churn",
+                    );
+                    assert!(
+                        m.holds(&warm),
+                        "warm worker must always hold its own untouched chain",
+                    );
+                }
+            }
+        }));
+    }
+
+    for r in readers {
+        r.join().expect("reader thread panicked under contention");
+    }
+    stop.store(true, std::sync::atomic::Ordering::Relaxed);
+    writer
+        .join()
+        .expect("writer thread panicked under contention");
+
+    // Only the warm chains remain: 16 chains x 3 nodes = 48 non-root nodes.
+    assert_eq!(
+        tree.node_count(),
+        (WARM_CHAINS * 3) as usize,
+        "writer's scratch churn must leave no residual nodes",
+    );
+    for c in 0..WARM_CHAINS {
+        let m = tree.match_prefix(None, &warm_chain(c));
+        assert_eq!(m.matched_blocks, 3);
+        assert!(m.holds(&warm));
+    }
+}

experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs

@@ -84,8 +84,8 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
         let mb = router_b.tree().match_prefix(None, &hashes);
         let converged = ma.matched_blocks == target
             && mb.matched_blocks == target
-            && ma.workers().contains(&key)
+            && ma.holds(&key)
-            && mb.workers().contains(&key);
+            && mb.holds(&key);
         if converged {
             // Both trees agree on count AND on the worker that holds the
             // prefix. This is what the Radix Tree provider reads to
@@ -253,32 +253,32 @@ async fn two_subscribers_merge_events_from_two_publishers() {
             && ay.matched_blocks == target_y
             && bx.matched_blocks == target_x
             && by.matched_blocks == target_y
-            && ax.workers().contains(&key_x)
+            && ax.holds(&key_x)
-            && ay.workers().contains(&key_y)
+            && ay.holds(&key_y)
-            && bx.workers().contains(&key_x)
+            && bx.holds(&key_x)
-            && by.workers().contains(&key_y);
+            && by.holds(&key_y);
         if converged {
             // Negative attribution: prefix X must not be attributed to
             // worker_y in either tree, and vice versa. A regression that
             // keyed events by arriving socket rather than announced
             // worker URL would set BOTH worker keys on each prefix.
             assert!(
-                !ax.workers().contains(&key_y),
+                !ax.holds(&key_y),
                 "router_a cross-attributed worker_y to prefix X: {:?}",
                 ax.workers(),
             );
             assert!(
-                !ay.workers().contains(&key_x),
+                !ay.holds(&key_x),
                 "router_a cross-attributed worker_x to prefix Y: {:?}",
                 ay.workers(),
             );
             assert!(
-                !bx.workers().contains(&key_y),
+                !bx.holds(&key_y),
                 "router_b cross-attributed worker_y to prefix X: {:?}",
                 bx.workers(),
             );
             assert!(
-                !by.workers().contains(&key_x),
+                !by.holds(&key_x),
                 "router_b cross-attributed worker_x to prefix Y: {:?}",
                 by.workers(),
             );