feat(kv-router): expose prefix hashes, policy class, and worker capacity to selection policies

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

@@ -379,6 +379,9 @@ Routing cache resources start only if a component requests `CACHE`. `serve_index
 | `block_size()` | Tokens in one KV block |
 | `tracks_prefill_tokens()` | Whether the request contributes to prefill-load tracking |
 | `session_context()` | Optional session metadata described below |
+| `prefix_hashes()` | Optional chained hashes of the prompt's complete KV blocks; entry `i` identifies blocks `0..=i`, so equal entries mean an exactly shared prefix. `None` when the frontend does not track active blocks, as in disaggregated prefill pools; unique to the request when the host does not assume KV reuse, as in disaggregated decode pools. The standalone selection service passes a hash-only request's precomputed hashes through unchecked |
+| `policy_class()` | Optional policy class: the caller's value, replaced by a request classifier's override. A caller's value is not validated against the policy profile |
+| `worker_capacity(worker)` | Capacity advertised in a worker's runtime config: `total_kv_blocks()` per data-parallel rank, `None` when not advertised or advertised as zero. `None` for a worker the host does not know. Each call looks the worker up, so policies that never call it pay nothing |
 | `affinity_target()` | Optional advisory worker and data-parallel rank from soft session affinity |
 | `expected_output_tokens()` | Optional expected output length |
 | `priority_jump()` | Scheduler priority boost. Queue policies treat negative values as zero |

lib/kv-router/src/plugins/worker_selection.rs

@@ -11,7 +11,7 @@ pub use context::WorkerSelectionContext;
 pub(crate) use inputs::{CacheSnapshot, CandidateData, WorkerCacheData};
 pub use inputs::{
     ScoredWorkerCandidate, WorkerCacheInput, WorkerCacheInputs, WorkerCandidate, WorkerCandidates,
-    WorkerInputView, WorkerInputs, WorkerLoadInput,
+    WorkerCapacityInput, WorkerInputView, WorkerInputs, WorkerLoadInput,
 };
 
 pub use super::registry::{

lib/kv-router/src/plugins/worker_selection/context.rs

@@ -3,13 +3,16 @@
 
 //! Request metadata available to worker-selection components.
 
-use super::SessionContext;
-use crate::protocols::{WorkerAffinityTarget, WorkerWithDpRank};
+use super::{SessionContext, WorkerCapacityInput};
+use crate::protocols::{WorkerAffinityTarget, WorkerId, WorkerWithDpRank};
 use crate::scheduling::SchedulingRequest;
 
 /// Request-level values available to custom filters, scorers, and pickers.
 pub struct WorkerSelectionContext<'a> {
     pub(crate) request: &'a SchedulingRequest,
+    /// Looks up a worker's runtime config, so rows carry no capacity a policy may never read.
+    /// `Sync` keeps the context shareable across threads that score candidates in parallel.
+    pub(crate) worker_capacity: &'a (dyn Fn(WorkerId) -> Option<WorkerCapacityInput> + Sync),
     pub(crate) request_blocks: u64,
     pub(crate) block_size: u32,
     pub(crate) track_prefill_tokens: bool,
@@ -43,6 +46,20 @@ impl WorkerSelectionContext<'_> {
         self.block_size
     }
 
+    /// Chained hashes of the prompt's complete KV blocks, borrowed from this request.
+    ///
+    /// Entry `i` identifies prompt blocks `0..=i`, so two requests share their first `i + 1`
+    /// blocks exactly when their entry `i` is equal. LoRA adapters, cache namespaces, and
+    /// multimodal content hash into separate domains. These are the host's active-sequence
+    /// tracking hashes; do not compare them with engine KV-event hashes. None when the frontend
+    /// does not track active blocks, as in disaggregated prefill pools. When the host does not
+    /// assume KV reuse, as in disaggregated decode pools, every entry is unique to the request.
+    /// The standalone selection service always tracks active blocks, and passes the caller's
+    /// precomputed sequence hashes through unchecked for hash-only requests.
+    pub fn prefix_hashes(&self) -> Option<&[u64]> {
+        self.request.token_seq.as_deref()
+    }
+
     /// Return whether this request contributes to prefill-load tracking.
     pub fn tracks_prefill_tokens(&self) -> bool {
         self.track_prefill_tokens
@@ -66,6 +83,21 @@ impl WorkerSelectionContext<'_> {
         self.request.expected_output_tokens
     }
 
+    /// Return the capacity advertised in `worker`'s runtime config, or None if the host does not
+    /// know the worker. Every data-parallel rank of a worker shares one config. Each call looks
+    /// the worker up in the host's worker table.
+    pub fn worker_capacity(&self, worker: WorkerWithDpRank) -> Option<WorkerCapacityInput> {
+        (self.worker_capacity)(worker.worker_id)
+    }
+
+    /// Return the policy class requested for this request: the caller's value, replaced by a
+    /// request classifier's override. This is the requested name; a profile with class
+    /// families may still resolve it to a family member for queueing. A caller's value is not
+    /// validated against the profile, so treat it as untrusted input.
+    pub fn policy_class(&self) -> Option<&str> {
+        self.request.policy_class.as_deref()
+    }
+
     /// Return the request's scheduler priority boost.
     pub fn priority_jump(&self) -> f64 {
         self.request.priority_jump
@@ -81,3 +113,14 @@ impl WorkerSelectionContext<'_> {
         self.router_temperature_override
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::WorkerSelectionContext;
+
+    #[test]
+    fn context_is_send_and_sync() {
+        fn assert_send_sync<T: Send + Sync>() {}
+        assert_send_sync::<WorkerSelectionContext<'static>>();
+    }
+}

lib/kv-router/src/plugins/worker_selection/inputs.rs

@@ -3,7 +3,8 @@
 
 //! Worker input groups, stored snapshots, and component-scoped borrowed views.
 
-use crate::protocols::{SharedCacheHits, WorkerWithDpRank};
+use crate::protocols::{SharedCacheHits, WorkerConfigLike, WorkerWithDpRank};
+use std::num::NonZeroU64;
 use std::ops::BitOr;
 
 /// Host-owned row materialized once for the union of component input requirements.
@@ -232,6 +233,30 @@ pub struct WorkerLoadInput {
     pub(crate) active_requests: usize,
 }
 
+/// Capacity advertised in one worker's runtime config.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct WorkerCapacityInput {
+    total_kv_blocks: Option<NonZeroU64>,
+}
+
+impl WorkerCapacityInput {
+    /// Zero is the value some backends advertise when they cannot report capacity.
+    pub(crate) fn from_config(config: &impl WorkerConfigLike) -> Self {
+        Self {
+            total_kv_blocks: config.total_kv_blocks().and_then(NonZeroU64::new),
+        }
+    }
+
+    /// Return the KV-cache capacity advertised for each of the worker's data-parallel ranks, in
+    /// blocks of [`WorkerSelectionContext::block_size`](super::WorkerSelectionContext::block_size)
+    /// tokens, or None if it was not advertised. This is configured capacity, not free blocks.
+    /// Some backends derive it from aggregate or representative-rank data, so treat it as an
+    /// estimate.
+    pub fn total_kv_blocks(&self) -> Option<u64> {
+        self.total_kv_blocks.map(NonZeroU64::get)
+    }
+}
+
 /// Borrowed, index-aligned view of one custom picker's requested worker inputs.
 #[derive(Clone, Copy)]
 pub struct WorkerInputView<'a> {

lib/kv-router/src/scheduling/queue.rs

@@ -937,6 +937,8 @@ impl<
                         "unknown policy class {policy_class:?}"
                     ))
                 })?;
+            // Worker selection reads the classifier's class from the request.
+            request.policy_class = Some(policy_class);
         }
         if let Some(cost) = scheduling_cost_tokens {
             metadata.snapshot.scheduling_cost_tokens = cost;
@@ -2676,6 +2678,7 @@ policy_classes:
             queue.profile.class(metadata.class_index).name,
             "bulk_cached"
         );
+        assert_eq!(request.policy_class.as_deref(), Some("bulk"));
         assert_eq!(metadata.snapshot.scheduling_cost_tokens, 7);
         assert_eq!(metadata.due_at, Some(due_at));
 

lib/kv-router/src/scheduling/selector/mod.rs

@@ -19,7 +19,9 @@ pub use crate::plugins::worker_selection::{
 #[cfg(any(test, feature = "bench"))]
 use reference::{DefaultWorkerPicker, DefaultWorkerScorer};
 
-use crate::plugins::worker_selection::{CacheSnapshot, CandidateData, WorkerCacheData};
+use crate::plugins::worker_selection::{
+    CacheSnapshot, CandidateData, WorkerCacheData, WorkerCapacityInput,
+};
 pub use policy::WorkerSelectionPolicy;
 use policy::{ComposedPolicyState, WorkerSelectionPolicyStateRef, collect_policy_candidates};
 #[cfg(any(test, feature = "bench"))]
@@ -134,8 +136,17 @@ struct MaterializedSelectionInput<'a> {
     cache_snapshot: CacheSnapshot<'a>,
 }
 
+#[cfg(test)]
+fn no_worker_capacity(_: WorkerId) -> Option<WorkerCapacityInput> {
+    None
+}
+
 impl<'a> MaterializedSelectionInput<'a> {
-    fn new(request: &'a SchedulingRequest, block_size: u32) -> Self {
+    fn new(
+        request: &'a SchedulingRequest,
+        block_size: u32,
+        worker_capacity: &'a (dyn Fn(WorkerId) -> Option<WorkerCapacityInput> + Sync),
+    ) -> Self {
         Self {
             request,
             cache_snapshot: CacheSnapshot {
@@ -146,6 +157,7 @@ impl<'a> MaterializedSelectionInput<'a> {
             },
             context: WorkerSelectionContext {
                 request,
+                worker_capacity,
                 request_blocks: request.request_blocks(block_size),
                 block_size,
                 track_prefill_tokens: request.track_prefill_tokens,
@@ -335,7 +347,7 @@ fn log_selection<C: WorkerConfigLike>(
 #[inline(always)]
 // DefaultWorkerSelector and SelectionService both converge here. Only the scorer/picker stage is
 // dispatched; eligibility outcomes and result construction stay host-owned and shared.
-fn select_worker_with_policy<C: WorkerConfigLike>(
+fn select_worker_with_policy<C: WorkerConfigLike + Sync>(
     worker_type: &'static str,
     state: WorkerSelectionPolicyStateRef<'_>,
     workers: &HashMap<WorkerId, C>,
@@ -357,7 +369,12 @@ fn select_worker_with_policy<C: WorkerConfigLike>(
         }
     }
 
-    let mut input = MaterializedSelectionInput::new(request, block_size);
+    let worker_capacity = |worker_id: WorkerId| {
+        workers
+            .get(&worker_id)
+            .map(WorkerCapacityInput::from_config)
+    };
+    let mut input = MaterializedSelectionInput::new(request, block_size, &worker_capacity);
     input.context.pinned_worker = eligibility.pinned_worker();
     let selected = match state {
         #[cfg(any(test, feature = "bench"))]

lib/kv-router/src/scheduling/selector/policy.rs

@@ -327,7 +327,7 @@ pub(super) fn collect_policy_candidates<C: WorkerConfigLike>(
     Ok(has_eligible_worker)
 }
 
-impl<C: WorkerConfigLike> WorkerSelector<C> for WorkerSelectionPolicy {
+impl<C: WorkerConfigLike + Sync> WorkerSelector<C> for WorkerSelectionPolicy {
     fn uses_exclusive_affinity_target(&self) -> bool {
         #[cfg(any(test, feature = "bench"))]
         if matches!(&self.state, WorkerSelectionPolicyState::Reference(..)) {
@@ -375,12 +375,12 @@ impl<C: WorkerConfigLike> WorkerSelector<C> for WorkerSelectionPolicy {
 
 #[cfg(test)]
 mod tests {
-    use crate::plugins::worker_selection::WorkerInputView;
+    use crate::plugins::worker_selection::{WorkerCapacityInput, WorkerInputView};
     use crate::protocols::WorkerWithDpRank;
     use crate::scheduling::SessionContext;
     use std::{
-        cell::Cell,
         collections::{HashMap, HashSet},
+        sync::atomic::{AtomicUsize, Ordering},
     };
 
     use rustc_hash::FxHashMap;
@@ -585,7 +585,7 @@ mod tests {
     fn custom_policy_skips_undeclared_preferred_taints() {
         struct CountingTaintConfig {
             taints: HashSet<String>,
-            taint_reads: Cell<usize>,
+            taint_reads: AtomicUsize,
         }
 
         impl WorkerConfigLike for CountingTaintConfig {
@@ -606,7 +606,7 @@ mod tests {
             }
 
             fn taints(&self) -> &HashSet<String> {
-                self.taint_reads.set(self.taint_reads.get() + 1);
+                self.taint_reads.fetch_add(1, Ordering::Relaxed);
                 &self.taints
             }
         }
@@ -628,7 +628,7 @@ mod tests {
             0,
             CountingTaintConfig {
                 taints: HashSet::from(["preferred".to_string()]),
-                taint_reads: Cell::new(0),
+                taint_reads: AtomicUsize::new(0),
             },
         )]);
         let mut request = base_request(16);
@@ -651,7 +651,7 @@ mod tests {
             .unwrap();
         // Eligibility checks required taints once. The preference multiplier must not perform a
         // second lookup when no policy component declares it.
-        assert_eq!(workers[&0].taint_reads.get(), 1);
+        assert_eq!(workers[&0].taint_reads.load(Ordering::Relaxed), 1);
     }
 
     #[test]
@@ -946,4 +946,97 @@ mod tests {
         assert!(inputs.contains(WorkerInputs::CACHE));
         assert!(inputs.contains(WorkerInputs::LOAD));
     }
+
+    #[test]
+    fn request_facts_and_worker_capacity_reach_policy_components() {
+        struct CapacityConfig(Option<u64>, Option<u64>);
+
+        impl WorkerConfigLike for CapacityConfig {
+            fn data_parallel_start_rank(&self) -> u32 {
+                0
+            }
+            fn data_parallel_size(&self) -> u32 {
+                1
+            }
+            fn max_num_batched_tokens(&self) -> Option<u64> {
+                self.1
+            }
+            fn total_kv_blocks(&self) -> Option<u64> {
+                self.0
+            }
+        }
+
+        struct RequestFactsScorer(std::sync::Arc<AtomicUsize>);
+
+        impl WorkerScorer for RequestFactsScorer {
+            fn score(
+                &mut self,
+                context: &WorkerSelectionContext<'_>,
+                candidates: WorkerCandidates<'_>,
+                costs: &mut [f64],
+            ) -> Result<(), WorkerSelectionPolicyError> {
+                assert_eq!(context.prefix_hashes(), Some(&[11, 22][..]));
+                assert_eq!(context.policy_class(), Some("batch"));
+                costs.fill(0.0);
+                assert_eq!(candidates.len(), 2);
+                self.0.fetch_add(1, Ordering::Relaxed);
+                Ok(())
+            }
+        }
+
+        struct CapacityPicker;
+
+        impl WorkerPicker for CapacityPicker {
+            fn pick(
+                &mut self,
+                context: &WorkerSelectionContext<'_>,
+                input: WorkerInputView<'_>,
+            ) -> Result<usize, WorkerSelectionPolicyError> {
+                let capacity = |row: usize| {
+                    context
+                        .worker_capacity(input.candidates()[row].worker())
+                        .expect("known worker")
+                };
+                let advertised = (0..input.candidates().len())
+                    .find(|&row| capacity(row).total_kv_blocks().is_some())
+                    .expect("one worker advertises capacity");
+                assert_eq!(capacity(advertised).total_kv_blocks(), Some(100));
+                // The other worker advertises zero, which backends use for unknown capacity.
+                let unknown = 1 - advertised;
+                assert_eq!(capacity(unknown), WorkerCapacityInput::default());
+                assert!(
+                    context
+                        .worker_capacity(WorkerWithDpRank::from_worker_id(7))
+                        .is_none()
+                );
+                Ok(advertised)
+            }
+        }
+
+        let workers = HashMap::from([
+            (0, CapacityConfig(Some(100), Some(8192))),
+            (1, CapacityConfig(Some(0), None)),
+        ]);
+        let mut request = base_request(32);
+        request.token_seq = Some(vec![11, 22]);
+        request.policy_class = Some("batch".to_string());
+        let scored = std::sync::Arc::new(AtomicUsize::new(0));
+        let policy = WorkerSelectionPolicy::new(
+            KvRouterConfig::default(),
+            "test",
+            vec![Box::new(RequestFactsScorer(std::sync::Arc::clone(&scored)))],
+            Box::new(CapacityPicker),
+        );
+
+        let selected = policy
+            .select_worker(WorkerSelectionInput::configured(
+                &workers,
+                &request,
+                request.eligibility(),
+                16,
+            ))
+            .unwrap();
+        assert_eq!(selected.worker, WorkerWithDpRank::from_worker_id(0));
+        assert_eq!(scored.load(Ordering::Relaxed), 1);
+    }
 }