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,16 @@ 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; refreshed every 100 ms 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 tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is 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. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived or never arrive. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing 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 tracks for that worker rank, including cached blocks no running request uses, refreshed in the background. Without the declaration, before the first refresh, while the indexer is backlogged, or with a remote indexer, it returns `None`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`.
+
 ## Pick a Starting Point
 
 | Crate | Use it for |

lib/kv-router/benches/worker_selection.rs

@@ -125,6 +125,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 +407,60 @@ 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 = 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) in [("without", &without), ("with", &with)] {
+            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(),
+                        )
+                    })
+                },
+            );
+        }
+    }
+    group.finish();
+}
+
 fn default_policy_wrapper(c: &mut Criterion) {
     let config = KvRouterConfig {
         router_temperature: 0.0,
@@ -469,6 +559,7 @@ criterion_group!(
     worker_selection,
     custom_worker_selection,
     unused_preferred_taint_metadata,
+    resident_blocks_read,
     default_policy_wrapper,
     worker_projection
 );

lib/kv-router/src/indexer/kv_indexer.rs

@@ -15,7 +15,8 @@ use super::{
     ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig,
     ContainsWorkerBlockRequest, DumpRequest, EventKind, FlushRequest, GetWorkersRequest,
     KvIndexerInterface, KvIndexerMetrics, KvRouterError, MatchDetails, MatchDetailsRequest,
-    MatchRequest, PreBoundEventCounters, RadixTree, RoutingDecisionRequest, panic_payload_message,
+    MatchRequest, PreBoundEventCounters, RadixTree, ResidentBlockCounts,
+    ResidentBlockCountsRequest, RoutingDecisionRequest, panic_payload_message,
 };
 use crate::indexer::pruning::{BlockEntry, PruneConfig, WorkerPruneManager};
 use crate::protocols::*;
@@ -336,6 +337,9 @@ pub struct KvIndexer {
     remove_worker_dp_rank_tx: mpsc::Sender<(WorkerId, DpRank)>,
     /// A sender for get workers requests.
     get_workers_tx: mpsc::Sender<GetWorkersRequest>,
+    /// A sender for per-rank block count requests.
+    #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))]
+    resident_block_counts_tx: mpsc::Sender<ResidentBlockCountsRequest>,
     /// A sender for dump requests.
     dump_tx: mpsc::Sender<DumpRequest>,
     /// A sender for flush requests.
@@ -430,6 +434,8 @@ impl KvIndexer {
         let (remove_worker_dp_rank_tx, remove_worker_dp_rank_rx) =
             mpsc::channel::<(WorkerId, DpRank)>(16);
         let (get_workers_tx, get_workers_rx) = mpsc::channel::<GetWorkersRequest>(16);
+        let (resident_block_counts_tx, resident_block_counts_rx) =
+            mpsc::channel::<ResidentBlockCountsRequest>(16);
         let (dump_tx, dump_rx) = mpsc::channel::<DumpRequest>(16);
         let (flush_tx, flush_rx) = mpsc::channel::<FlushRequest>(16);
         let (routing_tx, mut routing_rx) = mpsc::channel::<RoutingDecisionRequest>(2048);
@@ -460,6 +466,7 @@ impl KvIndexer {
                     let mut remove_worker_rx = remove_worker_rx;
                     let mut remove_worker_dp_rank_rx = remove_worker_dp_rank_rx;
                     let mut get_workers_rx = get_workers_rx;
+                    let mut resident_block_counts_rx = resident_block_counts_rx;
                     let mut dump_rx = dump_rx;
                     let mut flush_rx = flush_rx;
                     let mut trie = delegate.map_or_else(RadixTree::new, RadixTree::new_with_delegate);
@@ -498,6 +505,11 @@ impl KvIndexer {
                                 }
                             }
 
+                            // Ahead of mutations and matches: under steady traffic, later arms may never run.
+                            Some(req) = resident_block_counts_rx.recv() => {
+                                let _ = req.resp.send(trie.resident_block_counts());
+                            }
+
                             Some(mutation) = mutation_rx.recv() => {
                                 apply_mutation(&mut trie, mutation, &counters, &prune_manager);
                             }
@@ -656,6 +668,7 @@ impl KvIndexer {
             remove_worker_tx,
             remove_worker_dp_rank_tx,
             get_workers_tx,
+            resident_block_counts_tx,
             dump_tx,
             flush_tx,
             routing_tx,
@@ -774,6 +787,21 @@ impl KvIndexer {
     pub fn get_workers_sender(&self) -> mpsc::Sender<GetWorkersRequest> {
         self.get_workers_tx.clone()
     }
+
+    /// Distinct blocks the tree holds for each worker rank.
+    #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))]
+    pub(crate) async fn resident_block_counts(&self) -> Result<ResidentBlockCounts, KvRouterError> {
+        let (resp_tx, resp_rx) = oneshot::channel();
+        self.resident_block_counts_tx
+            .send(ResidentBlockCountsRequest { resp: resp_tx })
+            .await
+            .map_err(|_| KvRouterError::IndexerOffline)?;
+        // A request sent while the actor drops its receiver stays queued with resp_tx alive.
+        tokio::select! {
+            resp = resp_rx => resp.map_err(|_| KvRouterError::IndexerDroppedRequest),
+            _ = self.resident_block_counts_tx.closed() => Err(KvRouterError::IndexerOffline),
+        }
+    }
 }
 
 #[async_trait]

lib/kv-router/src/indexer/mod.rs

@@ -69,6 +69,7 @@ mod lower_tier_indexers;
 mod metrics;
 #[cfg(feature = "bench")]
 mod observation;
+mod resident_blocks;
 mod thread_pool;
 mod traits;
 mod types;
@@ -95,6 +96,7 @@ pub use lower_tier_indexers::*;
 pub use metrics::*;
 #[cfg(feature = "bench")]
 pub use observation::*;
+pub(crate) use resident_blocks::*;
 pub use thread_pool::*;
 pub use traits::*;
 pub use types::*;

lib/kv-router/src/indexer/radix_tree.rs

@@ -13,7 +13,7 @@ use std::{
 use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
 
 use super::compressed_radix::{NodeState, append_dump_events};
-use super::{EventWarningKind, MatchDetails, PreBoundEventCounters};
+use super::{EventWarningKind, MatchDetails, PreBoundEventCounters, ResidentBlockCounts};
 use crate::protocols::*;
 
 pub(crate) type SharedRadixBlock = Rc<RefCell<RadixBlock>>;
@@ -860,6 +860,14 @@ impl RadixTree {
         self.lookup.values().map(FxHashMap::len).sum()
     }
 
+    /// Distinct blocks held by each worker rank.
+    pub(crate) fn resident_block_counts(&self) -> ResidentBlockCounts {
+        self.lookup
+            .iter()
+            .map(|(worker, worker_lookup)| (*worker, worker_lookup.len()))
+            .collect()
+    }
+
     #[cfg(test)]
     pub(crate) fn tree_size_for_worker(&self, worker: WorkerWithDpRank) -> Option<usize> {
         self.lookup.get(&worker).map(FxHashMap::len)

lib/kv-router/src/indexer/resident_blocks.rs

@@ -0,0 +1,68 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Per-rank resident block counts, published off the request path.
+//!
+//! Event threads own the per-rank block maps, so reading them is a queued round
+//! trip. A poller makes that trip and publishes the result here, where worker
+//! selection can read it without waiting on the indexer.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use arc_swap::{ArcSwapOption, Guard};
+use rustc_hash::FxHashMap;
+
+use crate::protocols::WorkerWithDpRank;
+
+/// How often routing hosts republish [`ResidentBlockCounts`].
+#[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))]
+pub(crate) const RESIDENT_BLOCK_COUNTS_POLL_INTERVAL: Duration = Duration::from_millis(100);
+
+/// Blocks the indexer tracked for each worker rank when polled.
+#[derive(Debug)]
+pub(crate) struct ResidentBlockCounts {
+    counts: FxHashMap<WorkerWithDpRank, u64>,
+}
+
+impl ResidentBlockCounts {
+    /// Blocks tracked for `worker`; zero when the indexer tracks none for it.
+    pub(crate) fn get(&self, worker: WorkerWithDpRank) -> u64 {
+        self.counts.get(&worker).copied().unwrap_or(0)
+    }
+}
+
+impl FromIterator<(WorkerWithDpRank, usize)> for ResidentBlockCounts {
+    fn from_iter<I: IntoIterator<Item = (WorkerWithDpRank, usize)>>(iter: I) -> Self {
+        let mut counts = FxHashMap::default();
+        for (worker, blocks) in iter {
+            *counts.entry(worker).or_default() += blocks as u64;
+        }
+        Self { counts }
+    }
+}
+
+/// The latest published [`ResidentBlockCounts`], shared by the poller and its readers.
+#[derive(Clone, Default)]
+pub(crate) struct ResidentBlockCountsHandle {
+    latest: Arc<ArcSwapOption<ResidentBlockCounts>>,
+}
+
+impl ResidentBlockCountsHandle {
+    /// The latest counts, or `None` before the first poll, while the indexer is backlogged, or
+    /// after the poller stops. Hold the guard only for one selection.
+    pub(crate) fn load(&self) -> Guard<Option<Arc<ResidentBlockCounts>>> {
+        self.latest.load()
+    }
+
+    #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))]
+    pub(crate) fn publish(&self, counts: ResidentBlockCounts) {
+        self.latest.store(Some(Arc::new(counts)));
+    }
+
+    /// Withdraw the counts so readers see `None` rather than stale counts.
+    #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))]
+    pub(crate) fn clear(&self) {
+        self.latest.store(None);
+    }
+}