feat(radix-index): shared radix membership index service (WIP experiment)

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["model_gateway", "crates/protocols", "crates/reasoning_parser", "crates/tool_parser", "crates/workflow", "crates/tokenizer", "crates/auth", "crates/mcp", "crates/kv_index", "crates/data_connector", "crates/multimodal", "crates/mm_rdma", "crates/wasm", "crates/mesh", "crates/grpc_client", "crates/engine_zmq_client", "bindings/python", "bindings/golang", "clients/rust", "clients/openapi-gen", "crates/mock_worker"]
+members = ["model_gateway", "crates/protocols", "crates/reasoning_parser", "crates/tool_parser", "crates/workflow", "crates/tokenizer", "crates/auth", "crates/mcp", "crates/kv_index", "crates/data_connector", "crates/multimodal", "crates/mm_rdma", "crates/wasm", "crates/mesh", "crates/grpc_client", "crates/engine_zmq_client", "bindings/python", "bindings/golang", "clients/rust", "clients/openapi-gen", "crates/mock_worker", "crates/radix_index", "crates/radix_tree"]
 resolver = "2"
 
 [workspace.dependencies]

crates/kv_index/src/event_tree.rs

@@ -36,6 +36,18 @@ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
 /// Seed for XXH3 hashing.
 pub const XXH3_SEED: u64 = 1337;
 
+/// Rolling prefix hash over content hashes: `XXH3(prev || current)`, the
+/// same chaining `PositionalIndexer` computes internally. Exported so
+/// out-of-process publishers (the radix index service's placement feed)
+/// can synthesize byte-identical position chains for identical prefixes.
+#[inline]
+pub fn chain_prefix_hash(prev: SequenceHash, current: ContentHash) -> SequenceHash {
+    let mut bytes = [0u8; 16];
+    bytes[..8].copy_from_slice(&prev.0.to_le_bytes());
+    bytes[8..].copy_from_slice(&current.0.to_le_bytes());
+    SequenceHash(xxhash_rust::xxh3::xxh3_64_with_seed(&bytes, XXH3_SEED))
+}
+
 /// Shard count for the main index DashMap.
 /// Tuned iteratively — higher values reduce per-shard contention under concurrent
 /// reads+writes at the cost of more memory for shard locks.
@@ -187,6 +199,36 @@ pub fn compute_request_content_hashes(tokens: &[u32], block_size: usize) -> Vec<
         .collect()
 }
 
+/// Content hash of a raw byte block (position-independent), the byte-mode
+/// analogue of [`compute_content_hash`]. Same XXH3-64 hasher and seed,
+/// fed the bytes directly, so the two symbol kinds live in disjoint
+/// content-hash spaces without any shared state.
+pub fn compute_byte_content_hash(bytes: &[u8]) -> ContentHash {
+    use std::hash::Hasher;
+    let mut hasher = xxhash_rust::xxh3::Xxh3::with_seed(XXH3_SEED);
+    hasher.write(bytes);
+    ContentHash(hasher.finish())
+}
+
+/// Chunk raw request bytes by `byte_block` and compute a [`ContentHash`]
+/// per full block — the string-mode (`SymbolKind::Bytes`) analogue of
+/// [`compute_request_content_hashes`], for routing HTTP requests that
+/// carry text but no tokens. Partial trailing chunks are discarded (a
+/// worker only caches a full block); empty if `byte_block` is 0.
+pub fn compute_request_byte_content_hashes(bytes: &[u8], byte_block: usize) -> Vec<ContentHash> {
+    if byte_block == 0 {
+        tracing::warn!(
+            "compute_request_byte_content_hashes called with byte_block=0, returning empty"
+        );
+        return Vec::new();
+    }
+    bytes
+        .chunks(byte_block)
+        .filter(|chunk| chunk.len() == byte_block)
+        .map(compute_byte_content_hash)
+        .collect()
+}
+
 // ---------------------------------------------------------------------------
 // SeqEntry: optimizes for the common case (one seq_hash per position+content)
 // ---------------------------------------------------------------------------
@@ -2844,4 +2886,41 @@ mod tests {
         assert_eq!(indexer.current_size(), 0);
         assert_eq!(indexer.entry_count(), 0);
     }
+
+    #[test]
+    fn byte_content_hashes_are_deterministic_and_prefix_stable() {
+        // Determinism: same bytes → same chain.
+        let a = compute_request_byte_content_hashes(b"hello world foo", 4);
+        let b = compute_request_byte_content_hashes(b"hello world foo", 4);
+        assert_eq!(a, b);
+
+        // Full blocks only: 15 bytes / 4 = 3 full blocks, trailing 3 dropped.
+        assert_eq!(a.len(), 3);
+
+        // Prefix stability: a longer request that shares a byte prefix
+        // produces the SAME leading block hashes — this is the property
+        // string-mode affinity relies on.
+        let longer = compute_request_byte_content_hashes(b"hello world foo BAR BAZ", 4);
+        assert_eq!(&longer[..3], &a[..]);
+
+        // A differing byte in the first block changes that block's hash
+        // (and only shared-prefix blocks stay equal).
+        let diverged = compute_request_byte_content_hashes(b"Hello world foo", 4);
+        assert_ne!(diverged[0], a[0]);
+        assert_eq!(&diverged[1..], &a[1..]);
+    }
+
+    #[test]
+    fn byte_content_hashes_empty_on_zero_block() {
+        assert!(compute_request_byte_content_hashes(b"anything", 0).is_empty());
+        assert!(compute_request_byte_content_hashes(b"", 4).is_empty());
+    }
+
+    /// The token hasher's block_size==0 guard: return empty rather than
+    /// panicking on chunks(0). Mirrors the byte-mode guard.
+    #[test]
+    fn token_content_hashes_empty_on_zero_block() {
+        assert!(compute_request_content_hashes(&[1, 2, 3], 0).is_empty());
+        assert!(compute_request_content_hashes(&[], 4).is_empty());
+    }
 }

crates/kv_index/src/lib.rs

@@ -20,9 +20,10 @@ mod token_tree;
 
 pub use common::{MatchResult, TenantId};
 pub use event_tree::{
-    compute_content_hash, compute_request_content_hashes, ApplyError, ContentHash, OverlapScores,
-    PositionalIndexer, PruneStats, SequenceHash, StoredBlock, WorkerBlockMap, WorkerId,
-    WorkerIdExhausted,
+    chain_prefix_hash, compute_byte_content_hash, compute_content_hash,
+    compute_request_byte_content_hashes, compute_request_content_hashes, ApplyError, ContentHash,
+    OverlapScores, PositionalIndexer, PruneStats, SequenceHash, StoredBlock, WorkerBlockMap,
+    WorkerId, WorkerIdExhausted, XXH3_SEED,
 };
 pub use path_hash::{hash_node_path, hash_token_path, GLOBAL_EVICTION_HASH};
 // Re-export under names matching old tree.rs API for easier migration

crates/radix_index/Cargo.toml

@@ -0,0 +1,49 @@
+[package]
+name = "radix-index"
+version = "0.1.0"
+edition = "2021"
+license = "Apache-2.0"
+description = "Generic radix membership index service: block-quantized prefix chains, per-holder overlap scoring, pub in / sub out"
+
+[lib]
+name = "radix_index"
+path = "src/lib.rs"
+
+[[bin]]
+name = "radix-index-service"
+path = "src/bin/service.rs"
+
+[[bin]]
+name = "radix-index-bridge"
+path = "src/bin/bridge.rs"
+
+[[bin]]
+name = "radix-index-bench"
+path = "src/bin/bench.rs"
+
+[[bin]]
+name = "radix-index-loadbench"
+path = "src/bin/loadbench.rs"
+
+[dependencies]
+radix-tree = { path = "../radix_tree" }
+xxhash-rust = { version = "0.8", features = ["xxh3"] }
+smg-grpc-client = { path = "../grpc_client" }
+tonic.workspace = true
+tonic-prost.workspace = true
+prost.workspace = true
+tokio = { workspace = true, features = ["full"] }
+tokio-stream = "0.1"
+futures.workspace = true
+tracing.workspace = true
+tracing-subscriber = { workspace = true }
+rustc-hash = "2"
+
+[dev-dependencies]
+kv-index = { path = "../kv_index" }
+rand = { workspace = true }
+mock-worker = { path = "../mock_worker" }
+portpicker = "0.1"
+
+[build-dependencies]
+tonic-prost-build = "0.14.6"

crates/radix_index/README.md

@@ -0,0 +1,96 @@
+# radix-index
+
+A shared radix membership index for cache-aware routing: gateways ask
+"which worker already holds the longest prefix of this request?" without
+each gateway building and syncing its own tree.
+
+The data structure is SMG `kv_index`'s block-quantized positional
+index (a direct import — whether the service should instead get a
+ground-up generic radix tree is an open evaluation). This crate wraps
+it in a keyspace-partitioned engine, a gRPC surface, and the
+client/bridge pieces that feed and query it.
+
+## Interface: two verbs
+
+- **Publish** (client-streamed `Update`s, acked): a holder's block-hash
+  chain changed. Two feeds share the verb:
+  - **Event feed** — a bridge subscribes to engine KV events
+    (`SubscribeKvEvents`) and forwards Stored/Removed/Cleared batches,
+    sequenced per holder with epoch bumps on gap or backend loss.
+    Eviction is *observed*, so index state tracks engine truth.
+  - **Placement feed** — a gateway publishes "this request's chain now
+    (probably) resides on that worker" after each completed request
+    (`seq=0`, content-idempotent). No engine cooperation needed — this
+    is the path for engines/modes with no KV event stream. Inferred
+    state is bounded by idle TTL + per-holder capacity with tail-first
+    (prefix-closed) eviction.
+  The first sequenced or Removed-bearing update marks a holder
+  *event-fed*: placements for it are ignored from then on, so the
+  precise feed always wins.
+- **Subscribe** (bidirectional query stream): content-hash chain in,
+  per-holder matched-block counts out. The gateway enforces its own
+  deadline (2 ms in SMG) and falls back to its local policy on a miss —
+  answers are advisory, never load-bearing for correctness.
+
+`Pull` streams the whole state as synthetic `Update`s; a starting
+replica bootstraps from a sibling with it before declaring ready.
+
+## Replication: copy, don't agree
+
+No consensus. Writes are per-holder sequenced (event feed) or
+content-idempotent (placement feed), so replicas converge by applying
+the same updates in any interleaving. A replica relays each accepted
+`Publish` to its `--peers` best-effort; a wedged peer drops relayed
+updates rather than wedging ingest, and TTL plus re-placement plus
+bootstrap bound the divergence. Gateways stay stupid: one endpoint, no
+fan-out, reconnect on failure.
+
+## Keyspaces
+
+State is partitioned by `(model, symbol_kind, block_size)`. TOKENS at
+the engine page size is what SMG uses today; BYTES exists for text-mode
+(HTTP) feeds that hash normalized bytes instead of token ids.
+
+## Binaries
+
+- `radix-index-service` — the server. Flags:
+
+  | flag | default | meaning |
+  |---|---|---|
+  | `--bind` | `127.0.0.1` | listen address (set `0.0.0.0` in k8s) |
+  | `--port` | `40000` | gRPC port |
+  | `--metrics-port` | off | admin plane: `/metrics`, `/healthz`, `/readyz` |
+  | `--peers` | none | sibling replicas to relay Publishes to (comma-separated URLs) |
+  | `--bootstrap-from` | none | sibling to Pull state from before serving |
+  | `--inferred-ttl-secs` | `180` | idle TTL for placement-fed holders |
+  | `--event-ttl-secs` | `1800` | liveness backstop for EVENT-fed holders: silence past this soft-retires the holder (a lost departure signal must not leak it); `0` disables |
+  | `--default-capacity-blocks` | unbounded | RUNAWAY PROTECTION for holders that never sent `Added` — the index truncates only past 2x this value. Leave unbounded or set well above worker KV size: the placement feed carries no removal signal, so an index that races the worker's own eviction under-matches (measured); idle TTL is the freshness bound. |
+  | `--sweep-interval-secs` | `5` | idle-sweep cadence |
+  | `--apply-delay-stored-ms` / `--apply-delay-removed-ms` | `0` | staleness injection (experiments only) |
+
+  Stops gracefully on SIGTERM/ctrl-c. `/readyz` answers 503 until the
+  bootstrap pull completes — point the k8s readiness probe at it.
+
+- `radix-index-bridge` — per-fleet event bridge: engine KV event
+  streams in, index Updates out. Flags: `--workers` (comma-separated
+  worker URLs), `--index` (service URL), `--model`, `--block-size`
+  (MUST match the gateway's `--kv-indexer-block-size` — the keyspace
+  key includes it, and a mismatch silently splits the fleet's state
+  into two keyspaces; both default to the same shared value).
+- `radix-index-bench` — apply/query throughput and memory-per-entry
+  microbench.
+
+A reference StatefulSet lives in [`deploy/statefulset.yaml`](deploy/statefulset.yaml).
+
+## Gateway side (SMG)
+
+`--kv-indexer-url` + `--kv-indexer-block-size` on the gateway enable
+the remote path: a routing-time overlap prefetch (2 ms deadline,
+fast-fail while disconnected) feeding the cache-aware policy, and a
+placement publish of the prompt⊕output chain after each completed
+request. Flag off = every code path byte-identical to local behavior.
+
+## Metrics
+
+`/metrics` (Prometheus text): `radix_index_{keyspaces,holders,event_fed_holders,dropped_holders,blocks}`
+gauges and `radix_index_{applies,queries,relay_dropped}_total` counters.

crates/radix_index/build.rs

@@ -0,0 +1,9 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    println!("cargo:rerun-if-changed=proto/radix_index.proto");
+    tonic_prost_build::configure()
+        .build_server(true)
+        .build_client(true)
+        .protoc_arg("--experimental_allow_proto3_optional")
+        .compile_protos(&["proto/radix_index.proto"], &["proto"])?;
+    Ok(())
+}