feat(kv-router): expose per-rank resident block counts to selection plugins

docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx

@@ -370,7 +370,7 @@ Each worker policy is constructed once per routing partition. An empty bundle re
 
 Worker identity and ordinary request metadata are always available. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests.
 
-Routing cache resources start only if a component requests `CACHE`. `serve_indexer` and `enable_session_prefix_index` also require `CACHE`.
+Routing cache resources start only if a component requests `CACHE`. `serve_indexer`, `enable_session_prefix_index`, and `RESIDENT_BLOCKS` also require `CACHE`.
 
 ### Request Context
 
@@ -446,13 +446,21 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi
 | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks |
 | `LOAD` | `active_requests()` | Requests currently active on the worker |
 | `PREFERRED_TAINT` | `preferred_taint_multiplier()` | Optional cost multiplier from matching preferred routing taints |
+| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; current as of the last event the indexer applied and not clamped to capacity |
 
 The GPU, CPU, and disk overlap accessors return raw tier counts. A missing tier or worker entry is zero; Dynamo does not substitute its weighted effective-overlap estimate. Each custom scorer chooses how to combine the raw counts.
 
 A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. A picker reads index-aligned views through `WorkerInputView`; its cache view supports `get(row)` and `iter()`. Each component must declare every group it reads; another component requesting a group does not grant access.
 
 `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules.
 
+`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The count is unweighted, not clamped to `total_kv_blocks`, and measures occupancy rather than reusable cache: blocks that prefix matching can no longer reach, such as descendants of an evicted block, count until their own remove events arrive, and a lost remove event inflates the count until the rank is removed. A count is current as of the last event its indexer thread applied, the same freshness as the `CACHE` overlap scores, so it can trail queued events and change between two reads in one selection. A rank reads `None` until its indexer thread has reached the policy's request and holds a block for it, `0` after a clear, and `None` again once the rank or its worker is removed. A component also reads `None` when it did not request `WorkerInputs::RESIDENT_BLOCKS`, with a remote indexer, or after an indexer thread stops, so a worker whose KV events never arrive never looks empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this for every policy it builds for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason.
+
+What the count contains, and whether router replicas agree on it, depends on the routing mode:
+
+- **KV events:** the device-tier blocks the indexer currently tracks from engine events. Every router replica sees the same counts.
+- **Approximate routing:** the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. Each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it.
+
 ## Link the Policy Into Dynamo
 
 Both paths use the same policy crate, catalog, and YAML file. Choose the process that owns worker selection.

examples/router/custom-policy-example/README.md

@@ -24,6 +24,8 @@ Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (fi
 
 Preferred routing taints are optional candidate metadata. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading `preferred_taint_multiplier()` from a candidate; otherwise, that component receives `None`, even if another component requested it. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules.
 
+Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer holds for that worker rank, including cached blocks no running request uses, as of the last event the indexer applied. Without the declaration, with a remote indexer, or for a rank the indexer has held no blocks for or has removed, it returns `None`; a cleared rank returns `0`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`.
+
 ## Pick a Starting Point
 
 | Crate | Use it for |

lib/bench/kv_router/INDEXER_BENCH.md

@@ -158,6 +158,16 @@ print(f"runs={len(values)} mean={statistics.fmean(values):.3f} median={statistic
 PY
 ```
 
+### Resident block tracking cost
+
+To measure what `RESIDENT_BLOCKS` tracking costs the event threads, run the single
+CRTC trial with `--track-resident-blocks`, which keeps live per-rank block counters,
+and compare `achieved_block_ops_per_sec` against an untracked run. Add
+`--resident-blocks-reader-cpu <CPU>` to also read every counter in a loop from a
+thread pinned to that CPU, as selections do, so the result includes the cache-line
+traffic between event threads and readers. Keep that CPU outside the backend and
+issuer masks.
+
 ### Sweep
 
 Sweep cells use the same completion-correct implementation. Each cell regenerates

lib/bench/kv_router/mooncake_bench.rs

@@ -23,6 +23,7 @@ use mooncake_shared::{
 };
 use sha2::{Digest, Sha256};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 #[cfg(target_os = "linux")]
 const PRE_RUN_QUIESCENCE_MS: u64 = 5_000;
@@ -130,6 +131,16 @@ struct Args {
     #[clap(long, default_value = "1")]
     benchmark_runs: usize,
 
+    /// Keep live per-rank resident block counts on the event threads, as when a
+    /// selection policy reads `RESIDENT_BLOCKS` (concurrent-radix-tree-compressed only).
+    #[clap(long)]
+    track_resident_blocks: bool,
+
+    /// Also read every resident block counter in a loop from a thread pinned to this
+    /// CPU, as selections do. Implies `--track-resident-blocks`.
+    #[clap(long)]
+    resident_blocks_reader_cpu: Option<usize>,
+
     /// Indexer backend to benchmark. Defaults to concurrent-radix-tree-compressed
     /// with `--num-event-workers` event threads.
     #[clap(subcommand)]
@@ -270,6 +281,11 @@ async fn run_open_loop_for_config(
     let metrics = || Some(Arc::new(KvIndexerMetrics::new_unregistered()));
     let open_config = parse_open_loop_config(args)?;
     pin_current_thread_to_cpus(&open_config.backend_cpus)?;
+    let track_resident_blocks =
+        args.track_resident_blocks || args.resident_blocks_reader_cpu.is_some();
+    if track_resident_blocks && config.kind != MooncakeIndexerKind::ConcurrentRadixTreeCompressed {
+        anyhow::bail!("resident block tracking applies only to concurrent-radix-tree-compressed");
+    }
 
     match config.kind {
         MooncakeIndexerKind::NestedMap => {
@@ -288,7 +304,28 @@ async fn run_open_loop_for_config(
                 args.common.block_size,
                 metrics(),
             ));
-            run_backend(config.short_name(), indexer, trial, open_config).await
+            let counts = track_resident_blocks.then(|| indexer.bench_track_resident_blocks());
+            let stop = Arc::new(AtomicBool::new(false));
+            let reader = counts
+                .zip(args.resident_blocks_reader_cpu)
+                .map(|(counts, cpu)| {
+                    let stop = stop.clone();
+                    std::thread::spawn(move || {
+                        pin_current_thread_to_cpus(&[cpu])?;
+                        while !stop.load(Ordering::Relaxed) {
+                            std::hint::black_box(counts.bench_sum());
+                        }
+                        std::io::Result::Ok(())
+                    })
+                });
+            let result = run_backend(config.short_name(), indexer, trial, open_config).await;
+            stop.store(true, Ordering::Relaxed);
+            if let Some(reader) = reader {
+                reader
+                    .join()
+                    .map_err(|_| anyhow::anyhow!("resident block reader panicked"))??;
+            }
+            result
         }
         MooncakeIndexerKind::RadixTree | MooncakeIndexerKind::BranchShardedCrtc => {
             anyhow::bail!(

lib/kv-router/benches/worker_selection.rs

@@ -6,11 +6,18 @@
 //! Run with: `cargo bench -p dynamo-kv-router --bench worker_selection`
 
 use std::collections::{HashMap, HashSet};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 
 use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
+use dynamo_kv_router::indexer::{
+    KvIndexerInterface, ThreadPoolIndexer,
+    concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed,
+};
 use dynamo_kv_router::protocols::{
-    RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank,
+    ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, KvCacheStoreData,
+    KvCacheStoredBlockData, LocalBlockHash, RouterEvent, RoutingConstraints, WorkerConfigLike,
+    WorkerId, WorkerWithDpRank, compute_seq_hash_for_block,
 };
 use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode};
 use dynamo_kv_router::{
@@ -125,6 +132,42 @@ impl WorkerScorer for IgnoringPreferenceScorer {
     }
 }
 
+/// `IgnoringPreferenceScorer` plus each candidate's resident block count.
+struct ResidentBlocksScorer;
+
+impl WorkerScorer for ResidentBlocksScorer {
+    fn required_worker_inputs(&self) -> WorkerInputs {
+        WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::RESIDENT_BLOCKS
+    }
+
+    fn score(
+        &mut self,
+        context: &WorkerSelectionContext<'_>,
+        candidates: WorkerCandidates<'_>,
+        costs: &mut [f64],
+    ) -> Result<(), WorkerSelectionPolicyError> {
+        for (candidate, cost) in candidates.iter().zip(costs) {
+            let cache = candidate
+                .cache()
+                .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?;
+            let load = candidate
+                .load()
+                .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?;
+            // Real policies should treat None (no fresh counts) as neutral; failing keeps the bench honest.
+            let resident_blocks = candidate.resident_blocks().ok_or_else(|| {
+                WorkerSelectionPolicyError::failed("resident blocks input unavailable")
+            })?;
+            let uncached_blocks =
+                (context.request_blocks() as f64 - cache.device_overlap_blocks()).max(0.0);
+            let load_blocks = load.active_prefill_tokens() as f64 / context.block_size() as f64
+                + load.decode_cost_blocks()
+                + load.active_requests() as f64;
+            *cost = uncached_blocks + load_blocks + resident_blocks as f64;
+        }
+        Ok(())
+    }
+}
+
 struct LowestCostPicker;
 
 impl WorkerPicker for LowestCostPicker {
@@ -371,6 +414,158 @@ fn unused_preferred_taint_metadata(c: &mut Criterion) {
     group.finish();
 }
 
+fn resident_blocks_read(c: &mut Criterion) {
+    let config = KvRouterConfig {
+        router_temperature: 0.0,
+        ..Default::default()
+    };
+    let mut group = c.benchmark_group("resident_blocks_read");
+    group.warm_up_time(Duration::from_secs(2));
+    group.measurement_time(Duration::from_secs(5));
+    group.sample_size(50);
+
+    for worker_count in [2, 32, 1_024, 10_000] {
+        let (workers, request) = fixture(worker_count);
+        group.throughput(Throughput::Elements(worker_count as u64));
+        let without = WorkerSelectionPolicy::new(
+            config.clone(),
+            "prefill",
+            vec![Box::new(IgnoringPreferenceScorer)],
+            Box::new(LowestCostPicker),
+        );
+        let (with, writer) = WorkerSelectionPolicy::new(
+            config.clone(),
+            "prefill",
+            vec![Box::new(ResidentBlocksScorer)],
+            Box::new(LowestCostPicker),
+        )
+        .with_fixed_resident_block_counts(
+            workers
+                .keys()
+                .map(|&worker_id| (WorkerWithDpRank::from_worker_id(worker_id), 4_096)),
+        );
+        for (name, policy, writing) in [
+            ("without", &without, false),
+            ("with", &with, false),
+            // Another thread storing to every counter in a loop: the worst case for the
+            // cross-core misses event-thread writes cause readers.
+            ("with_writer", &with, true),
+        ] {
+            let stop = AtomicBool::new(false);
+            std::thread::scope(|scope| {
+                if writing {
+                    scope.spawn(|| {
+                        let mut blocks = 0;
+                        while !stop.load(Ordering::Relaxed) {
+                            writer.store_all(blocks);
+                            blocks += 1;
+                        }
+                    });
+                }
+                group.bench_with_input(
+                    BenchmarkId::new(name, worker_count),
+                    &worker_count,
+                    |b, _| {
+                        b.iter(|| {
+                            black_box(
+                                policy
+                                    .select_worker(WorkerSelectionInput::configured(
+                                        black_box(&workers),
+                                        black_box(&request),
+                                        request.eligibility(),
+                                        black_box(16),
+                                    ))
+                                    .unwrap(),
+                            )
+                        })
+                    },
+                );
+                stop.store(true, Ordering::Relaxed);
+            });
+        }
+    }
+    group.finish();
+}
+
+fn kv_event(worker: WorkerId, event_id: u64, data: KvCacheEventData) -> RouterEvent {
+    RouterEvent::new(
+        worker,
+        KvCacheEvent {
+            event_id,
+            dp_rank: 0,
+            data,
+        },
+    )
+}
+
+/// Event-thread cost of keeping resident block counts, untracked and tracked: 128 workers on 8
+/// event threads each store a 16-block chain one block per event, then remove it the same way.
+/// Single-block events are the cheapest to apply, so the per-event tracking cost is the largest
+/// share of them. Nothing reads the counters here; unpinned, a busy reader thread skews this
+/// bench. The pinned Mooncake replay's `--resident-blocks-reader-cpu` covers a reader, at that
+/// trace's larger events.
+fn resident_blocks_write(c: &mut Criterion) {
+    const WORKERS: u64 = 128;
+    const BLOCKS: u64 = 16;
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    let events: Vec<_> = (0..WORKERS)
+        .flat_map(|worker| {
+            let local: Vec<_> = (0..BLOCKS)
+                .map(|i| LocalBlockHash(worker * BLOCKS + i))
+                .collect();
+            let hashes: Vec<_> = compute_seq_hash_for_block(&local)
+                .into_iter()
+                .map(ExternalSequenceBlockHash)
+                .collect();
+            let stores = (0..hashes.len()).map(|i| {
+                let store = KvCacheStoreData {
+                    parent_hash: i.checked_sub(1).map(|parent| hashes[parent]),
+                    start_position: None,
+                    blocks: vec![KvCacheStoredBlockData {
+                        block_hash: hashes[i],
+                        tokens_hash: local[i],
+                        mm_extra_info: None,
+                    }],
+                };
+                kv_event(worker, i as u64, KvCacheEventData::Stored(store))
+            });
+            let removes = (0..hashes.len()).rev().map(|i| {
+                let remove = KvCacheRemoveData {
+                    block_hashes: vec![hashes[i]],
+                };
+                kv_event(worker, BLOCKS + i as u64, KvCacheEventData::Removed(remove))
+            });
+            stores.chain(removes).collect::<Vec<_>>()
+        })
+        .collect();
+    let mut group = c.benchmark_group("resident_blocks_write");
+    group.sample_size(20);
+    group.warm_up_time(Duration::from_secs(1));
+    group.measurement_time(Duration::from_secs(3));
+    group.throughput(Throughput::Elements(events.len() as u64));
+    for (name, tracked) in [("untracked", false), ("tracked", true)] {
+        let indexer = ThreadPoolIndexer::new(ConcurrentRadixTreeCompressed::new(), 8, 32);
+        if tracked {
+            indexer.bench_track_resident_blocks();
+        }
+        group.bench_function(name, |b| {
+            b.iter(|| {
+                runtime.block_on(async {
+                    for event in &events {
+                        indexer.apply_event(event.clone()).await;
+                    }
+                    indexer.flush().await;
+                })
+            })
+        });
+        indexer.shutdown();
+    }
+    group.finish();
+}
+
 fn default_policy_wrapper(c: &mut Criterion) {
     let config = KvRouterConfig {
         router_temperature: 0.0,
@@ -469,6 +664,8 @@ criterion_group!(
     worker_selection,
     custom_worker_selection,
     unused_preferred_taint_metadata,
+    resident_blocks_read,
+    resident_blocks_write,
     default_policy_wrapper,
     worker_projection
 );

lib/kv-router/src/indexer/concurrent_radix_tree_compressed/lane_lookup.rs

@@ -322,7 +322,6 @@ impl LaneLookup {
         });
     }
 
-    #[cfg(test)]
     pub(super) fn block_count(&self, worker: WorkerWithDpRank) -> Option<usize> {
         self.workers.get(&worker).map(BlockLookup::len)
     }