From bed2a1ae1e634670e047ce8d05c3df33ae436ded Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 14 Sep 2026 06:32:21 +0000 Subject: [PATCH 01/31] refactor(router): move default worker selection into builtin plugin Signed-off-by: Ishan Dhanani --- Cargo.lock | 8 + Cargo.toml | 1 + deploy/inference-gateway/ext-proc/Cargo.toml | 1 + .../inference-gateway/ext-proc/src/runner.rs | 6 +- .../ext-proc/src/selector.rs | 9 +- .../router/custom-worker-selection.mdx | 31 +- .../router/custom-policy-example/README.md | 6 + lib/bindings/python/Cargo.lock | 4 + lib/bindings/python/Cargo.toml | 3 +- lib/bindings/python/rust/lib.rs | 4 +- lib/kv-router/Cargo.toml | 3 +- lib/kv-router/benches/selection_core.rs | 6 + lib/kv-router/src/lib.rs | 4 +- lib/kv-router/src/plugins/registry.rs | 58 ++- lib/kv-router/src/plugins/worker_selection.rs | 60 +++ lib/kv-router/src/scheduling/AGENTS.md | 4 +- lib/kv-router/src/scheduling/local.rs | 18 +- lib/kv-router/src/scheduling/queue.rs | 16 +- lib/kv-router/src/scheduling/selector/mod.rs | 59 ++- .../src/scheduling/selector/policy.rs | 38 +- .../selector/{default.rs => reference.rs} | 1 + .../src/services/selection/core/mod.rs | 3 +- .../src/services/selection/core/tests.rs | 3 + .../src/services/selection/core/workers.rs | 19 +- .../src/services/selection/membership.rs | 6 + .../src/services/selection/service.rs | 12 + lib/llm/Cargo.toml | 1 + lib/llm/src/kv_router.rs | 8 +- lib/llm/src/kv_router/embedded.rs | 4 +- lib/llm/src/kv_router/plugins.rs | 8 +- lib/mocker/Cargo.toml | 1 + .../offline/extensions/kv_router/mod.rs | 8 +- lib/mocker/src/replay/router_shared.rs | 5 +- lib/router-plugins/builtin/Cargo.toml | 10 + .../builtin/benches/default_selection.rs | 121 +++++ lib/router-plugins/builtin/src/default.rs | 492 ++++++++++++++++++ lib/router-plugins/builtin/src/lib.rs | 27 +- .../builtin/tests/default_lifecycle.rs | 88 ++++ .../builtin/tests/default_policy.rs | 245 +++++++++ .../builtin/tests/support/mod.rs | 84 +++ 40 files changed, 1399 insertions(+), 86 deletions(-) rename lib/kv-router/src/scheduling/selector/{default.rs => reference.rs} (99%) create mode 100644 lib/router-plugins/builtin/benches/default_selection.rs create mode 100644 lib/router-plugins/builtin/src/default.rs create mode 100644 lib/router-plugins/builtin/tests/default_lifecycle.rs create mode 100644 lib/router-plugins/builtin/tests/default_policy.rs create mode 100644 lib/router-plugins/builtin/tests/support/mod.rs diff --git a/Cargo.lock b/Cargo.lock index de42eda60c5b..f58a30df1e70 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2464,9 +2464,14 @@ dependencies = [ name = "dynamo-custom-policy-builtin" version = "1.6.0" dependencies = [ + "criterion", "dynamo-kv-router", + "fastrand", + "parking_lot", + "rustc-hash 2.1.2", "serde", "tempfile", + "tokio", ] [[package]] @@ -2565,6 +2570,7 @@ dependencies = [ "axum 0.8.4", "bytes", "dashmap", + "dynamo-custom-policy-builtin", "dynamo-kv-router", "dynamo-llm", "dynamo-protocols", @@ -2691,6 +2697,7 @@ dependencies = [ "derive_builder", "dialoguer", "dynamo-bench", + "dynamo-custom-policy-builtin", "dynamo-kv-router", "dynamo-memory", "dynamo-mocker", @@ -2806,6 +2813,7 @@ dependencies = [ "bytes", "dashmap", "derive_builder", + "dynamo-custom-policy-builtin", "dynamo-data-gen", "dynamo-kv-router", "dynamo-tokens", diff --git a/Cargo.toml b/Cargo.toml index 04bd36aaa6d2..5a53f27cca84 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,7 @@ repository = "https://github.com/ai-dynamo/dynamo.git" keywords = ["llm", "genai", "inference", "nvidia", "distributed"] [workspace.dependencies] +dynamo-custom-policy-builtin = { path = "lib/router-plugins/builtin" } # Published separately by https://github.com/ai-dynamo/aisimulate. aisimulate-core = { version = "=0.12.0" } diff --git a/deploy/inference-gateway/ext-proc/Cargo.toml b/deploy/inference-gateway/ext-proc/Cargo.toml index de1d14186130..812c89c740bf 100644 --- a/deploy/inference-gateway/ext-proc/Cargo.toml +++ b/deploy/inference-gateway/ext-proc/Cargo.toml @@ -16,6 +16,7 @@ name = "dynamo-ext-proc" path = "src/bin/dynamo-ext-proc.rs" [dependencies] +dynamo-custom-policy-builtin = { workspace = true } async-trait = { workspace = true } dynamo-kv-router = { workspace = true, features = [ "metrics", diff --git a/deploy/inference-gateway/ext-proc/src/runner.rs b/deploy/inference-gateway/ext-proc/src/runner.rs index 102e6ef9bb90..2c8abc97af84 100644 --- a/deploy/inference-gateway/ext-proc/src/runner.rs +++ b/deploy/inference-gateway/ext-proc/src/runner.rs @@ -134,7 +134,11 @@ pub async fn run(policy_registry: Option) -> Resu { anyhow::bail!("linked worker-selection policies require DYN_EPP_MODE=standalone") } - run_inner(mode, policy_registry.unwrap_or_default()).await + run_inner( + mode, + policy_registry.unwrap_or_else(dynamo_custom_policy_builtin::default_registry), + ) + .await } fn init_tracing() { diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 2a96eec5e32b..790ea6d1510f 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -110,9 +110,12 @@ impl Selector { .await?; } - let mut builder = - SelectionServiceBuilder::new(kv_router_config, WorkerType::Aggregated, policy_registry) - .indexer_threads(cfg.selector_threads); + let mut builder = SelectionServiceBuilder::new( + kv_router_config, + WorkerType::Aggregated, + policy_registry.with_default_factory(dynamo_custom_policy_builtin::default_factory()), + ) + .indexer_threads(cfg.selector_threads); if let Some(peer_replication) = peer_replication { builder = builder.replica_sync(peer_replication.sync_port, Vec::new()); } diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 5827ada9569d..c2aec97b24e0 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -408,7 +408,10 @@ Request context and worker identity are always available. Dynamo calculates opti | Accessor | Meaning | |---|---| -| `request_blocks()` | Incoming prompt size in KV blocks | +| `prompt_tokens()` | Exact incoming prompt token count, borrowed from the request without allocation | +| `request_blocks()` | Incoming prompt size in KV blocks, rounded up | +| `has_tier_matches()` | Whether the current request lookup contains any tier-specific matches, before policy filters | +| `shared_cache_hits()` | Optional unweighted shared-cache ranges from the current lookup | | `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 | @@ -596,3 +599,29 @@ Holding a classification future pauses admission before dispatch; it does not in - If a signal adds work, storage, allocation, or another scan, run the worker-selection benchmark. The [example README](https://github.com/ai-dynamo/dynamo/blob/main/examples/router/custom-policy-example/README.md) contains the in-tree package names and build-check commands. For the built-in cost model, see [Routing Concepts](routing-concepts.md). For the standalone selection lifecycle, see [Standalone Selection Service](standalone-selection.md). + +## Default Policy Plugin + +The default worker selector lives in `dynamo-custom-policy-builtin`, beside the two-tier policy. Routing hosts install `default_registry()` even when they do not link a custom catalog. An omitted policy or the reserved `default` selection uses that factory. A Rust embedding that constructs `SelectionServiceBuilder` must supply a registry with a default factory or an explicit policy factory. + +To configure a separate default-policy instance, use the named provider: + +```yaml +worker_selection: + aggregated: tuned-default + instances: + - name: tuned-default + type: dynamo-default-cost-fn + parameters: + overlap_score_credit: 1.0 + prefill_load_scale: 1.0 + router_temperature: 0.0 +``` + +The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. The default plugin uses configured parameters and does not apply per-request score or temperature overrides. Plain disaggregated decode remains load-only when prefill tracking is off and conditional disaggregation is disabled. + +`WorkerCacheInput::accounting_cache_estimate()` returns the host's weighted block estimate and rounded token estimate. Missing estimates are zero. These estimates are not clamped to prompt length and do not measure physical GPU residency. The plugin can use them without changing the load that the host records. `WorkerLoadInput::is_available()` distinguishes a missing load observation from an observed idle worker. Both values describe the current selection snapshot. + +A picker can implement `pick_with_cost()` to return a row and an optional finite cost for host diagnostics. Existing `pick()` implementations remain valid. The host checks the row and cost before booking. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. + +Direct Rust callers and replay import `DefaultWorkerSelector` from `dynamo_custom_policy_builtin`. Scheduler actors use `default_policy()` directly. The core's old selector is available only as a test and benchmark reference, not in production builds. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 15087c3bf518..ac03bfb36969 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -372,3 +372,9 @@ python3 -m dynamo.mocker \ ``` Send the same `curl` request from a fourth terminal. The frontend log records separate prefill and decode selections. Each worker set runs its own filter, scorer, and picker. + +## Exact Prompt Length + +Filters, scorers, and pickers can read `context.prompt_tokens()` for the exact input token count. This returns the existing request value without a prompt copy or allocation. Use `context.request_blocks()` when the policy needs rounded KV blocks instead. For example, at 16 tokens per block, a 17-token prompt has 2 blocks. + +The builtin default now uses the public policy API too. Rust hosts install `dynamo_custom_policy_builtin::default_registry()` before adding custom providers. This supplies default selection for roles that the custom configuration does not replace. diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 54ee7cebaf90..5490b57f3a13 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2106,6 +2106,8 @@ name = "dynamo-custom-policy-builtin" version = "1.6.0" dependencies = [ "dynamo-kv-router", + "fastrand", + "parking_lot", "serde", ] @@ -2205,6 +2207,7 @@ dependencies = [ "derive-getters", "derive_builder", "dialoguer", + "dynamo-custom-policy-builtin", "dynamo-kv-router", "dynamo-memory", "dynamo-mocker", @@ -2308,6 +2311,7 @@ dependencies = [ "bytes", "dashmap", "derive_builder", + "dynamo-custom-policy-builtin", "dynamo-data-gen", "dynamo-kv-router", "dynamo-tokens", diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index 8ead4c72ec59..45eeb1ab6747 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -28,7 +28,6 @@ default = ["custom-policy"] # router-policy YAML without rebuilding. A private image replaces only the catalog dependency, so # its policies are added alongside the shipped ones rather than displacing them. custom-policy = [ - "dep:dynamo-custom-policy-builtin", "dep:dynamo-worker-selection-policy-catalog", "dynamo-kv-router/standalone-selection", ] @@ -98,7 +97,7 @@ uuid = { version = "1.18.1" } # kv-indexer / slot-tracker / shared kv-router types dynamo-kv-router = { path = "../../kv-router" } -dynamo-custom-policy-builtin = { path = "../../router-plugins/builtin", optional = true } +dynamo-custom-policy-builtin = { path = "../../router-plugins/builtin" } dynamo-worker-selection-policy-catalog = { path = "../../router-plugins/catalog", optional = true } clap = { version = "4.5", features = ["derive"], optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 5b10f891ebbb..fa4cf09a3b0c 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -439,13 +439,13 @@ pub(crate) fn linked_worker_selection_policy_registry() -> RouterPluginRegistry #[cfg(not(feature = "custom-policy"))] { - RouterPluginRegistry::default() + dynamo_custom_policy_builtin::default_registry() } } #[cfg(feature = "custom-policy")] fn register_core_with_router_plugins(m: &Bound<'_, PyModule>) -> PyResult<()> { - let mut registry = RouterPluginRegistry::default(); + let mut registry = dynamo_custom_policy_builtin::default_registry(); // The policies Dynamo ships register first, so a replaced catalog that reuses one of their // type names fails here instead of silently overriding it. dynamo_custom_policy_builtin::register(&mut registry) diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index 2045158ecf18..420b11e87d31 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -103,8 +103,9 @@ harness = false [[bench]] name = "worker_selection" harness = false +required-features = ["bench"] [[bench]] name = "selection_core" harness = false -required-features = ["standalone-selection"] +required-features = ["standalone-selection", "bench"] diff --git a/lib/kv-router/benches/selection_core.rs b/lib/kv-router/benches/selection_core.rs index b472c21b1a99..ade4d3f6c197 100644 --- a/lib/kv-router/benches/selection_core.rs +++ b/lib/kv-router/benches/selection_core.rs @@ -45,6 +45,12 @@ fn core_with_workers( 1, CancellationToken::new(), SelectionCacheConfig::default(), + std::sync::Arc::new(|config, role, _| { + dynamo_kv_router::WorkerSelectionPolicy::default( + config.clone(), + role.default_selector_label(), + ) + }), ) .expect("core"); runtime.block_on(async { diff --git a/lib/kv-router/src/lib.rs b/lib/kv-router/src/lib.rs index 2ce2cf9a14f9..bf2df89bdac0 100644 --- a/lib/kv-router/src/lib.rs +++ b/lib/kv-router/src/lib.rs @@ -76,7 +76,9 @@ pub use scheduling::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, SessionContext, WorkerSelectionInputTrigger, WorkerSelectionPolicyError, }; -pub use selector::{DefaultWorkerSelector, WorkerSelectionInput, WorkerSelector}; +#[cfg(any(test, feature = "bench"))] +pub use selector::DefaultWorkerSelector; +pub use selector::{WorkerSelectionInput, WorkerSelector}; // TODO(v1.7): Remove these compatibility re-exports; use crate::plugins instead. pub use plugins::worker_selection::{ ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputView, diff --git a/lib/kv-router/src/plugins/registry.rs b/lib/kv-router/src/plugins/registry.rs index 5a48e7055bd8..137a5c2e943a 100644 --- a/lib/kv-router/src/plugins/registry.rs +++ b/lib/kv-router/src/plugins/registry.rs @@ -13,7 +13,9 @@ use super::request_classifier::RequestClassifierRegistry; use super::request_classifier::{ RequestClassifierFactory, RequestClassifierProvider, RequestClassifierRegistryError, }; -use super::worker_selection::{WorkerSelectionPolicy, WorkerSelectionPolicyFactory}; +#[cfg(test)] +use super::worker_selection::WorkerSelectionPolicy; +use super::worker_selection::WorkerSelectionPolicyFactory; use crate::WorkerType; use crate::config::KvRouterConfig; use crate::scheduling::config::WorkerSelectionPolicySelections; @@ -69,6 +71,7 @@ impl WorkerSelectionPolicyProviderError { pub struct RouterPluginRegistry { providers: HashMap, request_classifiers: RequestClassifierRegistry, + default_factory: Option, } /// Existing catalogs can retain their worker-selection registration signature. @@ -80,6 +83,8 @@ pub type WorkerSelectionPolicyRegistry = RouterPluginRegistry; pub enum WorkerSelectionPolicyRegistryError { #[error("worker-selection policy type must not be empty")] EmptyName, + #[error("routing host must supply a default worker-selection policy factory")] + MissingDefault, #[error("worker-selection policy type 'default' is reserved for Dynamo's built-in selector")] ReservedDefault, #[error("worker-selection policy type {name:?} is already registered")] @@ -99,6 +104,18 @@ pub enum WorkerSelectionPolicyRegistryError { } impl RouterPluginRegistry { + /// Construct a registry with the host's required default policy. Custom registrations + /// cannot replace this factory. An empty registry still supports explicit custom policies. + pub fn new(default_factory: WorkerSelectionPolicyFactory) -> Self { + Self { default_factory: Some(default_factory), ..Self::default() } + } + + /// Supply the host default while preserving registered custom plugin types. + pub fn with_default_factory(mut self, factory: WorkerSelectionPolicyFactory) -> Self { + self.default_factory = Some(factory); + self + } + /// Resolve every configured plugin once before constructing any routers. pub fn resolve_plugins( &self, @@ -215,9 +232,22 @@ impl RouterPluginRegistry { self.resolve_selected_cached(policy_config, selected.encode.as_deref(), &mut resolved)?; if aggregated.is_none() && prefill.is_none() && decode.is_none() && encode.is_none() { - return Ok(None); + return Ok(self.default_factory.clone()); } + #[cfg(not(test))] + if self.default_factory.is_none() + && [ + aggregated.is_none(), + prefill.is_none(), + decode.is_none(), + encode.is_none(), + ] + .contains(&true) + { + return Err(WorkerSelectionPolicyRegistryError::MissingDefault); + } + let default_factory = self.default_factory.clone(); Ok(Some(Arc::new(move |config, worker_type, partition| { let selected = match worker_type { WorkerType::Aggregated => aggregated.as_ref(), @@ -227,10 +257,22 @@ impl RouterPluginRegistry { }; match selected { Some(factory) => factory(config, worker_type, partition), - None => WorkerSelectionPolicy::default( - config.clone(), - worker_type.default_selector_label(), - ), + None => { + #[cfg(test)] + if default_factory.is_none() { + return WorkerSelectionPolicy::default( + config.clone(), + worker_type.default_selector_label(), + ); + } + default_factory + .as_ref() + .expect("default factory validated at resolution")( + config, + worker_type, + partition, + ) + } } }))) } @@ -262,10 +304,10 @@ impl RouterPluginRegistry { selected: Option<&str>, ) -> Result, WorkerSelectionPolicyRegistryError> { let Some(selected) = selected else { - return Ok(None); + return Ok(self.default_factory.clone()); }; if selected == "default" { - return Ok(None); + return Ok(self.default_factory.clone()); } let instance = config .and_then(|config| config.instance(selected)) diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index b16c5db7cb1e..4b83fc85d364 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use crate::protocols::{WorkerAffinityTarget, WorkerWithDpRank}; use crate::scheduling::SchedulingRequest; +#[cfg(any(test, feature = "bench"))] use crate::scheduling::selector::LogitWeights; use crate::{KvRouterConfig, RoutingPartitionRef, WorkerType}; @@ -34,11 +35,15 @@ pub type WorkerSelectionPolicyFactory = Arc< /// Request-level values available to custom filters, scorers, and pickers. pub struct WorkerSelectionContext<'a> { pub(crate) request: &'a SchedulingRequest, + #[cfg(any(test, feature = "bench"))] pub(crate) request_id: &'a str, pub(crate) request_blocks: u64, pub(crate) block_size: u32, pub(crate) track_prefill_tokens: bool, + #[cfg(any(test, feature = "bench"))] pub(crate) weights: LogitWeights, + pub(crate) has_tier_matches: bool, + pub(crate) pinned_worker: Option, pub(crate) router_temperature_override: Option, } @@ -74,6 +79,7 @@ impl WorkerInputs { pub const PREFERRED_TAINT: Self = Self(1 << 2); /// Request host-owned active-request counts. pub const OCCUPANCY: Self = Self(1 << 5); + #[cfg(any(test, feature = "bench"))] pub(crate) const ALL: Self = Self(Self::CACHE.0 | Self::LOAD.0 | Self::PREFERRED_TAINT.0); pub const fn contains(self, other: Self) -> bool { @@ -97,6 +103,7 @@ impl BitOr for WorkerInputs { #[derive(Clone, Copy, Default)] pub struct WorkerCacheInput { pub(crate) effective_overlap_blocks: f64, + pub(crate) estimated_cached_tokens: usize, pub(crate) device_overlap_blocks: f64, pub(crate) host_overlap_blocks: f64, pub(crate) disk_overlap_blocks: f64, @@ -106,7 +113,9 @@ pub struct WorkerCacheInput { /// Active-load values for one worker. #[derive(Clone, Copy, Default)] pub struct WorkerLoadInput { + #[cfg(any(test, feature = "bench"))] pub(crate) raw_prefill_blocks: f64, + pub(crate) available: bool, pub(crate) active_prefill_tokens: usize, pub(crate) decode_cost_blocks: f64, pub(crate) active_requests: usize, @@ -154,6 +163,16 @@ pub trait WorkerFilter: Send { /// Selects one row after all filters and scorers run. pub trait WorkerPicker: Send { + /// Return a selected row and an optional policy-computed cost for host diagnostics. + /// The host validates both. Existing pickers use the summed scorer cost by default. + fn pick_with_cost( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result<(usize, Option), WorkerSelectionPolicyError> { + self.pick(context, input).map(|row| (row, None)) + } + /// Declare the optional worker-signal columns needed by this picker. fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::NONE @@ -169,6 +188,33 @@ pub trait WorkerPicker: Send { } impl WorkerSelectionContext<'_> { + /// The exact worker/rank imposed by the host for this selection, if any. + /// Includes explicit pins and eligible exclusive-affinity targets. This is + /// read-only routing metadata, not permission to change eligibility. + pub fn pinned_worker(&self) -> Option { + self.pinned_worker + } + + /// Exact incoming prompt length in tokens. Borrowed from this request; no rounding, + /// cache weighting, or additional storage is involved. + pub fn prompt_tokens(&self) -> usize { + self.request.isl_tokens + } + + /// Shared-cache ranges from this request's lookup snapshot, if present. + /// The ranges are unweighted block offsets. The host owns their lifetime and + /// policy inspection does not perform a lookup or change accounting. + pub fn shared_cache_hits(&self) -> Option<&crate::SharedCacheHits> { + self.request.shared_cache_hits.as_ref() + } + + /// Whether this request has any tier-specific cache matches, before worker filtering. + /// False means the host only supplied its accounting estimate (or no cache data). + /// The value describes the current lookup snapshot, not worker cache capacity. + pub fn has_tier_matches(&self) -> bool { + self.has_tier_matches + } + /// Return the incoming prompt size in KV blocks. pub fn request_blocks(&self) -> u64 { self.request_blocks @@ -297,6 +343,14 @@ impl ScoredWorkerCandidate { } impl WorkerCacheInput { + /// Host accounting estimate for this worker, in weighted KV blocks and rounded tokens. + /// Lower-tier matches use the host's cache weights. Missing estimates are zero; + /// neither value is clamped to prompt length. This is the current lookup snapshot, + /// not a count of physically resident GPU tokens. Policy scores do not alter it. + pub fn accounting_cache_estimate(&self) -> (f64, usize) { + (self.effective_overlap_blocks, self.estimated_cached_tokens) + } + /// Return device-resident prefix overlap in KV blocks. pub fn device_overlap_blocks(&self) -> f64 { self.device_overlap_blocks @@ -319,6 +373,12 @@ impl WorkerCacheInput { } impl WorkerLoadInput { + /// Whether the host supplied a load projection for this worker in this selection. + /// False distinguishes a missing observation from an observed idle worker. + pub fn is_available(&self) -> bool { + self.available + } + /// Return the tokens active in this worker's prefill stage. pub fn active_prefill_tokens(&self) -> usize { self.active_prefill_tokens diff --git a/lib/kv-router/src/scheduling/AGENTS.md b/lib/kv-router/src/scheduling/AGENTS.md index 8f80b3312e41..9186a7c38299 100644 --- a/lib/kv-router/src/scheduling/AGENTS.md +++ b/lib/kv-router/src/scheduling/AGENTS.md @@ -135,12 +135,12 @@ The `selector` module contains the public Rust contract for custom worker filter - Require callers to name each `WorkerInputs` group that they use. Do not add a public `ALL` shortcut. - Before you add a value to an existing input group, account for its calculation and retained-column cost for every policy that requests that group. - Do not pass the full `SchedulingRequest`, worker maps, router configuration internals, default-score weights, or host-owned eligibility and reservation state to custom policies. -- Keep `DefaultWorkerScorer` and `DefaultWorkerPicker` internal. External policies own their filters and both scoring and picking stages. +- The builtin default policy lives in `lib/router-plugins/builtin` and uses only the public policy API. Keep the reference selector test/benchmark-only. Plugins own their score calculation and worker choice. - Keep eligibility, picker-row validation, accounting, and reservation in the host path. Before each public API addition: 1. Search all accessors, re-exports, documentation, examples, and external-looking call sites. 2. Add one focused contract test that uses the new value through `WorkerFilter`, `WorkerScorer`, or `WorkerPicker`. -3. Update `docs/fern/pages/developer-guide/advanced-customizations/custom-worker-selection.mdx` and one canonical example. +3. Update `docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx` and one canonical example. 4. If the signal adds work, storage, allocation, or another scan to the selection path, run the worker-selection benchmark. diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index b850f9879374..c8103f377cc1 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -21,7 +21,7 @@ use super::queue::{ SchedulerQueue, }; use super::request_classifier::{RequestClassifierRuntime, RequestLifecycle}; -use super::selector::{DefaultWorkerSelector, WorkerSelector}; +use super::selector::WorkerSelector; use super::types::{ AdmissionAttempt, AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, PotentialLoad, @@ -45,8 +45,12 @@ enum WorkerConfigReconcileOutcome { Rejected, } -pub struct LocalScheduler -where +pub struct LocalScheduler< + P, + C, + Sel = super::selector::WorkerSelectionPolicy, + RF = NoopOverlapScoresRefresh, +> where P: SequencePublisher, C: WorkerConfigLike, Sel: WorkerSelector, @@ -906,7 +910,7 @@ mod tests { monitor_worker_configs: bool, prefill_load_estimator: Option>, ) -> ( - Arc>, + Arc>, Arc>, watch::Sender>, CancellationToken, @@ -928,7 +932,7 @@ mod tests { prefill_load_estimator: Option>, replica_sync: bool, ) -> ( - Arc>, + Arc>, Arc>, watch::Sender>, CancellationToken, @@ -976,7 +980,9 @@ mod tests { } async fn wait_for_pending_count( - scheduler: &Arc>, + scheduler: &Arc< + LocalScheduler, + >, expected: usize, ) { tokio::time::timeout(Duration::from_millis(250), async { diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index d120f186198c..4b017f05a6d9 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -24,7 +24,7 @@ use super::policy_queue::{PolicyQueue, QueueMetadata, QueueSnapshot}; use super::prefill_load::{PrefillLoadEstimator, effective_prefill_tokens}; use super::queue_admission::WorkerPlacement; use super::request_classifier::{ClassificationOverrides, ClassifyRequest}; -use super::selector::{DefaultWorkerSelector, WorkerSelectionInput, WorkerSelector}; +use super::selector::{WorkerSelectionInput, WorkerSelector}; use super::types::{ AdvisorySchedulingResponse, AdvisoryWorkerLoad, AttemptId, KvSchedulerError, NonMaxOverlapSelection, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, @@ -508,7 +508,7 @@ struct SchedulerQueueActor< pub struct SchedulerQueue< P: SequencePublisher, C: WorkerConfigLike, - Sel: WorkerSelector = DefaultWorkerSelector, + Sel: WorkerSelector = super::selector::WorkerSelectionPolicy, RF: OverlapScoresRefresh = NoopOverlapScoresRefresh, > { admission_tx: mpsc::Sender, @@ -2174,13 +2174,14 @@ mod tests { } } + #[allow(clippy::type_complexity)] fn make_queue( num_workers: usize, block_size: u32, isl: usize, threshold_frac: Option, ) -> ( - Arc>, + Arc>, Arc>, ) { let (queue, slots, _tx) = @@ -2245,7 +2246,7 @@ mod tests { threshold_frac: Option, prefill_load_estimator: Option>, ) -> ( - Arc>, + Arc>, Arc>, watch::Sender>, ) { @@ -2301,7 +2302,7 @@ mod tests { max_num_batched_tokens: usize, profile: PolicyProfile, ) -> ( - Arc>, + Arc>, Arc>, ) { let (queue, slots, _cfg_tx) = make_queue_with_profile_and_sender( @@ -2320,7 +2321,7 @@ mod tests { max_num_batched_tokens: usize, profile: PolicyProfile, ) -> ( - Arc>, + Arc>, Arc>, watch::Sender>, ) { @@ -2360,6 +2361,7 @@ mod tests { (queue, slots, cfg_tx) } + #[allow(clippy::type_complexity)] fn make_queue_with_providers( num_workers: usize, block_size: u32, @@ -2367,7 +2369,7 @@ mod tests { overloaded_worker_provider: Option, available_worker_provider: Option, ) -> ( - Arc>, + Arc>, Arc>, ) { let dp_range: HashMap = diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 61b4f5c1c33b..8938febe80cf 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -3,24 +3,27 @@ use std::collections::HashMap; -mod default; mod policy; +#[cfg(any(test, feature = "bench"))] +mod reference; -pub use default::DefaultWorkerSelector; +#[cfg(any(test, feature = "bench"))] +pub use reference::DefaultWorkerSelector; -use default::{DefaultWorkerPicker, DefaultWorkerScorer}; // TODO(v1.7): Remove these compatibility re-exports; use crate::plugins instead. pub use crate::plugins::worker_selection::{ ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, }; +#[cfg(any(test, feature = "bench"))] +use reference::{DefaultWorkerPicker, DefaultWorkerScorer}; pub use policy::WorkerSelectionPolicy; - -use default::{pick_default_worker, selection_weights}; use policy::{ CustomWorkerSelectionState, WorkerSelectionPolicyStateRef, collect_custom_candidates, }; +#[cfg(any(test, feature = "bench"))] +use reference::{pick_default_worker, selection_weights}; use super::config::KvRouterConfig; use super::filter::{RoutingEligibility, WorkerEligibilityError}; @@ -127,6 +130,7 @@ impl<'a, C: WorkerConfigLike> WorkerSelectionInput<'a, C> { } #[derive(Debug, Clone, Copy)] +#[cfg(any(test, feature = "bench"))] pub(crate) struct LogitWeights { overlap_score_credit: f64, overlap_score_credit_decay: f64, @@ -140,16 +144,26 @@ struct MaterializedSelectionInput<'a> { } impl<'a> MaterializedSelectionInput<'a> { - fn new(request: &'a SchedulingRequest, block_size: u32, weights: LogitWeights) -> Self { + fn new( + request: &'a SchedulingRequest, + block_size: u32, + #[cfg(any(test, feature = "bench"))] weights: LogitWeights, + ) -> Self { Self { request, context: WorkerSelectionContext { request, + #[cfg(any(test, feature = "bench"))] request_id: request.mode.request_id().unwrap_or("-"), request_blocks: request.request_blocks(block_size), block_size, track_prefill_tokens: request.track_prefill_tokens, + #[cfg(any(test, feature = "bench"))] weights, + pinned_worker: request.pinned_worker, + has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() + || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() + || !request.overlap.tier_overlap_blocks.disk.is_empty(), router_temperature_override: request .router_config_override .as_ref() @@ -168,6 +182,8 @@ impl<'a> MaterializedSelectionInput<'a> { worker, preferred_taint_multiplier, inputs, + #[cfg(any(test, feature = "bench"))] + false, |_, device_overlap_blocks| device_overlap_blocks, ) } @@ -177,6 +193,7 @@ impl<'a> MaterializedSelectionInput<'a> { worker: WorkerWithDpRank, preferred_taint_multiplier: Option, inputs: WorkerInputs, + #[cfg(any(test, feature = "bench"))] reference_formula: bool, select_device_overlap: impl FnOnce(f64, f64) -> f64, ) -> WorkerCandidate { let cached_tokens = if inputs.contains(WorkerInputs::CACHE) @@ -212,6 +229,7 @@ impl<'a> MaterializedSelectionInput<'a> { }; WorkerCacheInput { effective_overlap_blocks, + estimated_cached_tokens: cached_tokens, device_overlap_blocks, host_overlap_blocks: self .request @@ -235,7 +253,8 @@ impl<'a> MaterializedSelectionInput<'a> { WorkerCacheInput::default() }; let load = if inputs.contains(WorkerInputs::LOAD) { - let raw_prefill_tokens = if self.request.track_prefill_tokens { + #[cfg(any(test, feature = "bench"))] + let raw_prefill_tokens = if reference_formula && self.request.track_prefill_tokens { match worker_load { Some(load) => { // Preserve the legacy operation order when overlap exceeds the prompt. @@ -251,8 +270,11 @@ impl<'a> MaterializedSelectionInput<'a> { } else { 0 } as f64; + let available = worker_load.is_some(); let worker_load = worker_load.unwrap_or_default(); WorkerLoadInput { + available, + #[cfg(any(test, feature = "bench"))] raw_prefill_blocks: raw_prefill_tokens / self.context.block_size as f64, active_prefill_tokens: worker_load.active_prefill_tokens, decode_cost_blocks: worker_load.potential_decode_blocks() as f64, @@ -380,9 +402,19 @@ fn select_worker_with_policy( } } + #[cfg(any(test, feature = "bench"))] let weights = selection_weights(kv_router_config, request); - let input = MaterializedSelectionInput::new(request, block_size, weights); + #[cfg(not(any(test, feature = "bench")))] + let _ = kv_router_config; + let mut input = MaterializedSelectionInput::new( + request, + block_size, + #[cfg(any(test, feature = "bench"))] + weights, + ); + input.context.pinned_worker = eligibility.pinned_worker(); let selected = match state { + #[cfg(any(test, feature = "bench"))] WorkerSelectionPolicyStateRef::Default(picker) => { let scorer = DefaultWorkerScorer { kv_router_config, @@ -425,7 +457,7 @@ fn select_worker_with_policy( .contains(WorkerInputs::LOAD) .then_some(load_inputs.as_slice()), }; - let row = picker.pick(&input.context, picker_input)?; + let (row, selected_cost) = picker.pick_with_cost(&input.context, picker_input)?; let Some(candidate) = candidates.get(row) else { return Err(WorkerSelectionPolicyError::InvalidPickerRow { row, @@ -433,7 +465,14 @@ fn select_worker_with_policy( } .into()); }; - Some((candidate.worker, candidate.cost)) + let cost = selected_cost.unwrap_or(candidate.cost); + if !cost.is_finite() { + return Err(WorkerSelectionPolicyError::failed( + "picker returned non-finite cost", + ) + .into()); + } + Some((candidate.worker, cost)) } } }; diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 01213e781f67..e5c974b973db 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -4,14 +4,20 @@ use std::cell::RefCell; use std::collections::HashMap; +#[cfg(any(test, feature = "bench"))] +use super::DefaultWorkerPicker; use super::{ - DefaultWorkerPicker, MaterializedSelectionInput, WorkerSelectionInput, WorkerSelector, - select_worker_with_policy, + MaterializedSelectionInput, WorkerSelectionInput, WorkerSelector, select_worker_with_policy, +}; + +use crate::protocols::{ + WorkerConfigLike, WorkerId, WorkerSelectionResult, }; -use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult}; use crate::scheduling::config::KvRouterConfig; use crate::scheduling::filter::RoutingEligibility; -use crate::scheduling::types::{KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError}; +use crate::scheduling::types::{ + KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError, +}; use crate::plugins::worker_selection::{ ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputs, @@ -20,12 +26,14 @@ use crate::plugins::worker_selection::{ #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] pub(super) enum WorkerSelectionPolicyState { + #[cfg(any(test, feature = "bench"))] Default(DefaultWorkerPicker), /// Policy-local state owned and called serially by one scheduler queue actor. Custom(RefCell), } pub(super) enum WorkerSelectionPolicyStateRef<'a> { + #[cfg(any(test, feature = "bench"))] Default(&'a DefaultWorkerPicker), Custom(&'a RefCell), } @@ -50,6 +58,7 @@ pub struct WorkerSelectionPolicy { kv_router_config: KvRouterConfig, worker_label: &'static str, state: WorkerSelectionPolicyState, + exclusive_affinity: bool, } impl WorkerSelectionPolicy { @@ -87,6 +96,7 @@ impl WorkerSelectionPolicy { Self { kv_router_config, worker_label, + exclusive_affinity: false, state: WorkerSelectionPolicyState::Custom(RefCell::new(CustomWorkerSelectionState { filters, scorers, @@ -101,15 +111,24 @@ impl WorkerSelectionPolicy { } } + /// Ask the host to constrain selection to an eligible affinity target. + /// Explicit request pins remain mandatory regardless of this option. + pub fn with_exclusive_affinity(mut self, exclusive: bool) -> Self { + self.exclusive_affinity = exclusive; + self + } + /// Wrap Dynamo's built-in selector for a host that uses the policy selector type. /// /// `worker_label` selects the built-in scoring and logging contract. Typed hosts use /// [`crate::WorkerType::default_selector_label`] to preserve Dynamo's historical behavior. + #[cfg(any(test, feature = "bench"))] pub fn default(kv_router_config: KvRouterConfig, worker_label: &'static str) -> Self { let picker = DefaultWorkerPicker::new(); Self { kv_router_config, worker_label, + exclusive_affinity: false, state: WorkerSelectionPolicyState::Default(picker), } } @@ -269,11 +288,16 @@ pub(super) fn collect_custom_candidates( impl WorkerSelector for WorkerSelectionPolicy { fn uses_exclusive_affinity_target(&self) -> bool { - matches!(&self.state, WorkerSelectionPolicyState::Default(_)) + #[cfg(any(test, feature = "bench"))] + if matches!(&self.state, WorkerSelectionPolicyState::Default(_)) { + return true; + } + self.exclusive_affinity } fn required_worker_inputs(&self) -> WorkerInputs { match &self.state { + #[cfg(any(test, feature = "bench"))] WorkerSelectionPolicyState::Default(_) => WorkerInputs::CACHE | WorkerInputs::LOAD, WorkerSelectionPolicyState::Custom(state) => { let state = state.borrow(); @@ -289,6 +313,7 @@ impl WorkerSelector for WorkerSelectionPolicy { ) -> Result { let (workers, request, eligibility, block_size) = input.into_configured()?; let state = match &self.state { + #[cfg(any(test, feature = "bench"))] WorkerSelectionPolicyState::Default(picker) => { WorkerSelectionPolicyStateRef::Default(picker) } @@ -320,9 +345,6 @@ mod tests { use super::super::DefaultWorkerSelector; use super::super::test_support::*; use super::*; - use crate::SessionContext; - use crate::plugins::worker_selection::WorkerInputView; - use crate::protocols::WorkerWithDpRank; use crate::scheduling::WorkerSelectionInputTrigger; fn uses_exclusive_affinity(selector: &impl WorkerSelector) -> bool { diff --git a/lib/kv-router/src/scheduling/selector/default.rs b/lib/kv-router/src/scheduling/selector/reference.rs similarity index 99% rename from lib/kv-router/src/scheduling/selector/default.rs rename to lib/kv-router/src/scheduling/selector/reference.rs index ad39dfc005f3..81a415c86ac0 100644 --- a/lib/kv-router/src/scheduling/selector/default.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -261,6 +261,7 @@ fn default_row( worker, preferred_taint_multiplier, WorkerInputs::ALL, + true, |effective_overlap_blocks, device_overlap_blocks| { context.device_overlap(effective_overlap_blocks, device_overlap_blocks) }, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index bae4224752ff..05aa9e7b88d4 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -327,6 +327,7 @@ impl SelectionCore { indexer_threads: usize, cancel_token: CancellationToken, cache_config: SelectionCacheConfig, + policy_factory: WorkerSelectionPolicyFactory, ) -> anyhow::Result { kv_router_config .validate_config() @@ -338,7 +339,7 @@ impl SelectionCore { indexer_threads, cancel_token, None, - None, + Some(policy_factory), SelectionHost::default(), WorkerType::Aggregated, true, diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 3c2af0e9e2fd..b42c9d119623 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -39,6 +39,9 @@ fn local_core_with( indexer_threads, cancel_token, SelectionCacheConfig::default(), + std::sync::Arc::new(|config, role, _| { + crate::WorkerSelectionPolicy::default(config.clone(), role.default_selector_label()) + }), ) .expect("valid test config") } diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 2d87f3706c84..3488ebbfd91f 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -316,10 +316,21 @@ impl SelectionCore { block_size, )) }); - let selector = self.worker_selection_policy_factory.as_ref().map_or_else( - || WorkerSelectionPolicy::default(self.kv_router_config.clone(), worker_label), - |factory| factory(&self.kv_router_config, self.worker_type, key.as_ref()), - ); + let selector = match self.worker_selection_policy_factory.as_ref() { + Some(factory) => { + factory(&self.kv_router_config, self.worker_type, key.as_ref()) + } + #[cfg(test)] + None => { + WorkerSelectionPolicy::default(self.kv_router_config.clone(), worker_label) + } + #[cfg(not(test))] + None => { + return Err(SelectionError::BadRequest( + "routing host must supply a worker-selection policy factory".into(), + )); + } + }; let profile = self .kv_router_config .policy_profile(Some(&key.model_name)) diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index 9158303561fa..e1cb98f3758c 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -218,6 +218,12 @@ mod tests { 1, CancellationToken::new(), SelectionCacheConfig::default(), + std::sync::Arc::new(|config, role, _| { + crate::WorkerSelectionPolicy::default( + config.clone(), + role.default_selector_label(), + ) + }), ) .expect("valid test config"), ) diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 1dc026914029..05332d800a8f 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -163,6 +163,12 @@ impl SelectionServiceBuilder { .plugin_registry .resolve_for_worker_type(&self.kv_router_config, self.worker_type)?, }; + #[cfg(not(test))] + if worker_selection_policy_factory.is_none() { + return Err( + super::policy_registry::WorkerSelectionPolicyRegistryError::MissingDefault.into(), + ); + } let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); let indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; let recover_from_peers = !self.indexer_peers.is_empty(); @@ -317,6 +323,12 @@ impl SelectionService { indexer_threads, cancel_token.clone(), SelectionCacheConfig::default(), + std::sync::Arc::new(|config, role, _| { + crate::WorkerSelectionPolicy::default( + config.clone(), + role.default_selector_label(), + ) + }), ) .expect("valid test config"), ), diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index c847abc86e8f..084ae6e9dadf 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -71,6 +71,7 @@ harness = false required-features = ["media-ffmpeg"] [dependencies] +dynamo-custom-policy-builtin = { workspace = true } # repo dynamo-kv-router = { workspace = true, features = ["metrics", "runtime-protocols", "standalone-selection"] } dynamo-memory = { workspace = true } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 363adf9dd576..fe6aa772ca89 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -9,10 +9,12 @@ use std::{ }; use anyhow::Result; +#[cfg(test)] +use dynamo_kv_router::WorkerSelectionPolicy; use dynamo_kv_router::{ DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionRef, SessionPrefixIndexer, SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, - TrackingHashScope, WorkerSelectionPolicy, WorkerSelectionPolicyFactory, + TrackingHashScope, WorkerSelectionPolicyFactory, config::{KvRouterConfig, RouterConfigOverride, min_initial_workers_from_env}, indexer::{ ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, RoutingDecisionHashes, @@ -123,7 +125,7 @@ impl SelectionPolicySource { { Some(factory) => factory, None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { - WorkerSelectionPolicy::default(config.clone(), label) + dynamo_custom_policy_builtin::default_policy(config.clone(), label) }), }, ), @@ -3231,7 +3233,7 @@ mod tests { let constructions = Arc::clone(&constructions); Arc::new(move |config: &KvRouterConfig, _, _| { constructions.fetch_add(1, Ordering::SeqCst); - WorkerSelectionPolicy::default(config.clone(), "decode") + dynamo_custom_policy_builtin::default_policy(config.clone(), "decode") }) }; let config = KvRouterConfig::default(); diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 39a69a7c1131..f3585409f832 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -9,7 +9,7 @@ //! frontend retains transport, stream leases, and request-expiry ownership. use std::collections::HashMap; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; use std::time::Duration; use anyhow::{Context, Result}; @@ -28,7 +28,7 @@ use dynamo_kv_router::services::selection::{ HostReplication, HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, SelectionOperation, SelectionOutcome, SelectionPartition, SelectionRun, SelectionScheduler, SelectionService, SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, - WorkerRequest, WorkerSelectionPolicyRegistry, + WorkerRequest, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; use tokio_util::sync::CancellationToken; diff --git a/lib/llm/src/kv_router/plugins.rs b/lib/llm/src/kv_router/plugins.rs index 6d6354795702..35f5b0aa802d 100644 --- a/lib/llm/src/kv_router/plugins.rs +++ b/lib/llm/src/kv_router/plugins.rs @@ -14,12 +14,14 @@ static INSTALLED_PLUGINS: OnceLock = OnceLock::new(); /// Install the linked catalog once for frontend and embedded router construction. /// Returns `false` if a catalog has already been installed. pub fn install_router_plugin_registry(registry: RouterPluginRegistry) -> bool { - INSTALLED_PLUGINS.set(registry).is_ok() + INSTALLED_PLUGINS + .set(registry.with_default_factory(dynamo_custom_policy_builtin::default_factory())) + .is_ok() } -/// The installed catalog, or an empty registry for Dynamo's default policies. +/// The installed catalog, or the required builtin default policy. pub fn router_plugin_registry() -> RouterPluginRegistry { - INSTALLED_PLUGINS.get().cloned().unwrap_or_default() + INSTALLED_PLUGINS.get().cloned().unwrap_or_else(dynamo_custom_policy_builtin::default_registry) } /// Carries configured plugins through the host's shared router construction. diff --git a/lib/mocker/Cargo.toml b/lib/mocker/Cargo.toml index f5d35ed54b8d..ce60454d9871 100644 --- a/lib/mocker/Cargo.toml +++ b/lib/mocker/Cargo.toml @@ -22,6 +22,7 @@ test-support = [] zmq-events = ["dep:tmq"] [dependencies] +dynamo-custom-policy-builtin = { workspace = true } # repo aisimulate-core = { workspace = true } dynamo-data-gen = { workspace = true } diff --git a/lib/mocker/src/replay/offline/extensions/kv_router/mod.rs b/lib/mocker/src/replay/offline/extensions/kv_router/mod.rs index 4ea2d4de1258..dda39c90bc30 100644 --- a/lib/mocker/src/replay/offline/extensions/kv_router/mod.rs +++ b/lib/mocker/src/replay/offline/extensions/kv_router/mod.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use dynamo_custom_policy_builtin::DefaultWorkerSelector; use std::borrow::Cow; use std::collections::HashMap; use std::fmt; @@ -22,10 +23,9 @@ use dynamo_kv_router::scheduling::{ }; use dynamo_kv_router::sequences::topology::WorkerDpRange; use dynamo_kv_router::{ - ActiveSequencesMultiWorker, DefaultWorkerSelector, RadixTree, RoutingPartitionRef, - SchedulingRequest, SequenceRequest, SessionContext, TrackingHashAlgorithm, TrackingHashContext, - TrackingHashScope, WorkerLoadProjection, WorkerSelectionInput, WorkerSelector, - scheduling::TierOverlapBlocks, + ActiveSequencesMultiWorker, RadixTree, RoutingPartitionRef, SchedulingRequest, SequenceRequest, + SessionContext, TrackingHashAlgorithm, TrackingHashContext, TrackingHashScope, + WorkerLoadProjection, WorkerSelectionInput, WorkerSelector, scheduling::TierOverlapBlocks, }; use dynamo_tokens::SequenceHash; use rustc_hash::FxHashMap; diff --git a/lib/mocker/src/replay/router_shared.rs b/lib/mocker/src/replay/router_shared.rs index 5859782f94b5..363af69eba6f 100644 --- a/lib/mocker/src/replay/router_shared.rs +++ b/lib/mocker/src/replay/router_shared.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use dynamo_custom_policy_builtin::DefaultWorkerSelector; use std::collections::HashMap; use std::sync::Arc; @@ -11,9 +12,7 @@ use dynamo_kv_router::protocols::{ }; use dynamo_kv_router::scheduling::queue::DEFAULT_MAX_BATCHED_TOKENS; use dynamo_kv_router::sequences::SchedulerLoadSnapshot; -use dynamo_kv_router::{ - ActiveSequencesMultiWorker, DefaultWorkerSelector, LocalScheduler, SequencePublisher, -}; +use dynamo_kv_router::{ActiveSequencesMultiWorker, LocalScheduler, SequencePublisher}; #[derive(Clone, Copy, Debug, Default)] pub(super) struct ReplayNoopPublisher; diff --git a/lib/router-plugins/builtin/Cargo.toml b/lib/router-plugins/builtin/Cargo.toml index 68fdb47c1e13..6fbd92348367 100644 --- a/lib/router-plugins/builtin/Cargo.toml +++ b/lib/router-plugins/builtin/Cargo.toml @@ -12,8 +12,18 @@ repository.workspace = true description = "Worker-selection policies Dynamo ships and always links" [dependencies] +fastrand = "2" +parking_lot.workspace = true dynamo-kv-router = { workspace = true, features = ["standalone-selection"] } serde.workspace = true [dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt", "time"] } +dynamo-kv-router = { workspace = true, features = ["bench"] } +rustc-hash = "2.1.1" +criterion = "0.5" tempfile = { workspace = true } + +[[bench]] +name = "default_selection" +harness = false diff --git a/lib/router-plugins/builtin/benches/default_selection.rs b/lib/router-plugins/builtin/benches/default_selection.rs new file mode 100644 index 000000000000..74fb0bfc4b3a --- /dev/null +++ b/lib/router-plugins/builtin/benches/default_selection.rs @@ -0,0 +1,121 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#[path = "../tests/support/mod.rs"] +mod support; +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use dynamo_custom_policy_builtin::default_policy; +use dynamo_kv_router::{KvRouterConfig, WorkerSelectionInput, WorkerSelector}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; +struct CountingAllocator; +static COUNTING: AtomicBool = AtomicBool::new(false); +static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0); +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; +// SAFETY: every operation delegates to System with the original pointer and layout. +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if COUNTING.load(Ordering::Relaxed) { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + } + unsafe { System.alloc(layout) } + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, size: usize) -> *mut u8 { + if COUNTING.load(Ordering::Relaxed) { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + } + unsafe { System.realloc(ptr, layout, size) } + } +} +fn allocations(mut select: impl FnMut()) -> usize { + select(); // Warm retained buffers and thread-local RNG. + ALLOCATIONS.store(0, Ordering::Relaxed); + COUNTING.store(true, Ordering::Relaxed); + for _ in 0..100 { + select(); + } + COUNTING.store(false, Ordering::Relaxed); + ALLOCATIONS.load(Ordering::Relaxed) +} + +fn bench(c: &mut Criterion) { + for temperature in [0.0, 0.7] { + let mut group = c.benchmark_group(format!("default_selection/t{temperature}")); + group + .warm_up_time(Duration::from_millis(200)) + .measurement_time(Duration::from_millis(500)) + .sample_size(20); + for count in [8, 64, 256, 1024] { + let (workers, request) = support::fixture(count, 2048); + let config = KvRouterConfig { + router_temperature: temperature, + ..Default::default() + }; + let reference = + dynamo_kv_router::DefaultWorkerSelector::new(Some(config.clone()), "prefill"); + let plugin = default_policy(config, "prefill"); + let reference_allocs = allocations(|| { + black_box( + reference + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(), + ); + }); + let plugin_allocs = allocations(|| { + black_box( + plugin + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(), + ); + }); + eprintln!( + "allocations per 100 warm selections: temperature={temperature} workers={count} reference={reference_allocs} plugin={plugin_allocs}" + ); + group.bench_function(BenchmarkId::new("reference", count), |b| { + b.iter(|| { + black_box( + reference + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(), + ) + }) + }); + group.bench_function(BenchmarkId::new("plugin", count), |b| { + b.iter(|| { + black_box( + plugin + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(), + ) + }) + }); + } + group.finish(); + } +} +criterion_group!(benches, bench); +criterion_main!(benches); diff --git a/lib/router-plugins/builtin/src/default.rs b/lib/router-plugins/builtin/src/default.rs new file mode 100644 index 000000000000..cde4406c1805 --- /dev/null +++ b/lib/router-plugins/builtin/src/default.rs @@ -0,0 +1,492 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The default KV policy, implemented solely against the public worker-policy API. + +use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerSelectionResult}; +use dynamo_kv_router::{ + KvRouterConfig, KvSchedulerError, WorkerInputView, WorkerInputs, WorkerPicker, + WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, + WorkerSelectionPolicyError, WorkerSelectionPolicyFactory, WorkerSelector, +}; +use parking_lot::Mutex; +use std::sync::Arc; + +fn softmax_sample_index( + entries: &[T], + cost: impl Fn(&T) -> f64, + temperature: f64, + sample: f64, + probabilities: &mut Vec, +) -> usize { + assert!(!entries.is_empty(), "Empty entries for softmax sampling"); + debug_assert_ne!(temperature, 0.0); + + let (min_cost, max_cost) = entries + .iter() + .map(&cost) + .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), cost| { + (lo.min(cost), hi.max(cost)) + }); + + probabilities.clear(); + if min_cost == max_cost { + probabilities.resize(entries.len(), 1.0 / entries.len() as f64); + } else { + let range = max_cost - min_cost; + let magnitude = if range.is_finite() { + 1.0 + } else { + min_cost.abs().max(max_cost.abs()) + }; + let min_normalized = min_cost / magnitude; + let scale = -1.0 / ((max_cost / magnitude - min_normalized) * temperature); + let max_scaled = min_normalized * scale; + probabilities.extend( + entries + .iter() + .map(|entry| (cost(entry) / magnitude * scale - max_scaled).exp()), + ); + } + + let sum: f64 = probabilities.iter().sum(); + for probability in probabilities.iter_mut() { + *probability /= sum; + } + let mut cumulative = 0.0; + for (row, probability) in probabilities.iter().enumerate() { + cumulative += probability; + if sample <= cumulative { + return row; + } + } + entries.len() - 1 +} + +/// Construct the builtin default from configured policy parameters. +/// Per-request score overrides are not used. Request load-tracking remains host-owned. +pub fn default_policy(config: KvRouterConfig, worker_label: &'static str) -> WorkerSelectionPolicy { + policy_with_rng(config, worker_label, None) +} + +fn policy_with_rng( + config: KvRouterConfig, + worker_label: &'static str, + rng: Option>>, +) -> WorkerSelectionPolicy { + let picker = DefaultPicker { + plain_decode: false, + config: config.clone(), + worker_label, + rng, + entries: Vec::new(), + probabilities: Vec::new(), + }; + WorkerSelectionPolicy::new(config, worker_label, Vec::new(), Box::new(picker)) + .with_exclusive_affinity(true) +} + +/// Factory installed by routing hosts, including hosts without a custom catalog. +pub fn default_factory() -> WorkerSelectionPolicyFactory { + Arc::new(|config, role, _partition| policy_for_role(config.clone(), role)) +} + +fn policy_for_role( + config: KvRouterConfig, + role: dynamo_kv_router::WorkerType, +) -> WorkerSelectionPolicy { + let picker = DefaultPicker { + plain_decode: role == dynamo_kv_router::WorkerType::Decode + && !config.conditional_disagg_enabled, + config: config.clone(), + worker_label: role.default_selector_label(), + rng: None, + entries: Vec::new(), + probabilities: Vec::new(), + }; + WorkerSelectionPolicy::new( + config, + role.default_selector_label(), + Vec::new(), + Box::new(picker), + ) + .with_exclusive_affinity(true) +} + +struct DefaultPicker { + plain_decode: bool, + config: KvRouterConfig, + worker_label: &'static str, + rng: Option>>, + entries: Vec<(usize, f64)>, + probabilities: Vec, +} + +impl DefaultPicker { + fn choose( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result<(usize, Option), WorkerSelectionPolicyError> { + let cache = input + .cache() + .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; + let load = input + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + let candidates = input.candidates(); + if candidates.is_empty() { + return Err(WorkerSelectionPolicyError::failed("no eligible worker")); + } + let min_prefill = + if context.tracks_prefill_tokens() && self.config.overlap_score_credit_decay > 0.0 { + load.iter() + .map(|l| l.active_prefill_tokens()) + .min() + .unwrap_or(0) + } else { + 0 + }; + self.entries.clear(); + let direct_minimum = self.config.router_temperature == 0.0 && self.rng.is_none(); + let mut best_row = 0; + let mut best_cost = f64::INFINITY; + let mut ties = 0; + for (row, candidate) in candidates.iter().enumerate() { + let cache = &cache[row]; + let load = &load[row]; + let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); + let device = if context.has_tier_matches() { + cache.device_overlap_blocks() + } else { + estimated_overlap + }; + let shared = if context.has_tier_matches() { + cache.shared_beyond_device_blocks() + } else { + context + .shared_cache_hits() + .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)) + }; + let decay = if context.tracks_prefill_tokens() + && self.config.overlap_score_credit_decay > 0.0 + { + let excess = load.active_prefill_tokens().saturating_sub(min_prefill) as f64 + / context.block_size() as f64; + 1.0 / (1.0 + + self.config.overlap_score_credit_decay + * (excess / context.request_blocks() as f64)) + } else { + 1.0 + }; + // Plain disaggregated decode is load-only. Conditional decode retains cache credit. + let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { + 0.0 + } else { + self.config.overlap_score_credit + }; + let credit = overlap_credit * decay * device + + self.config.host_cache_hit_weight * cache.host_overlap_blocks() + + self.config.disk_cache_hit_weight * cache.disk_overlap_blocks() + + self.config.shared_cache_multiplier * shared as f64; + let request_cost = + self.config.decode_active_request_weight * load.active_requests() as f64; + let logit = if self.worker_label == "decode" + && !context.tracks_prefill_tokens() + && overlap_credit > 0.0 + { + (load.decode_cost_blocks() - credit).max(0.0) + request_cost + } else { + let raw_tokens = if !context.tracks_prefill_tokens() { + 0 + } else if load.is_available() { + let uncached = context.prompt_tokens().saturating_sub(cached_tokens); + (load.active_prefill_tokens() + uncached).saturating_add(cached_tokens) + } else { + context.prompt_tokens() + }; + let prefill = (raw_tokens as f64 / context.block_size() as f64 - credit).max(0.0); + self.config.prefill_load_scale * prefill + load.decode_cost_blocks() + request_cost + }; + let cost = logit * candidate.preferred_taint_multiplier().unwrap_or(1.0); + if !cost.is_finite() { + return Err(WorkerSelectionPolicyError::failed( + "default policy produced non-finite cost", + )); + } + if direct_minimum { + if cost < best_cost { + best_row = row; + best_cost = cost; + ties = 1; + } else if cost == best_cost { + ties += 1; + if fastrand::usize(0..ties) == 0 { + best_row = row; + } + } + } else { + self.entries.push((row, cost)); + } + } + if direct_minimum { + return Ok((best_row, Some(best_cost))); + } + if context.pinned_worker().is_some() { + let (row, cost) = self.entries[0]; + return Ok((row, Some(cost))); + } + // Canonical order is required only for deterministic replay, never for production ties. + if self.rng.is_some() { + self.entries.sort_unstable_by_key(|(row, _)| { + let worker = candidates[*row].worker(); + (worker.worker_id, worker.dp_rank) + }); + } + let mut rng = self.rng.as_ref().map(|rng| rng.lock()); + let selected = if self.config.router_temperature == 0.0 { + let mut best = 0; + let mut best_cost = f64::INFINITY; + let mut ties = 0; + for (index, (_, cost)) in self.entries.iter().enumerate() { + if *cost < best_cost { + best = index; + best_cost = *cost; + ties = 1; + } else if *cost == best_cost { + ties += 1; + let sample = match rng.as_mut() { + Some(rng) => rng.usize(0..ties), + None => fastrand::usize(0..ties), + }; + if sample == 0 { + best = index; + } + } + } + best + } else { + let sample = match rng.as_mut() { + Some(rng) => rng.f64(), + None => fastrand::f64(), + }; + softmax_sample_index( + &self.entries, + |(_, cost)| *cost, + self.config.router_temperature, + sample, + &mut self.probabilities, + ) + }; + let (row, cost) = self.entries[selected]; + Ok((row, Some(cost))) + } +} + +impl WorkerPicker for DefaultPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT + } + fn pick( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + self.choose(context, input).map(|(row, _)| row) + } + fn pick_with_cost( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result<(usize, Option), WorkerSelectionPolicyError> { + self.choose(context, input) + } +} + +/// Synchronous adapter for direct callers and replay. Scheduler actors use `default_policy` +/// directly and do not acquire this adapter's mutex. +pub struct DefaultWorkerSelector { + kv_router_config: KvRouterConfig, + worker_type: &'static str, + policy: Mutex, + rng: Option>>, +} +impl DefaultWorkerSelector { + pub fn new(config: Option, worker_type: &'static str) -> Self { + Self::with_rng(config.unwrap_or_default(), worker_type, None) + } + /// Construct a reproducible selector. Clones share its random stream, as before. + pub fn new_seeded( + config: Option, + worker_type: &'static str, + seed: u64, + ) -> Self { + Self::with_rng( + config.unwrap_or_default(), + worker_type, + Some(Arc::new(Mutex::new(fastrand::Rng::with_seed(seed)))), + ) + } + fn with_rng( + config: KvRouterConfig, + worker_type: &'static str, + rng: Option>>, + ) -> Self { + Self { + policy: Mutex::new(policy_with_rng(config.clone(), worker_type, rng.clone())), + kv_router_config: config, + worker_type, + rng, + } + } +} +impl Clone for DefaultWorkerSelector { + fn clone(&self) -> Self { + Self::with_rng( + self.kv_router_config.clone(), + self.worker_type, + self.rng.clone(), + ) + } +} +impl std::fmt::Debug for DefaultWorkerSelector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DefaultWorkerSelector") + .field("kv_router_config", &self.kv_router_config) + .field("worker_type", &self.worker_type) + .finish_non_exhaustive() + } +} +impl WorkerSelector for DefaultWorkerSelector { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD + } + fn uses_exclusive_affinity_target(&self) -> bool { + true + } + fn select_worker( + &self, + input: WorkerSelectionInput<'_, C>, + ) -> Result { + self.policy.lock().select_worker(input) + } +} + +#[derive(Default, serde::Deserialize)] +#[serde(default, deny_unknown_fields)] +struct Parameters { + overlap_score_credit: Option, + overlap_score_credit_decay: Option, + prefill_load_scale: Option, + decode_active_request_weight: Option, + host_cache_hit_weight: Option, + disk_cache_hit_weight: Option, + shared_cache_multiplier: Option, + router_temperature: Option, +} + +pub(super) fn register( + registry: &mut dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistry, +) -> Result<(), dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistryError> { + use dynamo_kv_router::services::selection::WorkerSelectionPolicyProviderError; + registry.register( + "dynamo-default-cost-fn", + Arc::new(|parameters| { + let parameters: Parameters = parameters.deserialize()?; + if parameters + .overlap_score_credit + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "overlap_score_credit must be finite and non-negative", + )); + } + if parameters + .overlap_score_credit_decay + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "overlap_score_credit_decay must be finite and non-negative", + )); + } + if parameters + .prefill_load_scale + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "prefill_load_scale must be finite and non-negative", + )); + } + if parameters + .decode_active_request_weight + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "decode_active_request_weight must be finite and non-negative", + )); + } + if parameters + .host_cache_hit_weight + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "host_cache_hit_weight must be finite and non-negative", + )); + } + if parameters + .disk_cache_hit_weight + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "disk_cache_hit_weight must be finite and non-negative", + )); + } + if parameters + .shared_cache_multiplier + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "shared_cache_multiplier must be finite and non-negative", + )); + } + if parameters + .router_temperature + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err(WorkerSelectionPolicyProviderError::new( + "router_temperature must be finite and non-negative", + )); + } + Ok(Arc::new( + move |config: &KvRouterConfig, role, _partition| { + let mut config = config.clone(); + if let Some(value) = parameters.overlap_score_credit { + config.overlap_score_credit = value; + } + if let Some(value) = parameters.overlap_score_credit_decay { + config.overlap_score_credit_decay = value; + } + if let Some(value) = parameters.prefill_load_scale { + config.prefill_load_scale = value; + } + if let Some(value) = parameters.decode_active_request_weight { + config.decode_active_request_weight = value; + } + if let Some(value) = parameters.host_cache_hit_weight { + config.host_cache_hit_weight = value; + } + if let Some(value) = parameters.disk_cache_hit_weight { + config.disk_cache_hit_weight = value; + } + if let Some(value) = parameters.shared_cache_multiplier { + config.shared_cache_multiplier = value; + } + if let Some(value) = parameters.router_temperature { + config.router_temperature = value; + } + policy_for_role(config, role) + }, + )) + }), + ) +} diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index 06e8993e9303..c6656c09c8a8 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -3,22 +3,19 @@ //! Worker-selection policies Dynamo ships. //! -//! Any build that enables the Python bindings' `custom-policy` feature links these, and that -//! feature is on by default, so a frontend deployment selects one through router-policy YAML -//! without rebuilding. A build with default features off, and the standalone EPP, link no catalog -//! at all and reject a configured policy type at startup. -//! -//! These do not go through the replaceable `dynamo-worker-selection-policy-catalog` alias, and they -//! register before it, so a custom image's catalog adds its policies alongside these rather than -//! displacing them. -//! -//! Each policy is one module here. To ship another, add a module and a line in [`register`], then -//! add a row to the policy table in the router configuration guide. Keep a policy in a single file -//! until it needs submodules, then promote it to a directory. If a policy ever needs a dependency -//! beyond `dynamo-kv-router`, put it behind its own default-on Cargo feature so a build can drop -//! it; every policy registered here is compiled into every artifact that links this crate. +//! Routing hosts always link the default through `default_registry`. The optional custom +//! catalog adds the named default and two-tier providers through `register`. The default +//! itself uses the same public candidate inputs and picker dispatch as external policies. +//! Sequence tracking, eligibility, and admission remain in dynamo-kv-router. +mod default; mod two_tier_cost_fn; +pub use default::{DefaultWorkerSelector, default_factory, default_policy}; + +/// Registry containing the required default only, without an optional policy catalog. +pub fn default_registry() -> RouterPluginRegistry { + RouterPluginRegistry::new(default_factory()) +} use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegistryError}; @@ -30,6 +27,8 @@ use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegis pub fn register( registry: &mut RouterPluginRegistry, ) -> Result<(), WorkerSelectionPolicyRegistryError> { + *registry = std::mem::take(registry).with_default_factory(default_factory()); + default::register(registry)?; two_tier_cost_fn::register(registry) } diff --git a/lib/router-plugins/builtin/tests/default_lifecycle.rs b/lib/router-plugins/builtin/tests/default_lifecycle.rs new file mode 100644 index 000000000000..9cb20387d43e --- /dev/null +++ b/lib/router-plugins/builtin/tests/default_lifecycle.rs @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use dynamo_custom_policy_builtin::default_factory; +use dynamo_kv_router::protocols::RoutingConstraints; +use dynamo_kv_router::services::selection::{ + PromptRequest, SelectAndReserveRequest, SelectionCacheConfig, SelectionCore, WorkerRequest, +}; +use dynamo_kv_router::{KvRouterConfig, RoutingPartitionId}; +use std::collections::{HashMap, HashSet}; + +fn worker(worker_id: u64) -> WorkerRequest { + WorkerRequest { + worker_id, + model_name: "model".to_string(), + routing_group: "default".to_string(), + endpoint: Some(format!("http://worker-{worker_id}:8000")), + kv_events_endpoint: None, + kv_events_endpoints: HashMap::new(), + replay_endpoint: None, + block_size: Some(4), + data_parallel_start_rank: None, + data_parallel_size: None, + max_num_batched_tokens: Some(1024), + total_kv_blocks: None, + stable_routing_id: None, + is_eagle: None, + taints: HashSet::new(), + topology_domains: HashMap::new(), + kv_transfer_domain: None, + kv_transfer_enforcement: None, + kv_transfer_preferred_weight: None, + router_hint_worker_type: None, + router_hint_source_control_endpoints: HashMap::new(), + kv_event_source_mode: None, + } +} + +fn reserve_request(selection_id: &str) -> SelectAndReserveRequest { + SelectAndReserveRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: Some(selection_id.to_string()), + prompt: PromptRequest { + token_ids: Some(vec![1, 2, 3, 4, 5]), + ..Default::default() + }, + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: None, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + } +} + +#[tokio::test] +async fn shared_core_books_and_releases_with_the_builtin_default() { + let core = SelectionCore::try_new_local( + KvRouterConfig { + use_kv_events: false, + router_queue_threshold: None, + ..Default::default() + }, + 1, + Default::default(), + SelectionCacheConfig::default(), + default_factory(), + ) + .unwrap(); + core.upsert_worker(worker(1)).await.unwrap(); + let partition = core + .partition(&RoutingPartitionId::new("model", "default")) + .unwrap(); + assert!(!partition.scheduler().has_request("request")); + core.select_and_reserve(reserve_request("request")) + .await + .unwrap(); + assert!(partition.scheduler().has_request("request")); + core.prefill_complete("request").await.unwrap(); + core.free_reservation("request").await.unwrap(); + assert!(!partition.scheduler().has_request("request")); + assert!(core.free_reservation("request").await.is_err()); + core.shutdown(); +} diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs new file mode 100644 index 000000000000..34538b1b2b14 --- /dev/null +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -0,0 +1,245 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +mod support; +use dynamo_custom_policy_builtin::{DefaultWorkerSelector, default_policy, default_registry}; +use dynamo_kv_router::protocols::WorkerWithDpRank; +use dynamo_kv_router::{ + KvRouterConfig, RoutingPartitionRef, WorkerInputView, WorkerInputs, WorkerPicker, + WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, + WorkerSelectionPolicyError, WorkerSelector, WorkerType, +}; +use support::*; + +#[test] +fn seeded_selection_matches_reference_across_cache_and_load_shapes() { + for temperature in [0.0, 0.7] { + for prompt in [1, 17, 127, 2048] { + for mode in 0..4 { + let (workers, mut request) = fixture(16, prompt); + let config = KvRouterConfig { + router_temperature: temperature, + overlap_score_credit_decay: 0.6, + host_cache_hit_weight: 0.25, + disk_cache_hit_weight: 0.1, + ..Default::default() + }; + match mode { + 1 => request.overlap.tier_overlap_blocks = Default::default(), + 2 => request.worker_loads.clear(), + 3 => request.track_prefill_tokens = false, + _ => {} + } + let reference = dynamo_kv_router::DefaultWorkerSelector::new_seeded( + Some(config.clone()), + "test", + 42, + ); + let plugin = DefaultWorkerSelector::new_seeded(Some(config), "test", 42); + for _ in 0..64 { + let input = WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + ); + let expected = reference.select_worker(input).unwrap(); + let actual = plugin.select_worker(input).unwrap(); + assert_eq!( + actual.worker, expected.worker, + "temperature={temperature} prompt={prompt} mode={mode}" + ); + assert_eq!(actual.cached_tokens, expected.cached_tokens); + assert_eq!( + actual.potential_decode_blocks, + expected.potential_decode_blocks + ); + } + } + } + } +} + +#[test] +fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { + struct Inspect; + impl WorkerPicker for Inspect { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD + } + fn pick( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + assert_eq!(context.prompt_tokens(), 17); + assert_eq!(context.request_blocks(), 2); + assert!(context.has_tier_matches()); + assert!(input.load().unwrap().iter().all(|load| load.is_available())); + for cache in input.cache().unwrap() { + let (blocks, tokens) = cache.accounting_cache_estimate(); + assert_eq!((blocks * 16.0) as usize, tokens); + } + Ok(0) + } + } + let (workers, request) = fixture(2, 17); + WorkerSelectionPolicy::new(KvRouterConfig::default(), "test", vec![], Box::new(Inspect)) + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); +} + +#[test] +fn configured_default_resolves_for_every_role_and_uses_exclusive_affinity() { + let registry = default_registry(); + for role in [ + WorkerType::Aggregated, + WorkerType::Prefill, + WorkerType::Decode, + WorkerType::Encode, + ] { + let config = KvRouterConfig::default(); + let factory = registry + .resolve_for_worker_type(&config, role) + .unwrap() + .unwrap(); + let policy = factory(&config, role, RoutingPartitionRef::new("model", "default")); + assert!( + >::uses_exclusive_affinity_target( + &policy + ) + ); + } +} + +#[test] +fn mandatory_pin_is_preserved() { + let (workers, mut request) = fixture(4, 17); + request.pinned_worker = Some(WorkerWithDpRank::new(3, 1)); + let selected = default_policy(KvRouterConfig::default(), "test") + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + assert_eq!(Some(selected.worker), request.pinned_worker); +} + +#[test] +fn non_finite_picker_cost_is_rejected() { + struct Invalid; + impl WorkerPicker for Invalid { + fn pick( + &mut self, + _: &WorkerSelectionContext<'_>, + _: WorkerInputView<'_>, + ) -> Result { + Ok(0) + } + fn pick_with_cost( + &mut self, + _: &WorkerSelectionContext<'_>, + _: WorkerInputView<'_>, + ) -> Result<(usize, Option), WorkerSelectionPolicyError> { + Ok((0, Some(f64::NAN))) + } + } + let (workers, request) = fixture(1, 17); + let error = + WorkerSelectionPolicy::new(KvRouterConfig::default(), "test", vec![], Box::new(Invalid)) + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap_err(); + assert!(error.to_string().contains("non-finite")); +} + +#[test] +fn configured_parameters_replace_request_score_overrides() { + use dynamo_kv_router::RouterConfigOverride; + let (workers, mut request) = fixture(2, 160); + for (worker, load) in &mut request.worker_loads { + load.active_prefill_tokens = 0; + load.active_requests = 0; + load.active_decode_blocks = if worker.worker_id == 0 { 8 } else { 0 }; + load.additional_active_blocks = 0; + request + .overlap + .tier_overlap_blocks + .device + .insert(*worker, if worker.worker_id == 0 { 10 } else { 0 }); + } + request.router_config_override = Some(RouterConfigOverride { + overlap_score_credit: Some(0.0), + router_temperature: Some(10.0), + ..Default::default() + }); + let config = KvRouterConfig { + overlap_score_credit: 2.0, + host_cache_hit_weight: 0.0, + disk_cache_hit_weight: 0.0, + router_temperature: 0.0, + ..Default::default() + }; + let policy = default_policy(config, "prefill"); + for _ in 0..32 { + let result = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + assert_eq!(result.worker.worker_id, 0); + } +} + +#[test] +fn named_default_parameters_resolve_and_reject_invalid_values() { + use std::io::Write; + for (parameter, succeeds) in [ + ("overlap_score_credit: 2.0", true), + ("overlap_score_credit: -1.0", false), + ("unknown: 1", false), + ] { + let mut file = tempfile::NamedTempFile::new().unwrap(); + write!(file, "worker_selection:\n aggregated: tuned\n instances:\n - name: tuned\n type: dynamo-default-cost-fn\n parameters:\n {parameter}\n").unwrap(); + let config = KvRouterConfig { + router_policy_config: Some(file.path().display().to_string()), + ..Default::default() + }; + let mut registry = default_registry(); + dynamo_custom_policy_builtin::register(&mut registry).unwrap(); + assert_eq!(registry.resolve(&config).is_ok(), succeeds); + } +} + +#[test] +fn pin_does_not_advance_seeded_random_stream() { + let (workers, mut request) = fixture(8, 127); + let config = KvRouterConfig { + router_temperature: 0.7, + ..Default::default() + }; + let reference = + dynamo_kv_router::DefaultWorkerSelector::new_seeded(Some(config.clone()), "prefill", 42); + let plugin = DefaultWorkerSelector::new_seeded(Some(config), "prefill", 42); + for pinned in [true, false, true, false, false] { + request.pinned_worker = pinned.then_some(WorkerWithDpRank::new(3, 1)); + let input = WorkerSelectionInput::configured(&workers, &request, request.eligibility(), 16); + assert_eq!( + reference.select_worker(input).unwrap().worker, + plugin.select_worker(input).unwrap().worker + ); + } +} diff --git a/lib/router-plugins/builtin/tests/support/mod.rs b/lib/router-plugins/builtin/tests/support/mod.rs new file mode 100644 index 000000000000..626dafd225ad --- /dev/null +++ b/lib/router-plugins/builtin/tests/support/mod.rs @@ -0,0 +1,84 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use dynamo_kv_router::WorkerLoadProjection; +use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank}; +use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode, SchedulingRequest}; +use std::collections::HashMap; + +#[derive(Clone, Copy)] +pub struct TestWorker; +impl WorkerConfigLike for TestWorker { + fn data_parallel_start_rank(&self) -> u32 { + 0 + } + fn data_parallel_size(&self) -> u32 { + 2 + } + fn max_num_batched_tokens(&self) -> Option { + None + } + fn total_kv_blocks(&self) -> Option { + Some(16384) + } +} + +pub fn fixture( + count: usize, + prompt_tokens: usize, +) -> (HashMap, SchedulingRequest) { + let mut request = SchedulingRequest { + mode: ScheduleMode::QueryOnly { request_id: None }, + token_seq: None, + isl_tokens: prompt_tokens, + lora_name: None, + expected_output_tokens: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + router_config_override: None, + track_prefill_tokens: true, + priority_jump: 0.0, + strict_priority: 0, + policy_class: None, + session_context: None, + overlap: OverlapSignals::default(), + kv_transfer_candidates: None, + retain_kv_transfer_chain: false, + shared_cache_hits: None, + worker_loads: Default::default(), + resp_tx: None, + }; + let workers = (0..count as u64).map(|id| (id, TestWorker)).collect(); + for id in 0..count as u64 { + for dp_rank in 0..2 { + let w = WorkerWithDpRank::new(id, dp_rank); + let overlap = (id as usize * 7 + dp_rank as usize) % 9; + request + .overlap + .tier_overlap_blocks + .device + .insert(w, overlap); + request.overlap.tier_overlap_blocks.host_pinned.insert(w, 2); + request.overlap.tier_overlap_blocks.disk.insert(w, 1); + request + .overlap + .effective_overlap_blocks + .insert(w, overlap as f64 + 0.5); + request + .overlap + .effective_cached_tokens + .insert(w, overlap * 16 + 8); + request.worker_loads.insert( + w, + WorkerLoadProjection { + active_requests: id as usize % 5, + active_prefill_tokens: (id as usize % 7) * 19, + active_decode_blocks: id as usize % 11, + additional_active_blocks: 3, + }, + ); + } + } + (workers, request) +} From cd1abdc2b6d4ac727397869bd14baf233159a7c1 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 14 Sep 2026 07:21:20 +0000 Subject: [PATCH 02/31] perf(router): simplify builtin selection and retain cache estimate maps Signed-off-by: Ishan Dhanani --- lib/kv-router/benches/worker_selection.rs | 4 +- lib/kv-router/src/plugins/registry.rs | 6 - lib/kv-router/src/plugins/worker_selection.rs | 8 -- lib/kv-router/src/scheduling/local.rs | 94 +++++++------- lib/kv-router/src/scheduling/overlap.rs | 10 +- .../src/scheduling/overlap_refresh.rs | 4 +- lib/kv-router/src/scheduling/queue.rs | 12 +- lib/kv-router/src/scheduling/selector/mod.rs | 60 ++------- .../src/scheduling/selector/reference.rs | 97 ++++++++------ lib/kv-router/src/scheduling/types.rs | 4 +- .../src/services/selection/core/queries.rs | 2 +- lib/llm/src/kv_router.rs | 2 +- lib/router-plugins/builtin/Cargo.toml | 4 + .../builtin/benches/default_lifecycle.rs | 92 ++++++++++++++ .../builtin/benches/default_selection.rs | 8 +- lib/router-plugins/builtin/src/default.rs | 118 ++++++------------ lib/router-plugins/builtin/src/lib.rs | 2 +- .../builtin/tests/default_policy.rs | 11 +- 18 files changed, 286 insertions(+), 252 deletions(-) create mode 100644 lib/router-plugins/builtin/benches/default_lifecycle.rs diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index a0e9a51d28f1..0ed136eedaf4 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -145,8 +145,8 @@ fn fixture_with_preferred_taints( const PREFERRED_TAINTS: [&str; 4] = ["rack-a", "zone-a", "gpu-a", "node-a"]; assert!(preferred_taint_count <= PREFERRED_TAINTS.len()); let mut workers = HashMap::with_capacity(worker_count); - let mut effective_overlap_blocks = HashMap::with_capacity(worker_count); - let mut effective_cached_tokens = HashMap::with_capacity(worker_count); + let mut effective_overlap_blocks = FxHashMap::default(); + let mut effective_cached_tokens = FxHashMap::default(); let mut worker_loads = FxHashMap::with_capacity_and_hasher(worker_count, Default::default()); for worker_id in 0..worker_count as WorkerId { diff --git a/lib/kv-router/src/plugins/registry.rs b/lib/kv-router/src/plugins/registry.rs index 137a5c2e943a..ba0e811364a4 100644 --- a/lib/kv-router/src/plugins/registry.rs +++ b/lib/kv-router/src/plugins/registry.rs @@ -104,12 +104,6 @@ pub enum WorkerSelectionPolicyRegistryError { } impl RouterPluginRegistry { - /// Construct a registry with the host's required default policy. Custom registrations - /// cannot replace this factory. An empty registry still supports explicit custom policies. - pub fn new(default_factory: WorkerSelectionPolicyFactory) -> Self { - Self { default_factory: Some(default_factory), ..Self::default() } - } - /// Supply the host default while preserving registered custom plugin types. pub fn with_default_factory(mut self, factory: WorkerSelectionPolicyFactory) -> Self { self.default_factory = Some(factory); diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index 4b83fc85d364..ce7b347b5cd9 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -21,8 +21,6 @@ use std::sync::Arc; use crate::protocols::{WorkerAffinityTarget, WorkerWithDpRank}; use crate::scheduling::SchedulingRequest; -#[cfg(any(test, feature = "bench"))] -use crate::scheduling::selector::LogitWeights; use crate::{KvRouterConfig, RoutingPartitionRef, WorkerType}; /// Factory that creates one worker-selection policy per routing partition. @@ -35,13 +33,9 @@ pub type WorkerSelectionPolicyFactory = Arc< /// Request-level values available to custom filters, scorers, and pickers. pub struct WorkerSelectionContext<'a> { pub(crate) request: &'a SchedulingRequest, - #[cfg(any(test, feature = "bench"))] - pub(crate) request_id: &'a str, pub(crate) request_blocks: u64, pub(crate) block_size: u32, pub(crate) track_prefill_tokens: bool, - #[cfg(any(test, feature = "bench"))] - pub(crate) weights: LogitWeights, pub(crate) has_tier_matches: bool, pub(crate) pinned_worker: Option, pub(crate) router_temperature_override: Option, @@ -113,8 +107,6 @@ pub struct WorkerCacheInput { /// Active-load values for one worker. #[derive(Clone, Copy, Default)] pub struct WorkerLoadInput { - #[cfg(any(test, feature = "bench"))] - pub(crate) raw_prefill_blocks: f64, pub(crate) available: bool, pub(crate) active_prefill_tokens: usize, pub(crate) decode_cost_blocks: f64, diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index c8103f377cc1..00e8f82d8b3f 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -455,8 +455,8 @@ where isl_tokens: usize, token_seq: Option>, tier_overlap_blocks: TierOverlapBlocks, - effective_overlap_blocks: HashMap, - effective_cached_tokens: HashMap, + effective_overlap_blocks: FxHashMap, + effective_cached_tokens: FxHashMap, router_config_override: Option<&super::config::RouterConfigOverride>, update_states: bool, lora_name: Option, @@ -502,8 +502,8 @@ where token_seq: Option>, block_hashes: Option>, tier_overlap_blocks: TierOverlapBlocks, - effective_overlap_blocks: HashMap, - effective_cached_tokens: HashMap, + effective_overlap_blocks: FxHashMap, + effective_cached_tokens: FxHashMap, router_config_override: Option<&super::config::RouterConfigOverride>, update_states: bool, lora_name: Option, @@ -546,8 +546,8 @@ where token_seq: Option>, block_hashes: Option>, tier_overlap_blocks: TierOverlapBlocks, - effective_overlap_blocks: HashMap, - effective_cached_tokens: HashMap, + effective_overlap_blocks: FxHashMap, + effective_cached_tokens: FxHashMap, router_config_override: Option<&super::config::RouterConfigOverride>, update_states: bool, lora_name: Option, @@ -814,7 +814,7 @@ where &self, token_seq: Option>, isl_tokens: usize, - effective_cached_tokens: HashMap, + effective_cached_tokens: FxHashMap, track_prefill_tokens: bool, ) -> Vec { let decay_now = Instant::now(); @@ -1044,7 +1044,8 @@ mod tests { scheduler.get_active_lora_counts(), HashMap::from([(String::from("adapter-a"), 1)]) ); - let loads = scheduler.get_potential_loads(Some(vec![1, 2, 3, 4]), 64, HashMap::new(), true); + let loads = + scheduler.get_potential_loads(Some(vec![1, 2, 3, 4]), 64, Default::default(), true); let worker_load = loads .iter() .find(|load| load.worker_id == response.best_worker.worker_id && load.dp_rank == 0) @@ -1064,7 +1065,7 @@ mod tests { .await .unwrap(); - let loads = scheduler.get_potential_loads(None, 0, HashMap::new(), false); + let loads = scheduler.get_potential_loads(None, 0, Default::default(), false); assert_eq!(loads[0].active_requests, 0); cancel_token.cancel(); } @@ -1210,8 +1211,8 @@ mod tests { 64, None, TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1249,8 +1250,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), Some(&crate::config::RouterConfigOverride { track_prefill_tokens: Some(false), ..Default::default() @@ -1298,8 +1299,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::from([(worker, 0.75)]), - HashMap::from([(worker, 48)]), + rustc_hash::FxHashMap::from_iter([(worker, 0.75)]), + rustc_hash::FxHashMap::from_iter([(worker, 48)]), None, true, None, @@ -1345,8 +1346,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1370,8 +1371,8 @@ mod tests { 64, Some(vec![5, 6, 7, 8]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1416,8 +1417,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1441,8 +1442,8 @@ mod tests { 64, Some(vec![5, 6, 7, 8]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1501,8 +1502,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1526,8 +1527,8 @@ mod tests { 64, Some(vec![5, 6, 7, 8]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1585,8 +1586,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1610,8 +1611,8 @@ mod tests { 64, Some(vec![5, 6, 7, 8]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1667,8 +1668,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, Some("adapter-a".to_string()), @@ -1728,7 +1729,8 @@ mod tests { .collect(); expected.sort_by_key(|load| (load.worker_id, load.dp_rank)); - let mut actual = scheduler.get_potential_loads(Some(token_seq), 128, HashMap::new(), true); + let mut actual = + scheduler.get_potential_loads(Some(token_seq), 128, Default::default(), true); actual.sort_by_key(|load| (load.worker_id, load.dp_rank)); assert_eq!(actual.len(), expected.len()); @@ -1771,8 +1773,8 @@ mod tests { 100, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -1789,7 +1791,7 @@ mod tests { tokio::time::advance(Duration::from_secs(6)).await; - let loads = scheduler.get_potential_loads(None, 0, HashMap::new(), true); + let loads = scheduler.get_potential_loads(None, 0, Default::default(), true); assert_eq!(loads.len(), 1); assert_eq!(loads[0].potential_prefill_tokens, 40); @@ -1804,7 +1806,7 @@ mod tests { assert_eq!( scheduler - .get_potential_loads(None, 64, HashMap::new(), true,) + .get_potential_loads(None, 64, Default::default(), true,) .len(), 1 ); @@ -1823,7 +1825,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), async { loop { if scheduler - .get_potential_loads(None, 64, HashMap::new(), true) + .get_potential_loads(None, 64, Default::default(), true) .len() == 3 { @@ -1882,7 +1884,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), async { loop { if scheduler - .get_potential_loads(None, 64, HashMap::new(), true) + .get_potential_loads(None, 64, Default::default(), true) .iter() .any(|load| load.worker_id == 1) { @@ -1937,7 +1939,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), async { loop { if scheduler - .get_potential_loads(None, 64, HashMap::new(), true) + .get_potential_loads(None, 64, Default::default(), true) .is_empty() { break; @@ -1959,7 +1961,7 @@ mod tests { assert_eq!( scheduler - .get_potential_loads(None, 64, HashMap::new(), true) + .get_potential_loads(None, 64, Default::default(), true) .len(), 1 ); @@ -1971,7 +1973,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(100)).await; - let loads = scheduler.get_potential_loads(None, 64, HashMap::new(), true); + let loads = scheduler.get_potential_loads(None, 64, Default::default(), true); assert_eq!(loads.len(), 1); assert_eq!(loads[0].worker_id, 0); @@ -1996,8 +1998,8 @@ mod tests { 64, Some(vec![11, 22]), TierOverlapBlocks::default(), - HashMap::new(), - HashMap::new(), + Default::default(), + Default::default(), None, true, None, @@ -2012,7 +2014,7 @@ mod tests { .await .unwrap(); - let loads = scheduler.get_potential_loads(None, 64, HashMap::new(), false); + let loads = scheduler.get_potential_loads(None, 64, Default::default(), false); assert_eq!(loads.len(), 1); assert_eq!(loads[0].potential_prefill_tokens, 64); diff --git a/lib/kv-router/src/scheduling/overlap.rs b/lib/kv-router/src/scheduling/overlap.rs index 217980b09aa8..2af77f2de30b 100644 --- a/lib/kv-router/src/scheduling/overlap.rs +++ b/lib/kv-router/src/scheduling/overlap.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use crate::config::{KvRouterConfig, RouterConfigOverride}; use crate::indexer::TieredMatchDetails; @@ -23,8 +23,8 @@ pub struct CacheHitEstimates { #[derive(Debug, Clone, Default)] pub struct OverlapSignals { pub tier_overlap_blocks: TierOverlapBlocks, - pub effective_overlap_blocks: HashMap, - pub effective_cached_tokens: HashMap, + pub effective_overlap_blocks: FxHashMap, + pub effective_cached_tokens: FxHashMap, } impl OverlapSignals { @@ -141,8 +141,8 @@ impl<'a> OverlapAnalysis<'a> { cache_hit_estimates_from_tiered_matches(self.config, self.block_size, self.tiered); OverlapSignals { tier_overlap_blocks: tier_overlap_blocks_from_tiered_matches(self.tiered), - effective_overlap_blocks: estimates.effective_overlap_blocks.into_iter().collect(), - effective_cached_tokens: estimates.cached_tokens.into_iter().collect(), + effective_overlap_blocks: estimates.effective_overlap_blocks, + effective_cached_tokens: estimates.cached_tokens, } } diff --git a/lib/kv-router/src/scheduling/overlap_refresh.rs b/lib/kv-router/src/scheduling/overlap_refresh.rs index cef255567898..0b6db53d6c01 100644 --- a/lib/kv-router/src/scheduling/overlap_refresh.rs +++ b/lib/kv-router/src/scheduling/overlap_refresh.rs @@ -277,8 +277,8 @@ mod tests { self.calls.fetch_add(1, Ordering::Relaxed); Some(RefreshedOverlap::from_overlap(OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::new(), - effective_cached_tokens: HashMap::new(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), })) } } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 4b017f05a6d9..73e30a4fc5cb 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -4421,11 +4421,11 @@ policy_classes: }), overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::from([ + effective_overlap_blocks: rustc_hash::FxHashMap::from_iter([ (WorkerWithDpRank::new(0, 0), 1.0), (WorkerWithDpRank::new(1, 0), 9.0), ]), - effective_cached_tokens: HashMap::from([ + effective_cached_tokens: rustc_hash::FxHashMap::from_iter([ (WorkerWithDpRank::new(0, 0), 16), (WorkerWithDpRank::new(1, 0), 144), ]), @@ -4519,8 +4519,8 @@ policy_classes: }), overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::from([(worker, 5.0)]), - effective_cached_tokens: HashMap::from([(worker, 80)]), + effective_overlap_blocks: rustc_hash::FxHashMap::from_iter([(worker, 5.0)]), + effective_cached_tokens: rustc_hash::FxHashMap::from_iter([(worker, 80)]), }, }, }); @@ -4563,8 +4563,8 @@ policy_classes: let refresher = Arc::new(BlockingRefresher::new(RefreshedOverlap::from_overlap( OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::from([(worker, 7.0)]), - effective_cached_tokens: HashMap::from([(worker, 56)]), + effective_overlap_blocks: rustc_hash::FxHashMap::from_iter([(worker, 7.0)]), + effective_cached_tokens: rustc_hash::FxHashMap::from_iter([(worker, 56)]), }, ))); let (queue, slots) = make_queue_with_blocking_refresher( diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 8938febe80cf..866d398ff5cf 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -23,7 +23,7 @@ use policy::{ CustomWorkerSelectionState, WorkerSelectionPolicyStateRef, collect_custom_candidates, }; #[cfg(any(test, feature = "bench"))] -use reference::{pick_default_worker, selection_weights}; +use reference::pick_default_worker; use super::config::KvRouterConfig; use super::filter::{RoutingEligibility, WorkerEligibilityError}; @@ -129,37 +129,20 @@ impl<'a, C: WorkerConfigLike> WorkerSelectionInput<'a, C> { } } -#[derive(Debug, Clone, Copy)] -#[cfg(any(test, feature = "bench"))] -pub(crate) struct LogitWeights { - overlap_score_credit: f64, - overlap_score_credit_decay: f64, - prefill_load_scale: f64, - shared_cache_multiplier: f64, -} - struct MaterializedSelectionInput<'a> { request: &'a SchedulingRequest, context: WorkerSelectionContext<'a>, } impl<'a> MaterializedSelectionInput<'a> { - fn new( - request: &'a SchedulingRequest, - block_size: u32, - #[cfg(any(test, feature = "bench"))] weights: LogitWeights, - ) -> Self { + fn new(request: &'a SchedulingRequest, block_size: u32) -> Self { Self { request, context: WorkerSelectionContext { request, - #[cfg(any(test, feature = "bench"))] - request_id: request.mode.request_id().unwrap_or("-"), request_blocks: request.request_blocks(block_size), block_size, track_prefill_tokens: request.track_prefill_tokens, - #[cfg(any(test, feature = "bench"))] - weights, pinned_worker: request.pinned_worker, has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() @@ -172,6 +155,7 @@ impl<'a> MaterializedSelectionInput<'a> { } } + #[inline(always)] fn row( &self, worker: WorkerWithDpRank, @@ -182,18 +166,16 @@ impl<'a> MaterializedSelectionInput<'a> { worker, preferred_taint_multiplier, inputs, - #[cfg(any(test, feature = "bench"))] - false, |_, device_overlap_blocks| device_overlap_blocks, ) } + #[inline(always)] fn row_with_device_overlap( &self, worker: WorkerWithDpRank, preferred_taint_multiplier: Option, inputs: WorkerInputs, - #[cfg(any(test, feature = "bench"))] reference_formula: bool, select_device_overlap: impl FnOnce(f64, f64) -> f64, ) -> WorkerCandidate { let cached_tokens = if inputs.contains(WorkerInputs::CACHE) @@ -253,29 +235,10 @@ impl<'a> MaterializedSelectionInput<'a> { WorkerCacheInput::default() }; let load = if inputs.contains(WorkerInputs::LOAD) { - #[cfg(any(test, feature = "bench"))] - let raw_prefill_tokens = if reference_formula && self.request.track_prefill_tokens { - match worker_load { - Some(load) => { - // Preserve the legacy operation order when overlap exceeds the prompt. - let uncached_tokens = super::prefill_load::effective_prefill_tokens( - self.request.isl_tokens, - cached_tokens, - ); - let projected_tokens = load.active_prefill_tokens + uncached_tokens; - projected_tokens.saturating_add(cached_tokens) - } - None => self.request.isl_tokens, - } - } else { - 0 - } as f64; let available = worker_load.is_some(); let worker_load = worker_load.unwrap_or_default(); WorkerLoadInput { available, - #[cfg(any(test, feature = "bench"))] - raw_prefill_blocks: raw_prefill_tokens / self.context.block_size as f64, active_prefill_tokens: worker_load.active_prefill_tokens, decode_cost_blocks: worker_load.potential_decode_blocks() as f64, active_requests: worker_load.active_requests, @@ -402,16 +365,9 @@ fn select_worker_with_policy( } } - #[cfg(any(test, feature = "bench"))] - let weights = selection_weights(kv_router_config, request); #[cfg(not(any(test, feature = "bench")))] let _ = kv_router_config; - let mut input = MaterializedSelectionInput::new( - request, - block_size, - #[cfg(any(test, feature = "bench"))] - weights, - ); + let mut input = MaterializedSelectionInput::new(request, block_size); input.context.pinned_worker = eligibility.pinned_worker(); let selected = match state { #[cfg(any(test, feature = "bench"))] @@ -500,7 +456,7 @@ fn select_worker_with_policy( #[cfg(test)] mod test_support { - use std::collections::{HashMap, HashSet}; + use std::collections::HashSet; use rustc_hash::FxHashMap; @@ -543,8 +499,8 @@ mod test_support { isl_tokens, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index 81a415c86ac0..ea1d1d8f7aa9 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -10,8 +10,8 @@ use parking_lot::Mutex; use super::policy::WorkerSelectionPolicyStateRef; use super::{ - LogitWeights, MaterializedSelectionInput, WorkerCandidate, WorkerInputs, - WorkerSelectionContext, WorkerSelectionInput, WorkerSelector, select_worker_with_policy, + MaterializedSelectionInput, WorkerCandidate, WorkerInputs, WorkerSelectionContext, + WorkerSelectionInput, WorkerSelector, select_worker_with_policy, }; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank}; use crate::scheduling::config::KvRouterConfig; @@ -88,6 +88,14 @@ fn softmax_sample_index( entries.len() - 1 } +#[derive(Debug, Clone, Copy)] +struct LogitWeights { + overlap_score_credit: f64, + overlap_score_credit_decay: f64, + prefill_load_scale: f64, + shared_cache_multiplier: f64, +} + /// Default implementation matching the Python _cost_function. pub struct DefaultWorkerSelector { pub kv_router_config: KvRouterConfig, @@ -102,6 +110,7 @@ pub(super) struct DefaultWorkerScorer { #[derive(Debug, Clone, Copy)] struct DefaultScoringContext { + weights: LogitWeights, min_active_prefill_tokens: usize, has_tier_overlap_blocks: bool, } @@ -192,7 +201,7 @@ impl DefaultWorkerScorer { } } -pub(super) fn selection_weights( +fn selection_weights( kv_router_config: &KvRouterConfig, request: &SchedulingRequest, ) -> LogitWeights { @@ -237,6 +246,7 @@ impl DefaultScoringContext { || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() || !request.overlap.tier_overlap_blocks.disk.is_empty(); Self { + weights, min_active_prefill_tokens, has_tier_overlap_blocks, } @@ -261,7 +271,6 @@ fn default_row( worker, preferred_taint_multiplier, WorkerInputs::ALL, - true, |effective_overlap_blocks, device_overlap_blocks| { context.device_overlap(effective_overlap_blocks, device_overlap_blocks) }, @@ -277,7 +286,7 @@ impl> DefaultWorkerScorer { formula_name: &'static str, ) -> f64 { let kv_router_config = self.kv_router_config.borrow(); - let weights = context.weights; + let weights = default_context.weights; let worker = row.worker; let cache = &row.cache; let load = &row.load; @@ -329,7 +338,7 @@ impl> DefaultWorkerScorer { // `SchedulerQueueActor` task, so the logging layer cannot attach request identity to // it, and this branch returns early without reaching them. tracing::debug!( - request_id = context.request_id, + request_id = context.request.mode.request_id().unwrap_or("-"), worker_type = self.worker_type, "{formula_name} for worker_id={} dp_rank={:?} with {effective_overlap_blocks:.2} effective cached blocks: {logit:.3} \ = max(0, decode_blocks - overlap_credit_blocks) + active_request_cost_blocks \ @@ -340,7 +349,21 @@ impl> DefaultWorkerScorer { return logit; } - let adjusted_prefill_blocks = (load.raw_prefill_blocks - overlap_credit_blocks).max(0.0); + // Keep reference-only arithmetic out of public worker-input rows, including + // benchmark builds: their layout must match the production plugin's inputs. + let raw_prefill_tokens = if context.track_prefill_tokens { + if load.available { + let cached_tokens = cache.estimated_cached_tokens; + let uncached_tokens = context.request.isl_tokens.saturating_sub(cached_tokens); + (load.active_prefill_tokens + uncached_tokens).saturating_add(cached_tokens) + } else { + context.request.isl_tokens + } + } else { + 0 + }; + let raw_prefill_blocks = raw_prefill_tokens as f64 / context.block_size as f64; + let adjusted_prefill_blocks = (raw_prefill_blocks - overlap_credit_blocks).max(0.0); let prefill_cost_blocks = weights.prefill_load_scale * adjusted_prefill_blocks; let logit = prefill_cost_blocks + decode_cost_blocks + active_request_cost_blocks; @@ -352,7 +375,7 @@ impl> DefaultWorkerScorer { // inside the macro so they cost nothing when DEBUG is disabled. if shared_beyond_device_blocks > 0 { tracing::debug!( - request_id = context.request_id, + request_id = context.request.mode.request_id().unwrap_or("-"), worker_type = self.worker_type, "{formula_name} for worker_id={} dp_rank={:?} with {effective_overlap_blocks:.2} effective cached blocks, \ {} shared blocks beyond device (multiplier={shared_cache_multiplier:.2}): {logit:.3} \ @@ -363,13 +386,13 @@ impl> DefaultWorkerScorer { worker.worker_id, worker.dp_rank, shared_beyond_device_blocks, - load.raw_prefill_blocks, + raw_prefill_blocks, shared_cache_multiplier = weights.shared_cache_multiplier, prefill_load_scale = weights.prefill_load_scale ); } else { tracing::debug!( - request_id = context.request_id, + request_id = context.request.mode.request_id().unwrap_or("-"), worker_type = self.worker_type, "{formula_name} for worker_id={} dp_rank={:?} with {effective_overlap_blocks:.2} effective cached blocks: {logit:.3} \ = prefill_load_scale * adjusted_prefill_blocks + decode_blocks + active_request_cost_blocks \ @@ -378,7 +401,7 @@ impl> DefaultWorkerScorer { overlap_credit_decay: {overlap_credit_decay:.3})", worker.worker_id, worker.dp_rank, - load.raw_prefill_blocks, + raw_prefill_blocks, prefill_load_scale = weights.prefill_load_scale ); } @@ -421,8 +444,12 @@ pub(super) fn pick_default_worker( request: &SchedulingRequest, eligibility: RoutingEligibility<'_>, ) -> Option<(WorkerWithDpRank, f64)> { - let default_context = - DefaultScoringContext::new(workers, request, eligibility, input.context.weights); + let default_context = DefaultScoringContext::new( + workers, + request, + eligibility, + selection_weights(scorer.kv_router_config, request), + ); if let Some(worker) = eligibility.pinned_worker() { let row = default_row(input, default_context, worker, None); return Some(( @@ -586,7 +613,7 @@ mod tests { weights: LogitWeights, ) -> f64 { let workers = HashMap::from([(worker.worker_id, TaintedWorkerConfig::default())]); - let input = MaterializedSelectionInput::new(request, block_size, weights); + let input = MaterializedSelectionInput::new(request, block_size); let default_context = DefaultScoringContext::new(&workers, request, request.eligibility(), weights); DefaultWorkerScorer::new(selector.kv_router_config.clone(), selector.worker_type) @@ -674,8 +701,8 @@ mod tests { isl_tokens: 16, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1077,8 +1104,8 @@ mod tests { isl_tokens: 16, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1136,8 +1163,8 @@ mod tests { isl_tokens: 16, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1213,8 +1240,8 @@ mod tests { isl_tokens: 16, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1288,8 +1315,8 @@ mod tests { isl_tokens: 16, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1359,8 +1386,8 @@ mod tests { isl_tokens: 16, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1412,11 +1439,11 @@ mod tests { let isl = 4usize; let worker0 = WorkerWithDpRank::from_worker_id(0); - let mut effective_overlap_blocks = HashMap::new(); + let mut effective_overlap_blocks = rustc_hash::FxHashMap::default(); effective_overlap_blocks.insert(worker0, 2.0); // worker1 has 0 overlap (not in map) - let mut effective_cached_tokens = HashMap::new(); + let mut effective_cached_tokens = rustc_hash::FxHashMap::default(); effective_cached_tokens.insert(worker0, 2); let mut tier_overlap_blocks = crate::scheduling::TierOverlapBlocks::default(); @@ -1493,7 +1520,7 @@ mod tests { let worker0 = WorkerWithDpRank::from_worker_id(0); let worker1 = WorkerWithDpRank::from_worker_id(1); - let mut effective_cached_tokens = HashMap::new(); + let mut effective_cached_tokens = rustc_hash::FxHashMap::default(); effective_cached_tokens.insert(worker0, 32); let mut tier_overlap_blocks = crate::scheduling::TierOverlapBlocks::default(); @@ -1524,7 +1551,7 @@ mod tests { isl_tokens: isl, overlap: OverlapSignals { tier_overlap_blocks, - effective_overlap_blocks: HashMap::new(), + effective_overlap_blocks: Default::default(), effective_cached_tokens, }, kv_transfer_candidates: None, @@ -1799,7 +1826,7 @@ mod tests { let worker0 = WorkerWithDpRank::from_worker_id(0); let worker1 = WorkerWithDpRank::from_worker_id(1); - let mut effective_overlap_blocks = HashMap::new(); + let mut effective_overlap_blocks = rustc_hash::FxHashMap::default(); effective_overlap_blocks.insert(worker0, 4.0); let config = KvRouterConfig { @@ -1827,7 +1854,7 @@ mod tests { overlap: OverlapSignals { tier_overlap_blocks: Default::default(), effective_overlap_blocks, - effective_cached_tokens: HashMap::new(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, @@ -1878,7 +1905,7 @@ mod tests { prefill_load_scale: 1.0, shared_cache_multiplier: 1.0, }; - let input = MaterializedSelectionInput::new(&request, 16, weights); + let input = MaterializedSelectionInput::new(&request, 16); let default_context = DefaultScoringContext::new(&workers, &request, request.eligibility(), weights); let custom_row = input.row(worker, None, WorkerInputs::CACHE); @@ -1899,7 +1926,7 @@ mod tests { let isl = 64usize; let worker0 = WorkerWithDpRank::from_worker_id(0); - let mut effective_overlap_blocks = HashMap::new(); + let mut effective_overlap_blocks = rustc_hash::FxHashMap::default(); effective_overlap_blocks.insert(worker0, 2.0); let config = KvRouterConfig::default(); @@ -1917,7 +1944,7 @@ mod tests { overlap: OverlapSignals { tier_overlap_blocks: Default::default(), effective_overlap_blocks, - effective_cached_tokens: HashMap::new(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index 00b86e50f9d1..9400daf832a0 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -603,8 +603,8 @@ mod tests { session_context: None, overlap: OverlapSignals { tier_overlap_blocks: Default::default(), - effective_overlap_blocks: HashMap::default(), - effective_cached_tokens: HashMap::default(), + effective_overlap_blocks: Default::default(), + effective_cached_tokens: Default::default(), }, kv_transfer_candidates: None, retain_kv_transfer_chain: false, diff --git a/lib/kv-router/src/services/selection/core/queries.rs b/lib/kv-router/src/services/selection/core/queries.rs index 66d862bd08e6..46413b7782b2 100644 --- a/lib/kv-router/src/services/selection/core/queries.rs +++ b/lib/kv-router/src/services/selection/core/queries.rs @@ -25,7 +25,7 @@ impl SelectionCore { routing_group: entry.key.routing_group.clone(), loads: entry .scheduler - .get_potential_loads(None, 0, HashMap::new(), false), + .get_potential_loads(None, 0, Default::default(), false), pending_count: entry.scheduler.pending_count(), pending_isl_tokens: entry.scheduler.pending_isl_tokens(), }); diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index fe6aa772ca89..e2812919a790 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1779,7 +1779,7 @@ impl KvRouter { Ok(self.selection.scheduler().get_potential_loads( maybe_seq_hashes, isl_tokens, - cache_hit_estimates.cached_tokens.into_iter().collect(), + cache_hit_estimates.cached_tokens, track_prefill_tokens, )) } diff --git a/lib/router-plugins/builtin/Cargo.toml b/lib/router-plugins/builtin/Cargo.toml index 6fbd92348367..f40f7023d22e 100644 --- a/lib/router-plugins/builtin/Cargo.toml +++ b/lib/router-plugins/builtin/Cargo.toml @@ -27,3 +27,7 @@ tempfile = { workspace = true } [[bench]] name = "default_selection" harness = false + +[[bench]] +name = "default_lifecycle" +harness = false diff --git a/lib/router-plugins/builtin/benches/default_lifecycle.rs b/lib/router-plugins/builtin/benches/default_lifecycle.rs new file mode 100644 index 000000000000..cf7daac88783 --- /dev/null +++ b/lib/router-plugins/builtin/benches/default_lifecycle.rs @@ -0,0 +1,92 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Includes load projection, selection, booking, and release in the shared core. +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use dynamo_custom_policy_builtin::default_factory; +use dynamo_kv_router::services::selection::{ + PromptRequest, SelectAndReserveRequest, SelectionCacheConfig, SelectionCore, WorkerRequest, +}; +use dynamo_kv_router::{KvRouterConfig, WorkerSelectionPolicy, WorkerSelectionPolicyFactory}; +use std::sync::Arc; +use std::time::Duration; + +fn bench(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let _guard = runtime.enter(); + let mut group = c.benchmark_group("default_lifecycle"); + group + .warm_up_time(Duration::from_millis(200)) + .measurement_time(Duration::from_secs(1)) + .sample_size(30) + .nresamples(1_000); + for count in [8, 64, 512] { + let reference: WorkerSelectionPolicyFactory = Arc::new(|config, role, _| { + WorkerSelectionPolicy::default(config.clone(), role.default_selector_label()) + }); + for (label, factory) in [("reference", reference), ("plugin", default_factory())] { + let core = SelectionCore::try_new_local( + KvRouterConfig { + use_kv_events: false, + router_queue_threshold: None, + ..Default::default() + }, + 1, + Default::default(), + SelectionCacheConfig::default(), + factory, + ) + .unwrap(); + runtime.block_on(async { + for worker_id in 0..count { + core.upsert_worker(WorkerRequest { + worker_id, + endpoint: Some(format!("http://worker-{worker_id}:8000")), + block_size: Some(16), + max_num_batched_tokens: Some(8192), + ..Default::default() + }) + .await + .unwrap(); + } + }); + group.bench_function(BenchmarkId::new(label, count), |b| { + b.iter(|| { + runtime.block_on(async { + let response = core + .select_and_reserve(SelectAndReserveRequest { + selection_id: Some("request".into()), + prompt: PromptRequest { + token_ids: Some((0..2048).collect()), + ..Default::default() + }, + model_name: "default".into(), + routing_group: "default".into(), + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: None, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: Default::default(), + }) + .await + .unwrap(); + criterion::black_box(response); + core.free_reservation("request").await.unwrap(); + }) + }); + }); + core.shutdown(); + } + } + group.finish(); +} +criterion_group!(benches, bench); +criterion_main!(benches); diff --git a/lib/router-plugins/builtin/benches/default_selection.rs b/lib/router-plugins/builtin/benches/default_selection.rs index 74fb0bfc4b3a..0564effb25a1 100644 --- a/lib/router-plugins/builtin/benches/default_selection.rs +++ b/lib/router-plugins/builtin/benches/default_selection.rs @@ -43,16 +43,22 @@ fn allocations(mut select: impl FnMut()) -> usize { } fn bench(c: &mut Criterion) { + let decay: f64 = std::env::var("DYN_BENCH_OVERLAP_DECAY") + .unwrap_or_else(|_| "0".into()) + .parse() + .unwrap(); for temperature in [0.0, 0.7] { let mut group = c.benchmark_group(format!("default_selection/t{temperature}")); group .warm_up_time(Duration::from_millis(200)) .measurement_time(Duration::from_millis(500)) - .sample_size(20); + .sample_size(30) + .nresamples(1_000); for count in [8, 64, 256, 1024] { let (workers, request) = support::fixture(count, 2048); let config = KvRouterConfig { router_temperature: temperature, + overlap_score_credit_decay: decay, ..Default::default() }; let reference = diff --git a/lib/router-plugins/builtin/src/default.rs b/lib/router-plugins/builtin/src/default.rs index cde4406c1805..dc34def72e52 100644 --- a/lib/router-plugins/builtin/src/default.rs +++ b/lib/router-plugins/builtin/src/default.rs @@ -66,16 +66,17 @@ fn softmax_sample_index( /// Construct the builtin default from configured policy parameters. /// Per-request score overrides are not used. Request load-tracking remains host-owned. pub fn default_policy(config: KvRouterConfig, worker_label: &'static str) -> WorkerSelectionPolicy { - policy_with_rng(config, worker_label, None) + policy_with_rng(config, worker_label, None, false) } fn policy_with_rng( config: KvRouterConfig, worker_label: &'static str, rng: Option>>, + plain_decode: bool, ) -> WorkerSelectionPolicy { let picker = DefaultPicker { - plain_decode: false, + plain_decode, config: config.clone(), worker_label, rng, @@ -95,22 +96,9 @@ fn policy_for_role( config: KvRouterConfig, role: dynamo_kv_router::WorkerType, ) -> WorkerSelectionPolicy { - let picker = DefaultPicker { - plain_decode: role == dynamo_kv_router::WorkerType::Decode - && !config.conditional_disagg_enabled, - config: config.clone(), - worker_label: role.default_selector_label(), - rng: None, - entries: Vec::new(), - probabilities: Vec::new(), - }; - WorkerSelectionPolicy::new( - config, - role.default_selector_label(), - Vec::new(), - Box::new(picker), - ) - .with_exclusive_affinity(true) + let plain_decode = + role == dynamo_kv_router::WorkerType::Decode && !config.conditional_disagg_enabled; + policy_with_rng(config, role.default_selector_label(), None, plain_decode) } struct DefaultPicker { @@ -333,7 +321,12 @@ impl DefaultWorkerSelector { rng: Option>>, ) -> Self { Self { - policy: Mutex::new(policy_with_rng(config.clone(), worker_type, rng.clone())), + policy: Mutex::new(policy_with_rng( + config.clone(), + worker_type, + rng.clone(), + false, + )), kv_router_config: config, worker_type, rng, @@ -393,69 +386,30 @@ pub(super) fn register( "dynamo-default-cost-fn", Arc::new(|parameters| { let parameters: Parameters = parameters.deserialize()?; - if parameters - .overlap_score_credit - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "overlap_score_credit must be finite and non-negative", - )); - } - if parameters - .overlap_score_credit_decay - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "overlap_score_credit_decay must be finite and non-negative", - )); - } - if parameters - .prefill_load_scale - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "prefill_load_scale must be finite and non-negative", - )); - } - if parameters - .decode_active_request_weight - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "decode_active_request_weight must be finite and non-negative", - )); - } - if parameters - .host_cache_hit_weight - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "host_cache_hit_weight must be finite and non-negative", - )); - } - if parameters - .disk_cache_hit_weight - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "disk_cache_hit_weight must be finite and non-negative", - )); - } - if parameters - .shared_cache_multiplier - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "shared_cache_multiplier must be finite and non-negative", - )); - } - if parameters - .router_temperature - .is_some_and(|value| !value.is_finite() || value < 0.0) - { - return Err(WorkerSelectionPolicyProviderError::new( - "router_temperature must be finite and non-negative", - )); + for (name, value) in [ + ("overlap_score_credit", parameters.overlap_score_credit), + ( + "overlap_score_credit_decay", + parameters.overlap_score_credit_decay, + ), + ("prefill_load_scale", parameters.prefill_load_scale), + ( + "decode_active_request_weight", + parameters.decode_active_request_weight, + ), + ("host_cache_hit_weight", parameters.host_cache_hit_weight), + ("disk_cache_hit_weight", parameters.disk_cache_hit_weight), + ( + "shared_cache_multiplier", + parameters.shared_cache_multiplier, + ), + ("router_temperature", parameters.router_temperature), + ] { + if value.is_some_and(|value| !value.is_finite() || value < 0.0) { + return Err(WorkerSelectionPolicyProviderError::new(format!( + "{name} must be finite and non-negative" + ))); + } } Ok(Arc::new( move |config: &KvRouterConfig, role, _partition| { diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index c6656c09c8a8..a8c17fb71282 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -14,7 +14,7 @@ pub use default::{DefaultWorkerSelector, default_factory, default_policy}; /// Registry containing the required default only, without an optional policy catalog. pub fn default_registry() -> RouterPluginRegistry { - RouterPluginRegistry::new(default_factory()) + RouterPluginRegistry::default().with_default_factory(default_factory()) } use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegistryError}; diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 34538b1b2b14..dd6d0a170172 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -14,7 +14,7 @@ use support::*; fn seeded_selection_matches_reference_across_cache_and_load_shapes() { for temperature in [0.0, 0.7] { for prompt in [1, 17, 127, 2048] { - for mode in 0..4 { + for mode in 0..8 { let (workers, mut request) = fixture(16, prompt); let config = KvRouterConfig { router_temperature: temperature, @@ -23,7 +23,14 @@ fn seeded_selection_matches_reference_across_cache_and_load_shapes() { disk_cache_hit_weight: 0.1, ..Default::default() }; - match mode { + if mode >= 4 { + request.shared_cache_hits = + Some(dynamo_kv_router::SharedCacheHits::from_ranges(vec![ + 1..3, + 5..12, + ])); + } + match mode % 4 { 1 => request.overlap.tier_overlap_blocks = Default::default(), 2 => request.worker_loads.clear(), 3 => request.track_prefill_tokens = false, From c95eaff2863670d309df631da70d78eec66bac0a Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 18 Sep 2026 05:07:09 +0000 Subject: [PATCH 03/31] fix(router): update replay test map and KVBM lockfile Signed-off-by: Ishan Dhanani --- lib/bindings/kvbm/Cargo.lock | 12 ++++++++++++ lib/mocker/src/replay/online/router.rs | 2 +- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index eee0b73ed25a..51341927b903 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -2054,6 +2054,16 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "dynamo-custom-policy-builtin" +version = "1.6.0" +dependencies = [ + "dynamo-kv-router", + "fastrand", + "parking_lot", + "serde", +] + [[package]] name = "dynamo-data-gen" version = "1.6.0" @@ -2150,6 +2160,7 @@ dependencies = [ "derive-getters", "derive_builder", "dialoguer", + "dynamo-custom-policy-builtin", "dynamo-kv-router", "dynamo-memory", "dynamo-mocker", @@ -2248,6 +2259,7 @@ dependencies = [ "bytes", "dashmap", "derive_builder", + "dynamo-custom-policy-builtin", "dynamo-data-gen", "dynamo-kv-router", "dynamo-tokens", diff --git a/lib/mocker/src/replay/online/router.rs b/lib/mocker/src/replay/online/router.rs index f8489409d045..b03d6c39bcfe 100644 --- a/lib/mocker/src/replay/online/router.rs +++ b/lib/mocker/src/replay/online/router.rs @@ -372,7 +372,7 @@ impl KvReplayRouter { self.scheduler.get_potential_loads( None, isl_tokens, - std::collections::HashMap::new(), + Default::default(), track_prefill_tokens, ) } From 50af464c288906416b2f81106b5c32338ead9c17 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 18 Sep 2026 05:15:11 +0000 Subject: [PATCH 04/31] fix(bench): specify relay phase clock type Signed-off-by: Ishan Dhanani --- lib/bench/kv_router/dc_ckf_relay_bench.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/bench/kv_router/dc_ckf_relay_bench.rs b/lib/bench/kv_router/dc_ckf_relay_bench.rs index eb68905e8901..116995d4022d 100644 --- a/lib/bench/kv_router/dc_ckf_relay_bench.rs +++ b/lib/bench/kv_router/dc_ckf_relay_bench.rs @@ -992,7 +992,7 @@ async fn drive_phase( tokio::task::spawn_blocking(move || { let ready = Arc::new(Barrier::new(issuer_count + 1)); let start = Arc::new(Barrier::new(issuer_count + 1)); - let started = Arc::new(OnceLock::new()); + let started = Arc::new(OnceLock::::new()); let (elapsed, outputs) = std::thread::scope(|scope| { let mut handles = Vec::with_capacity(issuer_count); for (issuer, shard) in shards.into_iter().enumerate() { From a054cb14fa8c03743ba4ba0127285317e5bb17d0 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 18 Sep 2026 05:42:55 +0000 Subject: [PATCH 05/31] fix(router): distinguish builtin defaults from explicit policies Signed-off-by: Ishan Dhanani --- deploy/inference-gateway/ext-proc/src/peer_discovery.rs | 6 ++---- lib/bindings/python/rust/lib.rs | 9 +++++++++ lib/kv-router/src/plugins/registry.rs | 8 +++++++- 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 2d2270e858ec..271989cd10f3 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -268,9 +268,7 @@ mod tests { use axum::{Json, Router}; use dynamo_kv_router::WorkerType; use dynamo_kv_router::config::kv_router_config_from_dynamo_env; - use dynamo_kv_router::services::selection::{ - SelectionServiceBuilder, WorkerSelectionPolicyRegistry, - }; + use dynamo_kv_router::services::selection::SelectionServiceBuilder; use kube::Config; use tokio::net::TcpListener; use tokio::sync::Notify; @@ -285,7 +283,7 @@ mod tests { SelectionServiceBuilder::new( kv_router_config_from_dynamo_env(), WorkerType::Aggregated, - WorkerSelectionPolicyRegistry::default(), + dynamo_custom_policy_builtin::default_registry(), ) .indexer_threads(1) .replica_sync(listener_port, Vec::new()) diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index fa4cf09a3b0c..dae7cddac38c 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -430,6 +430,15 @@ pub(crate) fn router_plugins(config: &KvRouterConfig) -> anyhow::Result RouterPluginRegistry { #[cfg(feature = "custom-policy")] diff --git a/lib/kv-router/src/plugins/registry.rs b/lib/kv-router/src/plugins/registry.rs index ba0e811364a4..90554fa0fdfa 100644 --- a/lib/kv-router/src/plugins/registry.rs +++ b/lib/kv-router/src/plugins/registry.rs @@ -116,7 +116,13 @@ impl RouterPluginRegistry { config: &KvRouterConfig, ) -> Result { Ok(super::RouterPlugins { - worker_selection: self.resolve(config)?, + // The builtin default does not opt a frontend into custom-plugin restrictions. + // Embedded routers still obtain it from the registry when constructing a pool. + worker_selection: if config.selected_worker_selection_policy_instance()?.is_some() { + self.resolve(config)? + } else { + None + }, request_classifier: self.resolve_request_classifier(config)?, }) } From 01badd39aa3f09c67f35bc22c252cc82ddedbbd7 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 18 Sep 2026 19:56:44 +0000 Subject: [PATCH 06/31] refactor(router): split builtin policy scoring and picking Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 46 +- .../router/custom-policy-example/README.md | 4 +- .../simple-filter-score-pick/src/lib.rs | 2 +- .../simple-filter-score-pick/src/scorer.rs | 27 +- lib/kv-router/src/plugins/worker_selection.rs | 25 +- lib/kv-router/src/scheduling/selector/mod.rs | 11 +- .../src/scheduling/selector/policy.rs | 167 ++++--- lib/router-plugins/builtin/src/default.rs | 446 ------------------ lib/router-plugins/builtin/src/default/mod.rs | 53 +++ .../builtin/src/default/parameters.rs | 89 ++++ .../builtin/src/default/picker.rs | 164 +++++++ .../builtin/src/default/scorer.rs | 124 +++++ .../builtin/src/default/selector.rs | 88 ++++ lib/router-plugins/builtin/src/lib.rs | 2 +- .../builtin/tests/default_lifecycle.rs | 76 +++ .../builtin/tests/default_policy.rs | 32 -- .../builtin/tests/scorer_prepare.rs | 265 +++++++++++ 17 files changed, 1035 insertions(+), 586 deletions(-) delete mode 100644 lib/router-plugins/builtin/src/default.rs create mode 100644 lib/router-plugins/builtin/src/default/mod.rs create mode 100644 lib/router-plugins/builtin/src/default/parameters.rs create mode 100644 lib/router-plugins/builtin/src/default/picker.rs create mode 100644 lib/router-plugins/builtin/src/default/scorer.rs create mode 100644 lib/router-plugins/builtin/src/default/selector.rs create mode 100644 lib/router-plugins/builtin/tests/scorer_prepare.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index c2aec97b24e0..9bd5f486948e 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -11,13 +11,14 @@ subtitle: Build Rust filters, scorers, and pickers for the Dynamo frontend or EP ## How It Works -This feature replaces the worker-ranking part of Dynamo's routing pipeline. A `WorkerFilter` can exclude a host-eligible worker, a `WorkerScorer` assigns a finite cost to each remaining worker, Dynamo adds costs from all configured scorers, and a `WorkerPicker` chooses one row from the scored candidates. Filters run in declaration order before scoring, and rejecting every worker returns an error. Lower costs rank first by convention, but the picker can implement deterministic selection, sampling, tie-breaking, or policy-local state. Dynamo continues to own discovery, eligibility, score validation, accounting, and reservation. +This feature replaces the worker-ranking part of Dynamo's routing pipeline. A `WorkerFilter` can exclude a host-eligible worker, a `WorkerScorer` assigns a finite cost to each remaining worker, Dynamo adds costs from all configured scorers, and a `WorkerPicker` chooses one row from the scored candidates. Filters run in declaration order. Each scorer can then prepare an aggregate from all surviving workers before scoring each worker. Rejecting every worker returns an error. Lower costs rank first by convention, but the picker can implement deterministic selection, sampling, tie-breaking, or policy-local state. Dynamo continues to own discovery, eligibility, score validation, accounting, and reservation. ```mermaid flowchart LR Request["Incoming request"] --> Eligibility["Dynamo eligibility filters"] Eligibility --> Filters["WorkerFilter(s)
keep or reject each worker"] - Filters --> Scorers["WorkerScorer(s)
one cost per worker"] + Filters --> Prepare["WorkerScorer.prepare
all surviving workers"] + Prepare --> Scorers["WorkerScorer.score
one cost per worker"] Scorers --> Sum["Dynamo sums scorer costs"] Sum --> Picker["WorkerPicker
one candidate row"] Picker --> Validate["Dynamo validates the row"] @@ -37,11 +38,11 @@ Write your own policy when no shipped policy expresses the rule you need. The re | Stage | Input | Output | Use this stage for | |---|---|---|---| | `WorkerFilter` | Request context and one host-eligible worker | Keep or reject | The rule is a hard requirement, not a ranking preference | -| `WorkerScorer` | Request context and one eligible worker | One finite cost | The rule ranks workers or adds a penalty | +| `WorkerScorer` | Request context; all surviving workers in `prepare`, one worker in `score` | One finite cost per worker | The rule ranks workers or adds a penalty | | `WorkerPicker` | All eligible workers and their total costs | One row index | The rule samples, breaks ties, tracks policy state, or ignores total cost | | Policy factory | Router configuration, worker type, and routing partition | One `WorkerSelectionPolicy` | Prefill, decode, models, or routing groups need different components | -An external policy owns its filters, scorers, and picker. Dynamo's default scorer and picker are internal and can change with the built-in routing algorithm. +An external policy owns its filters, scorers, and picker. The default scorer and picker are private components of `dynamo-custom-policy-builtin`. They implement the same public traits as external policies. ## Public Plugin API and Compatibility @@ -143,7 +144,7 @@ Pass filters to `WorkerSelectionPolicy::new_with_filters`. If the policy has no A scorer expresses a preference without excluding a worker. It returns one finite cost for one worker. Lower total cost is better by convention. -This scorer uses the current number of active requests as its cost: +This scorer first finds the lowest active-request count among surviving workers, then scores each worker by the number of requests above that minimum. Subtracting the same minimum from each cost preserves the ranking: ```rust use dynamo_kv_router::plugins::worker_selection::{ @@ -151,13 +152,34 @@ use dynamo_kv_router::plugins::worker_selection::{ WorkerSelectionPolicyError, }; -struct ActiveRequestsScorer; +#[derive(Default)] +struct ActiveRequestsScorer { + minimum: usize, +} impl WorkerScorer for ActiveRequestsScorer { + /// Requests load inputs for the active-request count. fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::LOAD } + /// Prepare a request-local minimum after all filters have run. + fn prepare( + &mut self, + _context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + self.minimum = usize::MAX; + for candidate in candidates { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + self.minimum = self.minimum.min(load.active_requests()); + } + Ok(()) + } + + /// Return a relative load cost. Subtracting the minimum preserves the ranking. fn score( &mut self, _context: &WorkerSelectionContext<'_>, @@ -166,11 +188,13 @@ impl WorkerScorer for ActiveRequestsScorer { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - Ok(load.active_requests() as f64) + Ok((load.active_requests() - self.minimum) as f64) } } ``` +`prepare` is optional and defaults to no work. Dynamo calls it once per scorer after filtering, before any `score` call. It receives a borrowed slice with the same declared inputs and snapshot used for scoring. Recompute request-local aggregates on each call. The slice order is unspecified, and the slice cannot be retained. Dynamo skips preparation when no candidates survive; a preparation error stops selection before picking or booking. + A policy can stack multiple scorers. Dynamo calls them in declaration order and adds their costs. Dynamo rejects a non-finite contribution or total. ### Pick a Worker @@ -241,7 +265,7 @@ fn provider( let filters: Vec> = vec![ Box::new(MinimumDeviceOverlapFilter { minimum_blocks }), ]; - let scorers: Vec> = vec![Box::new(ActiveRequestsScorer)]; + let scorers: Vec> = vec![Box::new(ActiveRequestsScorer::default())]; WorkerSelectionPolicy::new_with_filters( config.clone(), worker_type.as_str(), @@ -593,7 +617,7 @@ Holding a classification future pauses admission before dispatch; it does not in - Return a valid picker row. - Treat candidate order as unspecified. - Request only the signal groups that the component reads. -- Keep blocking I/O and panics out of `keep`, `score`, and `pick`. +- Keep blocking I/O and panics out of `keep`, `prepare`, `score`, and `pick`. - Keep policy state local to the factory-created policy unless cross-partition sharing is a deliberate requirement. - Build the policy against the same Dynamo revision as the frontend or EPP. - If a signal adds work, storage, allocation, or another scan, run the worker-selection benchmark. @@ -622,6 +646,8 @@ The provider also accepts `overlap_score_credit_decay`, `decode_active_request_w `WorkerCacheInput::accounting_cache_estimate()` returns the host's weighted block estimate and rounded token estimate. Missing estimates are zero. These estimates are not clamped to prompt length and do not measure physical GPU residency. The plugin can use them without changing the load that the host records. `WorkerLoadInput::is_available()` distinguishes a missing load observation from an observed idle worker. Both values describe the current selection snapshot. -A picker can implement `pick_with_cost()` to return a row and an optional finite cost for host diagnostics. Existing `pick()` implementations remain valid. The host checks the row and cost before booking. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. +The default scorer owns the cache and load formula, including preferred-taint costs. Its `prepare` method finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. Direct Rust callers and replay import `DefaultWorkerSelector` from `dynamo_custom_policy_builtin`. Scheduler actors use `default_policy()` directly. The core's old selector is available only as a test and benchmark reference, not in production builds. + +The builtin crate keeps composition in `default/mod.rs`, the cost formula in `default/scorer.rs`, selection in `default/picker.rs`, startup parameters in `default/parameters.rs`, and the synchronous adapter in `default/selector.rs`. The two-tier policy remains in its own module. Preparation reads the existing snapshot; it adds no tracker synchronization. The routing core still projects load, selects and books serially, and processes request lifecycle updates. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index ac03bfb36969..decef548ce5c 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -31,7 +31,7 @@ Preferred routing taints are optional candidate metadata. A filter, scorer, or p | `disagg-filter-score-pick` | Prefill and decode workers each need the complete policy flow | | `simple-stacked-score-pick` | Multiple scorer costs compose before one picker runs | -The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap and scores active requests. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. +The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap, computes the least-loaded surviving worker in `WorkerScorer::prepare`, and scores active requests above that minimum. Preparation uses the same request snapshot as scoring and resets the minimum for every selection. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. The [`soft-pin-repin` policy](soft-pin-repin/README.md) documents its load threshold, soft-binding behavior, and two-Mocker `A -> B -> B` walkthrough. @@ -110,7 +110,7 @@ fn provider( config.clone(), worker_type.as_str(), filters, - vec![Box::new(ActiveRequestsScorer)], + vec![Box::new(ActiveRequestsScorer::default())], Box::new(RequestAwarePicker), ) })) diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs index a63e0e6504d0..f9439a8de1ff 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs @@ -56,7 +56,7 @@ fn provider( config.clone(), worker_type.as_str(), filters, - vec![Box::new(ActiveRequestsScorer)], + vec![Box::new(ActiveRequestsScorer::default())], Box::new(RequestAwarePicker), ) }, diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs index a9429e55de94..50840db7a340 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs @@ -7,8 +7,11 @@ use dynamo_kv_router::plugins::worker_selection::{ WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, }; -/// Scores each candidate by its current number of active requests. -pub(crate) struct ActiveRequestsScorer; +/// Scores active requests above the least-loaded surviving candidate. +#[derive(Default)] +pub(crate) struct ActiveRequestsScorer { + minimum: usize, +} impl WorkerScorer for ActiveRequestsScorer { /// Requests load inputs for the active-request count. @@ -16,7 +19,23 @@ impl WorkerScorer for ActiveRequestsScorer { WorkerInputs::LOAD } - /// Returns the active-request count as a lower-is-better cost. + /// Prepare a request-local minimum after all filters have run. + fn prepare( + &mut self, + _context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + self.minimum = usize::MAX; + for candidate in candidates { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + self.minimum = self.minimum.min(load.active_requests()); + } + Ok(()) + } + + /// Return a relative load cost. Subtracting the minimum preserves the ranking. fn score( &mut self, _context: &WorkerSelectionContext<'_>, @@ -25,6 +44,6 @@ impl WorkerScorer for ActiveRequestsScorer { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - Ok(load.active_requests() as f64) + Ok((load.active_requests() - self.minimum) as f64) } } diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index ce7b347b5cd9..3955b936d4a1 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -128,6 +128,21 @@ pub trait WorkerScorer: Send { WorkerInputs::NONE } + /// Prepare once for this selection using every candidate that survived host eligibility + /// and policy filters. All scorers prepare before any candidate is scored. This is not + /// called for an empty candidate set; returning an error stops selection before picking. + /// + /// The slice borrows the same request snapshot and declared worker inputs used by `score`. + /// Its order is unspecified, and it cannot be retained after this call. Reset request-local + /// aggregates here rather than carrying them between selections. The default does no work. + fn prepare( + &mut self, + _context: &WorkerSelectionContext<'_>, + _candidates: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + Ok(()) + } + /// Return one finite, lower-is-better cost contribution for an eligible worker row. fn score( &mut self, @@ -155,16 +170,6 @@ pub trait WorkerFilter: Send { /// Selects one row after all filters and scorers run. pub trait WorkerPicker: Send { - /// Return a selected row and an optional policy-computed cost for host diagnostics. - /// The host validates both. Existing pickers use the summed scorer cost by default. - fn pick_with_cost( - &mut self, - context: &WorkerSelectionContext<'_>, - input: WorkerInputView<'_>, - ) -> Result<(usize, Option), WorkerSelectionPolicyError> { - self.pick(context, input).map(|row| (row, None)) - } - /// Declare the optional worker-signal columns needed by this picker. fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::NONE diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 866d398ff5cf..3c40b0e4f1db 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -413,7 +413,7 @@ fn select_worker_with_policy( .contains(WorkerInputs::LOAD) .then_some(load_inputs.as_slice()), }; - let (row, selected_cost) = picker.pick_with_cost(&input.context, picker_input)?; + let row = picker.pick(&input.context, picker_input)?; let Some(candidate) = candidates.get(row) else { return Err(WorkerSelectionPolicyError::InvalidPickerRow { row, @@ -421,14 +421,7 @@ fn select_worker_with_policy( } .into()); }; - let cost = selected_cost.unwrap_or(candidate.cost); - if !cost.is_finite() { - return Err(WorkerSelectionPolicyError::failed( - "picker returned non-finite cost", - ) - .into()); - } - Some((candidate.worker, cost)) + Some((candidate.worker, candidate.cost)) } } }; diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index e5c974b973db..55b41076bbd8 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -45,6 +45,7 @@ pub(super) struct CustomWorkerSelectionState { pub(super) filter_inputs: WorkerInputs, pub(super) scorer_picker_inputs: WorkerInputs, pub(super) picker_inputs: WorkerInputs, + unscored_candidates: Vec, pub(super) candidates: Vec, pub(super) cache_inputs: Vec, pub(super) load_inputs: Vec, @@ -52,8 +53,8 @@ pub(super) struct CustomWorkerSelectionState { /// Native scorer/picker composition for [`WorkerSelector`]. /// -/// SelectionService constructs the concrete default state unless a caller explicitly supplies -/// custom scorer and picker implementations through [`Self::new`]. +/// Routing hosts supply a policy factory. Both builtin and external policies compose +/// scorers and a picker through [`Self::new`], with optional filters through [`Self::new_with_filters`]. pub struct WorkerSelectionPolicy { kv_router_config: KvRouterConfig, worker_label: &'static str, @@ -104,6 +105,7 @@ impl WorkerSelectionPolicy { filter_inputs, scorer_picker_inputs, picker_inputs, + unscored_candidates: Vec::new(), candidates: Vec::new(), cache_inputs: Vec::new(), load_inputs: Vec::new(), @@ -135,28 +137,14 @@ impl WorkerSelectionPolicy { } #[inline(always)] -#[allow(clippy::too_many_arguments)] -fn push_scored_candidate( - context: &WorkerSelectionContext<'_>, +fn push_picker_candidate( candidate: &WorkerCandidate, - scorers: &mut [Box], + cost: f64, picker_inputs: WorkerInputs, candidates: &mut Vec, cache_inputs: &mut Vec, load_inputs: &mut Vec, -) -> Result<(), KvSchedulerError> { - let mut cost = 0.0; - for (scorer_index, scorer) in scorers.iter_mut().enumerate() { - let contribution = scorer.score(context, candidate)?; - cost += contribution; - if !contribution.is_finite() || !cost.is_finite() { - return Err(WorkerSelectionPolicyError::NonFiniteCost { - scorer_index, - row: candidates.len(), - } - .into()); - } - } +) { candidates.push(ScoredWorkerCandidate { worker: candidate.worker, cost, @@ -168,7 +156,68 @@ fn push_scored_candidate( if picker_inputs.contains(WorkerInputs::LOAD) { load_inputs.push(candidate.load); } - Ok(()) +} + +impl CustomWorkerSelectionState { + // Keep row construction and storage together to avoid passing a full row through a call. + #[inline(always)] + fn push_candidate(&mut self, candidate: WorkerCandidate) { + if self.scorers.is_empty() { + // Picker-only policies do not need a second copy of the worker inputs. + push_picker_candidate( + &candidate, + 0.0, + self.picker_inputs, + &mut self.candidates, + &mut self.cache_inputs, + &mut self.load_inputs, + ); + } else { + self.unscored_candidates.push(candidate); + } + } + + fn score_candidates( + &mut self, + context: &WorkerSelectionContext<'_>, + ) -> Result<(), KvSchedulerError> { + let Self { + scorers, + picker_inputs, + unscored_candidates, + candidates, + cache_inputs, + load_inputs, + .. + } = self; + if unscored_candidates.is_empty() { + return Ok(()); + } + for scorer in scorers.iter_mut() { + scorer.prepare(context, unscored_candidates)?; + } + for (row, candidate) in unscored_candidates.iter().enumerate() { + let mut cost = 0.0; + for (scorer_index, scorer) in scorers.iter_mut().enumerate() { + let contribution = scorer.score(context, candidate)?; + cost += contribution; + if !contribution.is_finite() || !cost.is_finite() { + return Err( + WorkerSelectionPolicyError::NonFiniteCost { scorer_index, row }.into(), + ); + } + } + push_picker_candidate( + candidate, + cost, + *picker_inputs, + candidates, + cache_inputs, + load_inputs, + ); + } + Ok(()) + } } pub(super) fn collect_custom_candidates( @@ -178,24 +227,15 @@ pub(super) fn collect_custom_candidates( request: &SchedulingRequest, eligibility: RoutingEligibility<'_>, ) -> Result { - let CustomWorkerSelectionState { - filters, - scorers, - filter_inputs, - scorer_picker_inputs, - picker_inputs, - candidates, - cache_inputs, - load_inputs, - .. - } = state; - candidates.clear(); - cache_inputs.clear(); - load_inputs.clear(); - if filters.is_empty() { + state.unscored_candidates.clear(); + state.candidates.clear(); + state.cache_inputs.clear(); + state.load_inputs.clear(); + if state.filters.is_empty() { let materialize_preferred_taint = eligibility.pinned_worker().is_none() - && scorer_picker_inputs.contains(WorkerInputs::PREFERRED_TAINT); - let mut error = None; + && state + .scorer_picker_inputs + .contains(WorkerInputs::PREFERRED_TAINT); eligibility.any_eligible_worker_rank(workers, |worker, config| { let preferred_taint_multiplier = if materialize_preferred_taint { request @@ -204,30 +244,21 @@ pub(super) fn collect_custom_candidates( } else { None }; - let candidate = input.row(worker, preferred_taint_multiplier, *scorer_picker_inputs); - if let Err(policy_error) = push_scored_candidate( - &input.context, - &candidate, - scorers, - *picker_inputs, - candidates, - cache_inputs, - load_inputs, - ) { - error = Some(policy_error); - return true; - } + let candidate = input.row( + worker, + preferred_taint_multiplier, + state.scorer_picker_inputs, + ); + state.push_candidate(candidate); false }); - if let Some(error) = error { - return Err(error); - } - return Ok(!candidates.is_empty()); + state.score_candidates(&input.context)?; + return Ok(!state.candidates.is_empty()); } - let additional_inputs = scorer_picker_inputs.without(*filter_inputs); + let additional_inputs = state.scorer_picker_inputs.without(state.filter_inputs); let materialize_filter_preferred_taint = eligibility.pinned_worker().is_none() - && filter_inputs.contains(WorkerInputs::PREFERRED_TAINT); + && state.filter_inputs.contains(WorkerInputs::PREFERRED_TAINT); let materialize_additional_preferred_taint = eligibility.pinned_worker().is_none() && additional_inputs.contains(WorkerInputs::PREFERRED_TAINT); let mut has_eligible_worker = false; @@ -241,8 +272,12 @@ pub(super) fn collect_custom_candidates( } else { None }; - let filter_candidate = input.row(worker, filter_preferred_taint_multiplier, *filter_inputs); - for filter in filters.iter_mut() { + let filter_candidate = input.row( + worker, + filter_preferred_taint_multiplier, + state.filter_inputs, + ); + for filter in &mut state.filters { match filter.keep(&input.context, &filter_candidate) { Ok(true) => {} Ok(false) => return false, @@ -265,24 +300,14 @@ pub(super) fn collect_custom_candidates( additional_preferred_taint_multiplier, additional_inputs, ); - let candidate = filter_candidate.with_inputs_from(&additional, *scorer_picker_inputs); - if let Err(policy_error) = push_scored_candidate( - &input.context, - &candidate, - scorers, - *picker_inputs, - candidates, - cache_inputs, - load_inputs, - ) { - error = Some(policy_error); - return true; - } + let candidate = filter_candidate.with_inputs_from(&additional, state.scorer_picker_inputs); + state.push_candidate(candidate); false }); if let Some(error) = error { return Err(error); } + state.score_candidates(&input.context)?; Ok(has_eligible_worker) } diff --git a/lib/router-plugins/builtin/src/default.rs b/lib/router-plugins/builtin/src/default.rs deleted file mode 100644 index dc34def72e52..000000000000 --- a/lib/router-plugins/builtin/src/default.rs +++ /dev/null @@ -1,446 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! The default KV policy, implemented solely against the public worker-policy API. - -use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerSelectionResult}; -use dynamo_kv_router::{ - KvRouterConfig, KvSchedulerError, WorkerInputView, WorkerInputs, WorkerPicker, - WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, - WorkerSelectionPolicyError, WorkerSelectionPolicyFactory, WorkerSelector, -}; -use parking_lot::Mutex; -use std::sync::Arc; - -fn softmax_sample_index( - entries: &[T], - cost: impl Fn(&T) -> f64, - temperature: f64, - sample: f64, - probabilities: &mut Vec, -) -> usize { - assert!(!entries.is_empty(), "Empty entries for softmax sampling"); - debug_assert_ne!(temperature, 0.0); - - let (min_cost, max_cost) = entries - .iter() - .map(&cost) - .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), cost| { - (lo.min(cost), hi.max(cost)) - }); - - probabilities.clear(); - if min_cost == max_cost { - probabilities.resize(entries.len(), 1.0 / entries.len() as f64); - } else { - let range = max_cost - min_cost; - let magnitude = if range.is_finite() { - 1.0 - } else { - min_cost.abs().max(max_cost.abs()) - }; - let min_normalized = min_cost / magnitude; - let scale = -1.0 / ((max_cost / magnitude - min_normalized) * temperature); - let max_scaled = min_normalized * scale; - probabilities.extend( - entries - .iter() - .map(|entry| (cost(entry) / magnitude * scale - max_scaled).exp()), - ); - } - - let sum: f64 = probabilities.iter().sum(); - for probability in probabilities.iter_mut() { - *probability /= sum; - } - let mut cumulative = 0.0; - for (row, probability) in probabilities.iter().enumerate() { - cumulative += probability; - if sample <= cumulative { - return row; - } - } - entries.len() - 1 -} - -/// Construct the builtin default from configured policy parameters. -/// Per-request score overrides are not used. Request load-tracking remains host-owned. -pub fn default_policy(config: KvRouterConfig, worker_label: &'static str) -> WorkerSelectionPolicy { - policy_with_rng(config, worker_label, None, false) -} - -fn policy_with_rng( - config: KvRouterConfig, - worker_label: &'static str, - rng: Option>>, - plain_decode: bool, -) -> WorkerSelectionPolicy { - let picker = DefaultPicker { - plain_decode, - config: config.clone(), - worker_label, - rng, - entries: Vec::new(), - probabilities: Vec::new(), - }; - WorkerSelectionPolicy::new(config, worker_label, Vec::new(), Box::new(picker)) - .with_exclusive_affinity(true) -} - -/// Factory installed by routing hosts, including hosts without a custom catalog. -pub fn default_factory() -> WorkerSelectionPolicyFactory { - Arc::new(|config, role, _partition| policy_for_role(config.clone(), role)) -} - -fn policy_for_role( - config: KvRouterConfig, - role: dynamo_kv_router::WorkerType, -) -> WorkerSelectionPolicy { - let plain_decode = - role == dynamo_kv_router::WorkerType::Decode && !config.conditional_disagg_enabled; - policy_with_rng(config, role.default_selector_label(), None, plain_decode) -} - -struct DefaultPicker { - plain_decode: bool, - config: KvRouterConfig, - worker_label: &'static str, - rng: Option>>, - entries: Vec<(usize, f64)>, - probabilities: Vec, -} - -impl DefaultPicker { - fn choose( - &mut self, - context: &WorkerSelectionContext<'_>, - input: WorkerInputView<'_>, - ) -> Result<(usize, Option), WorkerSelectionPolicyError> { - let cache = input - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; - let load = input - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - let candidates = input.candidates(); - if candidates.is_empty() { - return Err(WorkerSelectionPolicyError::failed("no eligible worker")); - } - let min_prefill = - if context.tracks_prefill_tokens() && self.config.overlap_score_credit_decay > 0.0 { - load.iter() - .map(|l| l.active_prefill_tokens()) - .min() - .unwrap_or(0) - } else { - 0 - }; - self.entries.clear(); - let direct_minimum = self.config.router_temperature == 0.0 && self.rng.is_none(); - let mut best_row = 0; - let mut best_cost = f64::INFINITY; - let mut ties = 0; - for (row, candidate) in candidates.iter().enumerate() { - let cache = &cache[row]; - let load = &load[row]; - let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); - let device = if context.has_tier_matches() { - cache.device_overlap_blocks() - } else { - estimated_overlap - }; - let shared = if context.has_tier_matches() { - cache.shared_beyond_device_blocks() - } else { - context - .shared_cache_hits() - .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)) - }; - let decay = if context.tracks_prefill_tokens() - && self.config.overlap_score_credit_decay > 0.0 - { - let excess = load.active_prefill_tokens().saturating_sub(min_prefill) as f64 - / context.block_size() as f64; - 1.0 / (1.0 - + self.config.overlap_score_credit_decay - * (excess / context.request_blocks() as f64)) - } else { - 1.0 - }; - // Plain disaggregated decode is load-only. Conditional decode retains cache credit. - let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { - 0.0 - } else { - self.config.overlap_score_credit - }; - let credit = overlap_credit * decay * device - + self.config.host_cache_hit_weight * cache.host_overlap_blocks() - + self.config.disk_cache_hit_weight * cache.disk_overlap_blocks() - + self.config.shared_cache_multiplier * shared as f64; - let request_cost = - self.config.decode_active_request_weight * load.active_requests() as f64; - let logit = if self.worker_label == "decode" - && !context.tracks_prefill_tokens() - && overlap_credit > 0.0 - { - (load.decode_cost_blocks() - credit).max(0.0) + request_cost - } else { - let raw_tokens = if !context.tracks_prefill_tokens() { - 0 - } else if load.is_available() { - let uncached = context.prompt_tokens().saturating_sub(cached_tokens); - (load.active_prefill_tokens() + uncached).saturating_add(cached_tokens) - } else { - context.prompt_tokens() - }; - let prefill = (raw_tokens as f64 / context.block_size() as f64 - credit).max(0.0); - self.config.prefill_load_scale * prefill + load.decode_cost_blocks() + request_cost - }; - let cost = logit * candidate.preferred_taint_multiplier().unwrap_or(1.0); - if !cost.is_finite() { - return Err(WorkerSelectionPolicyError::failed( - "default policy produced non-finite cost", - )); - } - if direct_minimum { - if cost < best_cost { - best_row = row; - best_cost = cost; - ties = 1; - } else if cost == best_cost { - ties += 1; - if fastrand::usize(0..ties) == 0 { - best_row = row; - } - } - } else { - self.entries.push((row, cost)); - } - } - if direct_minimum { - return Ok((best_row, Some(best_cost))); - } - if context.pinned_worker().is_some() { - let (row, cost) = self.entries[0]; - return Ok((row, Some(cost))); - } - // Canonical order is required only for deterministic replay, never for production ties. - if self.rng.is_some() { - self.entries.sort_unstable_by_key(|(row, _)| { - let worker = candidates[*row].worker(); - (worker.worker_id, worker.dp_rank) - }); - } - let mut rng = self.rng.as_ref().map(|rng| rng.lock()); - let selected = if self.config.router_temperature == 0.0 { - let mut best = 0; - let mut best_cost = f64::INFINITY; - let mut ties = 0; - for (index, (_, cost)) in self.entries.iter().enumerate() { - if *cost < best_cost { - best = index; - best_cost = *cost; - ties = 1; - } else if *cost == best_cost { - ties += 1; - let sample = match rng.as_mut() { - Some(rng) => rng.usize(0..ties), - None => fastrand::usize(0..ties), - }; - if sample == 0 { - best = index; - } - } - } - best - } else { - let sample = match rng.as_mut() { - Some(rng) => rng.f64(), - None => fastrand::f64(), - }; - softmax_sample_index( - &self.entries, - |(_, cost)| *cost, - self.config.router_temperature, - sample, - &mut self.probabilities, - ) - }; - let (row, cost) = self.entries[selected]; - Ok((row, Some(cost))) - } -} - -impl WorkerPicker for DefaultPicker { - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT - } - fn pick( - &mut self, - context: &WorkerSelectionContext<'_>, - input: WorkerInputView<'_>, - ) -> Result { - self.choose(context, input).map(|(row, _)| row) - } - fn pick_with_cost( - &mut self, - context: &WorkerSelectionContext<'_>, - input: WorkerInputView<'_>, - ) -> Result<(usize, Option), WorkerSelectionPolicyError> { - self.choose(context, input) - } -} - -/// Synchronous adapter for direct callers and replay. Scheduler actors use `default_policy` -/// directly and do not acquire this adapter's mutex. -pub struct DefaultWorkerSelector { - kv_router_config: KvRouterConfig, - worker_type: &'static str, - policy: Mutex, - rng: Option>>, -} -impl DefaultWorkerSelector { - pub fn new(config: Option, worker_type: &'static str) -> Self { - Self::with_rng(config.unwrap_or_default(), worker_type, None) - } - /// Construct a reproducible selector. Clones share its random stream, as before. - pub fn new_seeded( - config: Option, - worker_type: &'static str, - seed: u64, - ) -> Self { - Self::with_rng( - config.unwrap_or_default(), - worker_type, - Some(Arc::new(Mutex::new(fastrand::Rng::with_seed(seed)))), - ) - } - fn with_rng( - config: KvRouterConfig, - worker_type: &'static str, - rng: Option>>, - ) -> Self { - Self { - policy: Mutex::new(policy_with_rng( - config.clone(), - worker_type, - rng.clone(), - false, - )), - kv_router_config: config, - worker_type, - rng, - } - } -} -impl Clone for DefaultWorkerSelector { - fn clone(&self) -> Self { - Self::with_rng( - self.kv_router_config.clone(), - self.worker_type, - self.rng.clone(), - ) - } -} -impl std::fmt::Debug for DefaultWorkerSelector { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DefaultWorkerSelector") - .field("kv_router_config", &self.kv_router_config) - .field("worker_type", &self.worker_type) - .finish_non_exhaustive() - } -} -impl WorkerSelector for DefaultWorkerSelector { - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::CACHE | WorkerInputs::LOAD - } - fn uses_exclusive_affinity_target(&self) -> bool { - true - } - fn select_worker( - &self, - input: WorkerSelectionInput<'_, C>, - ) -> Result { - self.policy.lock().select_worker(input) - } -} - -#[derive(Default, serde::Deserialize)] -#[serde(default, deny_unknown_fields)] -struct Parameters { - overlap_score_credit: Option, - overlap_score_credit_decay: Option, - prefill_load_scale: Option, - decode_active_request_weight: Option, - host_cache_hit_weight: Option, - disk_cache_hit_weight: Option, - shared_cache_multiplier: Option, - router_temperature: Option, -} - -pub(super) fn register( - registry: &mut dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistry, -) -> Result<(), dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistryError> { - use dynamo_kv_router::services::selection::WorkerSelectionPolicyProviderError; - registry.register( - "dynamo-default-cost-fn", - Arc::new(|parameters| { - let parameters: Parameters = parameters.deserialize()?; - for (name, value) in [ - ("overlap_score_credit", parameters.overlap_score_credit), - ( - "overlap_score_credit_decay", - parameters.overlap_score_credit_decay, - ), - ("prefill_load_scale", parameters.prefill_load_scale), - ( - "decode_active_request_weight", - parameters.decode_active_request_weight, - ), - ("host_cache_hit_weight", parameters.host_cache_hit_weight), - ("disk_cache_hit_weight", parameters.disk_cache_hit_weight), - ( - "shared_cache_multiplier", - parameters.shared_cache_multiplier, - ), - ("router_temperature", parameters.router_temperature), - ] { - if value.is_some_and(|value| !value.is_finite() || value < 0.0) { - return Err(WorkerSelectionPolicyProviderError::new(format!( - "{name} must be finite and non-negative" - ))); - } - } - Ok(Arc::new( - move |config: &KvRouterConfig, role, _partition| { - let mut config = config.clone(); - if let Some(value) = parameters.overlap_score_credit { - config.overlap_score_credit = value; - } - if let Some(value) = parameters.overlap_score_credit_decay { - config.overlap_score_credit_decay = value; - } - if let Some(value) = parameters.prefill_load_scale { - config.prefill_load_scale = value; - } - if let Some(value) = parameters.decode_active_request_weight { - config.decode_active_request_weight = value; - } - if let Some(value) = parameters.host_cache_hit_weight { - config.host_cache_hit_weight = value; - } - if let Some(value) = parameters.disk_cache_hit_weight { - config.disk_cache_hit_weight = value; - } - if let Some(value) = parameters.shared_cache_multiplier { - config.shared_cache_multiplier = value; - } - if let Some(value) = parameters.router_temperature { - config.router_temperature = value; - } - policy_for_role(config, role) - }, - )) - }), - ) -} diff --git a/lib/router-plugins/builtin/src/default/mod.rs b/lib/router-plugins/builtin/src/default/mod.rs new file mode 100644 index 000000000000..6f268dda6f0d --- /dev/null +++ b/lib/router-plugins/builtin/src/default/mod.rs @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The default KV policy, composed from a scorer and a picker using the public plugin API. + +mod parameters; +mod picker; +mod scorer; +mod selector; + +pub(super) use parameters::register; +pub use selector::DefaultWorkerSelector; + +use dynamo_kv_router::{KvRouterConfig, WorkerSelectionPolicy, WorkerSelectionPolicyFactory}; +use parking_lot::Mutex; +use std::sync::Arc; + +/// Construct the builtin default from configured policy parameters. +/// Per-request score overrides are not used. Request load-tracking remains host-owned. +pub fn default_policy(config: KvRouterConfig, worker_label: &'static str) -> WorkerSelectionPolicy { + policy_with_rng(config, worker_label, None, false) +} + +fn policy_with_rng( + config: KvRouterConfig, + worker_label: &'static str, + rng: Option>>, + plain_decode: bool, +) -> WorkerSelectionPolicy { + let scorer = scorer::DefaultScorer::new(config.clone(), worker_label, plain_decode); + let picker = picker::DefaultPicker::new(config.router_temperature, rng); + WorkerSelectionPolicy::new( + config, + worker_label, + vec![Box::new(scorer)], + Box::new(picker), + ) + .with_exclusive_affinity(true) +} + +/// Factory installed by routing hosts, including hosts without a custom catalog. +pub fn default_factory() -> WorkerSelectionPolicyFactory { + Arc::new(|config, role, _partition| policy_for_role(config.clone(), role)) +} + +fn policy_for_role( + config: KvRouterConfig, + role: dynamo_kv_router::WorkerType, +) -> WorkerSelectionPolicy { + let plain_decode = + role == dynamo_kv_router::WorkerType::Decode && !config.conditional_disagg_enabled; + policy_with_rng(config, role.default_selector_label(), None, plain_decode) +} diff --git a/lib/router-plugins/builtin/src/default/parameters.rs b/lib/router-plugins/builtin/src/default/parameters.rs new file mode 100644 index 000000000000..6ac4fcbd8bd5 --- /dev/null +++ b/lib/router-plugins/builtin/src/default/parameters.rs @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Startup parameters and provider registration for the default policy. + +use dynamo_kv_router::KvRouterConfig; +use std::sync::Arc; + +use super::policy_for_role; + +#[derive(Default, serde::Deserialize)] +#[serde(default, deny_unknown_fields)] +struct Parameters { + overlap_score_credit: Option, + overlap_score_credit_decay: Option, + prefill_load_scale: Option, + decode_active_request_weight: Option, + host_cache_hit_weight: Option, + disk_cache_hit_weight: Option, + shared_cache_multiplier: Option, + router_temperature: Option, +} + +pub(crate) fn register( + registry: &mut dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistry, +) -> Result<(), dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistryError> { + use dynamo_kv_router::services::selection::WorkerSelectionPolicyProviderError; + registry.register( + "dynamo-default-cost-fn", + Arc::new(|parameters| { + let parameters: Parameters = parameters.deserialize()?; + for (name, value) in [ + ("overlap_score_credit", parameters.overlap_score_credit), + ( + "overlap_score_credit_decay", + parameters.overlap_score_credit_decay, + ), + ("prefill_load_scale", parameters.prefill_load_scale), + ( + "decode_active_request_weight", + parameters.decode_active_request_weight, + ), + ("host_cache_hit_weight", parameters.host_cache_hit_weight), + ("disk_cache_hit_weight", parameters.disk_cache_hit_weight), + ( + "shared_cache_multiplier", + parameters.shared_cache_multiplier, + ), + ("router_temperature", parameters.router_temperature), + ] { + if value.is_some_and(|value| !value.is_finite() || value < 0.0) { + return Err(WorkerSelectionPolicyProviderError::new(format!( + "{name} must be finite and non-negative" + ))); + } + } + Ok(Arc::new( + move |config: &KvRouterConfig, role, _partition| { + let mut config = config.clone(); + if let Some(value) = parameters.overlap_score_credit { + config.overlap_score_credit = value; + } + if let Some(value) = parameters.overlap_score_credit_decay { + config.overlap_score_credit_decay = value; + } + if let Some(value) = parameters.prefill_load_scale { + config.prefill_load_scale = value; + } + if let Some(value) = parameters.decode_active_request_weight { + config.decode_active_request_weight = value; + } + if let Some(value) = parameters.host_cache_hit_weight { + config.host_cache_hit_weight = value; + } + if let Some(value) = parameters.disk_cache_hit_weight { + config.disk_cache_hit_weight = value; + } + if let Some(value) = parameters.shared_cache_multiplier { + config.shared_cache_multiplier = value; + } + if let Some(value) = parameters.router_temperature { + config.router_temperature = value; + } + policy_for_role(config, role) + }, + )) + }), + ) +} diff --git a/lib/router-plugins/builtin/src/default/picker.rs b/lib/router-plugins/builtin/src/default/picker.rs new file mode 100644 index 000000000000..ce0e3bfa98c7 --- /dev/null +++ b/lib/router-plugins/builtin/src/default/picker.rs @@ -0,0 +1,164 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Minimum-cost selection, tie-breaking, and temperature sampling of scored candidates. + +use dynamo_kv_router::{ + WorkerInputView, WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicyError, +}; +use parking_lot::Mutex; +use std::sync::Arc; + +fn softmax_sample_index( + entries: &[T], + cost: impl Fn(&T) -> f64, + temperature: f64, + sample: f64, + probabilities: &mut Vec, +) -> usize { + assert!(!entries.is_empty(), "Empty entries for softmax sampling"); + debug_assert_ne!(temperature, 0.0); + + let (min_cost, max_cost) = entries + .iter() + .map(&cost) + .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), cost| { + (lo.min(cost), hi.max(cost)) + }); + + probabilities.clear(); + if min_cost == max_cost { + probabilities.resize(entries.len(), 1.0 / entries.len() as f64); + } else { + let range = max_cost - min_cost; + let magnitude = if range.is_finite() { + 1.0 + } else { + min_cost.abs().max(max_cost.abs()) + }; + let min_normalized = min_cost / magnitude; + let scale = -1.0 / ((max_cost / magnitude - min_normalized) * temperature); + let max_scaled = min_normalized * scale; + probabilities.extend( + entries + .iter() + .map(|entry| (cost(entry) / magnitude * scale - max_scaled).exp()), + ); + } + + let sum: f64 = probabilities.iter().sum(); + for probability in probabilities.iter_mut() { + *probability /= sum; + } + let mut cumulative = 0.0; + for (row, probability) in probabilities.iter().enumerate() { + cumulative += probability; + if sample <= cumulative { + return row; + } + } + entries.len() - 1 +} + +pub(super) struct DefaultPicker { + temperature: f64, + rng: Option>>, + entries: Vec<(usize, f64)>, + probabilities: Vec, +} + +impl DefaultPicker { + pub(super) fn new(temperature: f64, rng: Option>>) -> Self { + Self { + temperature, + rng, + entries: Vec::new(), + probabilities: Vec::new(), + } + } +} + +impl WorkerPicker for DefaultPicker { + fn pick( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + let candidates = input.candidates(); + if candidates.is_empty() { + return Err(WorkerSelectionPolicyError::failed("no eligible worker")); + } + if context.pinned_worker().is_some() { + return Ok(0); + } + if self.temperature == 0.0 && self.rng.is_none() { + let mut best_row = 0; + let mut best_cost = f64::INFINITY; + let mut ties = 0; + for (row, candidate) in candidates.iter().enumerate() { + let cost = candidate.cost(); + if cost < best_cost { + best_row = row; + best_cost = cost; + ties = 1; + } else if cost == best_cost { + ties += 1; + if fastrand::usize(0..ties) == 0 { + best_row = row; + } + } + } + return Ok(best_row); + } + self.entries.clear(); + self.entries.extend( + candidates + .iter() + .enumerate() + .map(|(row, candidate)| (row, candidate.cost())), + ); + // Canonical order is required only for deterministic replay, never for production ties. + if self.rng.is_some() { + self.entries.sort_unstable_by_key(|(row, _)| { + let worker = candidates[*row].worker(); + (worker.worker_id, worker.dp_rank) + }); + } + let mut rng = self.rng.as_ref().map(|rng| rng.lock()); + let selected = if self.temperature == 0.0 { + let mut best = 0; + let mut best_cost = f64::INFINITY; + let mut ties = 0; + for (index, (_, cost)) in self.entries.iter().enumerate() { + if *cost < best_cost { + best = index; + best_cost = *cost; + ties = 1; + } else if *cost == best_cost { + ties += 1; + let sample = match rng.as_mut() { + Some(rng) => rng.usize(0..ties), + None => fastrand::usize(0..ties), + }; + if sample == 0 { + best = index; + } + } + } + best + } else { + let sample = match rng.as_mut() { + Some(rng) => rng.f64(), + None => fastrand::f64(), + }; + softmax_sample_index( + &self.entries, + |(_, cost)| *cost, + self.temperature, + sample, + &mut self.probabilities, + ) + }; + Ok(self.entries[selected].0) + } +} diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs new file mode 100644 index 000000000000..b6a0df464a8f --- /dev/null +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Cache and load cost calculation. The host owns snapshots and validates finite scores. + +use dynamo_kv_router::{ + KvRouterConfig, WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicyError, +}; + +pub(super) struct DefaultScorer { + config: KvRouterConfig, + worker_label: &'static str, + plain_decode: bool, + min_prefill: usize, +} + +impl DefaultScorer { + pub(super) fn new( + config: KvRouterConfig, + worker_label: &'static str, + plain_decode: bool, + ) -> Self { + Self { + config, + worker_label, + plain_decode, + min_prefill: 0, + } + } +} + +impl WorkerScorer for DefaultScorer { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT + } + + fn prepare( + &mut self, + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + self.min_prefill = 0; + if context.tracks_prefill_tokens() && self.config.overlap_score_credit_decay > 0.0 { + let mut minimum = usize::MAX; + for candidate in candidates { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + minimum = minimum.min(load.active_prefill_tokens()); + } + self.min_prefill = minimum; + } + Ok(()) + } + + fn score( + &mut self, + context: &WorkerSelectionContext<'_>, + candidate: &WorkerCandidate, + ) -> Result { + let cache = candidate + .cache() + .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); + let device = if context.has_tier_matches() { + cache.device_overlap_blocks() + } else { + estimated_overlap + }; + let shared = if context.has_tier_matches() { + cache.shared_beyond_device_blocks() + } else { + context + .shared_cache_hits() + .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)) + }; + let decay = + if context.tracks_prefill_tokens() && self.config.overlap_score_credit_decay > 0.0 { + let excess = load + .active_prefill_tokens() + .saturating_sub(self.min_prefill) as f64 + / context.block_size() as f64; + 1.0 / (1.0 + + self.config.overlap_score_credit_decay + * (excess / context.request_blocks() as f64)) + } else { + 1.0 + }; + // Plain disaggregated decode is load-only. Conditional decode retains cache credit. + let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { + 0.0 + } else { + self.config.overlap_score_credit + }; + let credit = overlap_credit * decay * device + + self.config.host_cache_hit_weight * cache.host_overlap_blocks() + + self.config.disk_cache_hit_weight * cache.disk_overlap_blocks() + + self.config.shared_cache_multiplier * shared as f64; + let request_cost = self.config.decode_active_request_weight * load.active_requests() as f64; + let logit = if self.worker_label == "decode" + && !context.tracks_prefill_tokens() + && overlap_credit > 0.0 + { + (load.decode_cost_blocks() - credit).max(0.0) + request_cost + } else { + let raw_tokens = if !context.tracks_prefill_tokens() { + 0 + } else if load.is_available() { + let uncached = context.prompt_tokens().saturating_sub(cached_tokens); + (load.active_prefill_tokens() + uncached).saturating_add(cached_tokens) + } else { + context.prompt_tokens() + }; + let prefill = (raw_tokens as f64 / context.block_size() as f64 - credit).max(0.0); + self.config.prefill_load_scale * prefill + load.decode_cost_blocks() + request_cost + }; + let cost = logit * candidate.preferred_taint_multiplier().unwrap_or(1.0); + Ok(cost) + } +} diff --git a/lib/router-plugins/builtin/src/default/selector.rs b/lib/router-plugins/builtin/src/default/selector.rs new file mode 100644 index 000000000000..3742aa4fd85a --- /dev/null +++ b/lib/router-plugins/builtin/src/default/selector.rs @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Synchronous adapter for direct selection and deterministic replay. + +use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerSelectionResult}; +use dynamo_kv_router::{ + KvRouterConfig, KvSchedulerError, WorkerInputs, WorkerSelectionInput, WorkerSelectionPolicy, + WorkerSelector, +}; +use parking_lot::Mutex; +use std::sync::Arc; + +use super::policy_with_rng; + +/// Synchronous adapter for direct callers and replay. Scheduler actors use `default_policy` +/// directly and do not acquire this adapter's mutex. +pub struct DefaultWorkerSelector { + kv_router_config: KvRouterConfig, + worker_type: &'static str, + policy: Mutex, + rng: Option>>, +} +impl DefaultWorkerSelector { + pub fn new(config: Option, worker_type: &'static str) -> Self { + Self::with_rng(config.unwrap_or_default(), worker_type, None) + } + /// Construct a reproducible selector. Clones share its random stream, as before. + pub fn new_seeded( + config: Option, + worker_type: &'static str, + seed: u64, + ) -> Self { + Self::with_rng( + config.unwrap_or_default(), + worker_type, + Some(Arc::new(Mutex::new(fastrand::Rng::with_seed(seed)))), + ) + } + fn with_rng( + config: KvRouterConfig, + worker_type: &'static str, + rng: Option>>, + ) -> Self { + Self { + policy: Mutex::new(policy_with_rng( + config.clone(), + worker_type, + rng.clone(), + false, + )), + kv_router_config: config, + worker_type, + rng, + } + } +} +impl Clone for DefaultWorkerSelector { + fn clone(&self) -> Self { + Self::with_rng( + self.kv_router_config.clone(), + self.worker_type, + self.rng.clone(), + ) + } +} +impl std::fmt::Debug for DefaultWorkerSelector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DefaultWorkerSelector") + .field("kv_router_config", &self.kv_router_config) + .field("worker_type", &self.worker_type) + .finish_non_exhaustive() + } +} +impl WorkerSelector for DefaultWorkerSelector { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD + } + fn uses_exclusive_affinity_target(&self) -> bool { + true + } + fn select_worker( + &self, + input: WorkerSelectionInput<'_, C>, + ) -> Result { + self.policy.lock().select_worker(input) + } +} diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index a8c17fb71282..ced2c8edcfbd 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -5,7 +5,7 @@ //! //! Routing hosts always link the default through `default_registry`. The optional custom //! catalog adds the named default and two-tier providers through `register`. The default -//! itself uses the same public candidate inputs and picker dispatch as external policies. +//! itself uses the same public candidate inputs and scorer/picker dispatch as external policies. //! Sequence tracking, eligibility, and admission remain in dynamo-kv-router. mod default; diff --git a/lib/router-plugins/builtin/tests/default_lifecycle.rs b/lib/router-plugins/builtin/tests/default_lifecycle.rs index 9cb20387d43e..144d076c9abd 100644 --- a/lib/router-plugins/builtin/tests/default_lifecycle.rs +++ b/lib/router-plugins/builtin/tests/default_lifecycle.rs @@ -86,3 +86,79 @@ async fn shared_core_books_and_releases_with_the_builtin_default() { assert!(core.free_reservation("request").await.is_err()); core.shutdown(); } + +#[tokio::test] +async fn preparation_error_does_not_book_and_allows_retry() { + use dynamo_kv_router::{ + WorkerCandidate, WorkerInputView, WorkerPicker, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicy, WorkerSelectionPolicyError, + }; + use std::sync::Arc; + + struct FailOnce(bool); + impl WorkerScorer for FailOnce { + fn prepare( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + if std::mem::take(&mut self.0) { + return Err(WorkerSelectionPolicyError::failed("preparation failed")); + } + Ok(()) + } + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &WorkerCandidate, + ) -> Result { + Ok(0.0) + } + } + struct First; + impl WorkerPicker for First { + fn pick( + &mut self, + _: &WorkerSelectionContext<'_>, + _: WorkerInputView<'_>, + ) -> Result { + Ok(0) + } + } + let core = SelectionCore::try_new_local( + KvRouterConfig { + use_kv_events: false, + router_queue_threshold: None, + ..Default::default() + }, + 1, + Default::default(), + SelectionCacheConfig::default(), + Arc::new(|config, role, _| { + WorkerSelectionPolicy::new( + config.clone(), + role.default_selector_label(), + vec![Box::new(FailOnce(true))], + Box::new(First), + ) + }), + ) + .unwrap(); + core.upsert_worker(worker(1)).await.unwrap(); + let partition = core + .partition(&RoutingPartitionId::new("model", "default")) + .unwrap(); + let error = core + .select_and_reserve(reserve_request("request")) + .await + .unwrap_err(); + assert!(error.to_string().contains("preparation failed")); + assert!(!partition.scheduler().has_request("request")); + core.select_and_reserve(reserve_request("request")) + .await + .unwrap(); + assert!(partition.scheduler().has_request("request")); + core.free_reservation("request").await.unwrap(); + assert!(!partition.scheduler().has_request("request")); + core.shutdown(); +} diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index dd6d0a170172..5d18f3ab3f16 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -138,38 +138,6 @@ fn mandatory_pin_is_preserved() { assert_eq!(Some(selected.worker), request.pinned_worker); } -#[test] -fn non_finite_picker_cost_is_rejected() { - struct Invalid; - impl WorkerPicker for Invalid { - fn pick( - &mut self, - _: &WorkerSelectionContext<'_>, - _: WorkerInputView<'_>, - ) -> Result { - Ok(0) - } - fn pick_with_cost( - &mut self, - _: &WorkerSelectionContext<'_>, - _: WorkerInputView<'_>, - ) -> Result<(usize, Option), WorkerSelectionPolicyError> { - Ok((0, Some(f64::NAN))) - } - } - let (workers, request) = fixture(1, 17); - let error = - WorkerSelectionPolicy::new(KvRouterConfig::default(), "test", vec![], Box::new(Invalid)) - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) - .unwrap_err(); - assert!(error.to_string().contains("non-finite")); -} - #[test] fn configured_parameters_replace_request_score_overrides() { use dynamo_kv_router::RouterConfigOverride; diff --git a/lib/router-plugins/builtin/tests/scorer_prepare.rs b/lib/router-plugins/builtin/tests/scorer_prepare.rs new file mode 100644 index 000000000000..a7e760ce3816 --- /dev/null +++ b/lib/router-plugins/builtin/tests/scorer_prepare.rs @@ -0,0 +1,265 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Exercise scorer preparation from an external crate using only public plugin inputs. +mod support; + +use std::sync::{Arc, Mutex}; + +use dynamo_kv_router::{ + KvRouterConfig, WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, WorkerPicker, + WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, + WorkerSelectionPolicyError, WorkerSelector, +}; +use support::fixture; + +#[derive(Debug, PartialEq)] +enum Call { + Prepare(usize, Vec), + Score(usize), + Pick, +} + +struct RelativeLoadScorer { + index: usize, + minimum: usize, + calls: Arc>>, +} + +impl WorkerScorer for RelativeLoadScorer { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } + + fn prepare( + &mut self, + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + assert_eq!(context.prompt_tokens(), 17); + assert!(candidates.iter().all(|c| c.cache().is_none())); + let mut ids: Vec<_> = candidates.iter().map(|c| c.worker().worker_id).collect(); + ids.sort_unstable(); + self.calls + .lock() + .unwrap() + .push(Call::Prepare(self.index, ids)); + self.minimum = candidates + .iter() + .map(|c| c.load().unwrap().active_requests()) + .min() + .unwrap(); + Ok(()) + } + + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + candidate: &WorkerCandidate, + ) -> Result { + self.calls.lock().unwrap().push(Call::Score(self.index)); + Ok((candidate.load().unwrap().active_requests() - self.minimum) as f64) + } +} + +struct ExcludeWorkerOne; +impl WorkerFilter for ExcludeWorkerOne { + fn keep( + &mut self, + _: &WorkerSelectionContext<'_>, + candidate: &WorkerCandidate, + ) -> Result { + Ok(candidate.worker().worker_id != 1) + } +} + +struct InspectCosts(Arc>>); +impl WorkerPicker for InspectCosts { + fn pick( + &mut self, + _: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + assert!(input.load().is_none()); + assert!(input.cache().is_none()); + self.0.lock().unwrap().push(Call::Pick); + for candidate in input.candidates() { + // Each scorer contributes the same relative count; the host adds both. + let expected = if candidate.worker().worker_id == 3 { + 0.0 + } else { + 6.0 + }; + assert_eq!(candidate.cost(), expected); + } + Ok(input + .candidates() + .iter() + .position(|c| c.cost() == 0.0) + .unwrap()) + } +} + +#[test] +fn prepares_all_scorers_from_surviving_workers_and_resets_each_selection() { + let (workers, mut request) = fixture(4, 17); + request.allowed_worker_ids = Some([1, 2, 3].into_iter().collect()); + let calls = Arc::new(Mutex::new(Vec::new())); + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + vec![Box::new(ExcludeWorkerOne)], + (0..2) + .map(|index| { + Box::new(RelativeLoadScorer { + index, + minimum: usize::MAX, + calls: calls.clone(), + }) as Box + }) + .collect(), + Box::new(InspectCosts(calls.clone())), + ); + // Excluded workers have zero load. Neither may lower the preparation minimum. + for round in 0..3 { + for (worker, load) in &mut request.worker_loads { + load.active_requests = match worker.worker_id { + 2 => 10 + round * 10 + 3, + 3 => 10 + round * 10, + _ => 0, + }; + } + if round == 2 { + request.allowed_worker_ids = Some([3].into_iter().collect()); + } + let result = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + assert_eq!(result.worker.worker_id, 3); + let mut calls = calls.lock().unwrap(); + let expected = if round == 2 { + vec![3, 3] + } else { + vec![2, 2, 3, 3] + }; + assert_eq!(calls[0], Call::Prepare(0, expected.clone())); + assert_eq!(calls[1], Call::Prepare(1, expected.clone())); + assert_eq!(calls.len(), 2 + expected.len() * 2 + 1); + for pair in calls[2..calls.len() - 1].chunks_exact(2) { + assert_eq!(pair, [Call::Score(0), Call::Score(1)]); + } + assert_eq!(calls.last(), Some(&Call::Pick)); + calls.clear(); + } +} + +struct FailPreparation(Arc>); +impl WorkerScorer for FailPreparation { + fn prepare( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &[WorkerCandidate], + ) -> Result<(), WorkerSelectionPolicyError> { + *self.0.lock().unwrap() += 1; + Err(WorkerSelectionPolicyError::failed("prepare failed")) + } + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &WorkerCandidate, + ) -> Result { + panic!("preparation failure must stop scoring") + } +} +struct NeverPick; +impl WorkerPicker for NeverPick { + fn pick( + &mut self, + _: &WorkerSelectionContext<'_>, + _: WorkerInputView<'_>, + ) -> Result { + panic!("failed or empty selections must stop before picking") + } +} + +#[test] +fn skips_preparation_for_empty_sets_and_aborts_on_preparation_error() { + let (workers, mut request) = fixture(2, 17); + let calls = Arc::new(Mutex::new(0)); + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + vec![Box::new(ExcludeWorkerOne)], + vec![Box::new(FailPreparation(calls.clone()))], + Box::new(NeverPick), + ); + for allowed in [vec![], vec![1]] { + request.allowed_worker_ids = Some(allowed.into_iter().collect()); + assert!( + policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .is_err() + ); + assert_eq!(*calls.lock().unwrap(), 0); + } + request.allowed_worker_ids = None; + let error = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap_err(); + assert!(error.to_string().contains("prepare failed")); + assert_eq!(*calls.lock().unwrap(), 1); +} + +#[test] +fn rejects_nonfinite_contributions_and_overflow_before_picking() { + struct Constant(f64); + impl WorkerScorer for Constant { + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &WorkerCandidate, + ) -> Result { + Ok(self.0) + } + } + let (workers, request) = fixture(1, 17); + for costs in [ + vec![f64::NAN], + vec![f64::INFINITY], + vec![f64::MAX, f64::MAX], + ] { + let policy = WorkerSelectionPolicy::new( + KvRouterConfig::default(), + "test", + costs + .into_iter() + .map(|cost| Box::new(Constant(cost)) as Box) + .collect(), + Box::new(NeverPick), + ); + let error = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap_err(); + assert!(error.to_string().contains("non-finite")); + } +} From e2996bd5ba21405f7797abfb897a4d039865366f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 18 Sep 2026 23:38:12 +0000 Subject: [PATCH 07/31] perf(router): remove default score-pick hot-loop overhead Signed-off-by: Ishan Dhanani --- .../src/scheduling/selector/policy.rs | 53 +++-- lib/router-plugins/builtin/src/default/mod.rs | 11 +- .../builtin/src/default/scorer.rs | 220 ++++++++++++++---- .../builtin/tests/default_policy.rs | 96 +++++++- .../builtin/tests/scorer_prepare.rs | 71 ++++++ 5 files changed, 376 insertions(+), 75 deletions(-) diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 55b41076bbd8..e48d12d3b568 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -162,17 +162,18 @@ impl CustomWorkerSelectionState { // Keep row construction and storage together to avoid passing a full row through a call. #[inline(always)] fn push_candidate(&mut self, candidate: WorkerCandidate) { - if self.scorers.is_empty() { + // Build the picker's rows alongside the input snapshot. The scoring loop then only + // writes costs, without growing vectors or copying optional columns across trait calls. + push_picker_candidate( + &candidate, + 0.0, + self.picker_inputs, + &mut self.candidates, + &mut self.cache_inputs, + &mut self.load_inputs, + ); + if !self.scorers.is_empty() { // Picker-only policies do not need a second copy of the worker inputs. - push_picker_candidate( - &candidate, - 0.0, - self.picker_inputs, - &mut self.candidates, - &mut self.cache_inputs, - &mut self.load_inputs, - ); - } else { self.unscored_candidates.push(candidate); } } @@ -183,20 +184,35 @@ impl CustomWorkerSelectionState { ) -> Result<(), KvSchedulerError> { let Self { scorers, - picker_inputs, unscored_candidates, candidates, - cache_inputs, - load_inputs, .. } = self; if unscored_candidates.is_empty() { return Ok(()); } + debug_assert_eq!(unscored_candidates.len(), candidates.len()); for scorer in scorers.iter_mut() { scorer.prepare(context, unscored_candidates)?; } - for (row, candidate) in unscored_candidates.iter().enumerate() { + if let [scorer] = scorers.as_mut_slice() { + // One scorer needs neither a per-row scorer loop nor separate contribution/total + // checks. Keep the initial addition to preserve the sign of a zero total. + for (row, (candidate, scored)) in unscored_candidates.iter().zip(candidates).enumerate() + { + let cost = scorer.score(context, candidate)?; + if !cost.is_finite() { + return Err(WorkerSelectionPolicyError::NonFiniteCost { + scorer_index: 0, + row, + } + .into()); + } + scored.cost = 0.0 + cost; + } + return Ok(()); + } + for (row, (candidate, scored)) in unscored_candidates.iter().zip(candidates).enumerate() { let mut cost = 0.0; for (scorer_index, scorer) in scorers.iter_mut().enumerate() { let contribution = scorer.score(context, candidate)?; @@ -207,14 +223,7 @@ impl CustomWorkerSelectionState { ); } } - push_picker_candidate( - candidate, - cost, - *picker_inputs, - candidates, - cache_inputs, - load_inputs, - ); + scored.cost = cost; } Ok(()) } diff --git a/lib/router-plugins/builtin/src/default/mod.rs b/lib/router-plugins/builtin/src/default/mod.rs index 6f268dda6f0d..3d4268090f2f 100644 --- a/lib/router-plugins/builtin/src/default/mod.rs +++ b/lib/router-plugins/builtin/src/default/mod.rs @@ -27,15 +27,10 @@ fn policy_with_rng( rng: Option>>, plain_decode: bool, ) -> WorkerSelectionPolicy { - let scorer = scorer::DefaultScorer::new(config.clone(), worker_label, plain_decode); + let scorer = scorer::build(&config, worker_label, plain_decode); let picker = picker::DefaultPicker::new(config.router_temperature, rng); - WorkerSelectionPolicy::new( - config, - worker_label, - vec![Box::new(scorer)], - Box::new(picker), - ) - .with_exclusive_affinity(true) + WorkerSelectionPolicy::new(config, worker_label, vec![scorer], Box::new(picker)) + .with_exclusive_affinity(true) } /// Factory installed by routing hosts, including hosts without a custom catalog. diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index b6a0df464a8f..c640ce5951d4 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -8,29 +8,123 @@ use dynamo_kv_router::{ WorkerSelectionPolicyError, }; -pub(super) struct DefaultScorer { - config: KvRouterConfig, +/// Resolve optional terms once so disabled weights add neither branches nor conversions to +/// each worker score. The shared-cache specialization also omits range traversal code. +pub(super) fn build( + config: &KvRouterConfig, worker_label: &'static str, plain_decode: bool, +) -> Box { + match ( + config.decode_active_request_weight != 0.0, + config.shared_cache_multiplier != 0.0, + ) { + (false, false) => Box::new(DefaultScorer::::new( + config, + worker_label, + plain_decode, + )), + (false, true) => Box::new(DefaultScorer::::new( + config, + worker_label, + plain_decode, + )), + (true, false) => Box::new(DefaultScorer::::new( + config, + worker_label, + plain_decode, + )), + (true, true) => Box::new(DefaultScorer::::new( + config, + worker_label, + plain_decode, + )), + } +} + +struct DefaultScorer { + overlap_score_credit: f64, + overlap_score_credit_decay: f64, + host_cache_hit_weight: f64, + disk_cache_hit_weight: f64, + shared_cache_multiplier: f64, + decode_active_request_weight: f64, + prefill_load_scale: f64, + is_decode: bool, + plain_decode: bool, + prepared: PreparedRequest, +} + +/// Values shared by every worker score in one selection. Preparation keeps their conversions +/// and role checks outside the per-worker trait call without changing floating-point arithmetic. +#[derive(Default)] +struct PreparedRequest { min_prefill: usize, + block_size: IntegerDivisor, + request_blocks: IntegerDivisor, + overlap_credit: f64, + use_decay: bool, + subtract_from_decode: bool, +} + +/// Division by a power-of-two integer has an exact reciprocal. Both inputs to this helper +/// originate as unsigned token/block counts, so neither result can underflow or overflow. +/// Other divisors retain ordinary division to preserve the score's rounding exactly. +#[derive(Clone, Copy)] +struct IntegerDivisor { + value: f64, + reciprocal: Option, +} + +impl IntegerDivisor { + fn new(value: u64) -> Self { + Self { + value: value as f64, + reciprocal: value.is_power_of_two().then(|| 1.0 / value as f64), + } + } + + #[inline] + fn divide(self, numerator: f64) -> f64 { + match self.reciprocal { + Some(reciprocal) => numerator * reciprocal, + None => numerator / self.value, + } + } +} + +impl Default for IntegerDivisor { + fn default() -> Self { + Self::new(1) + } } -impl DefaultScorer { +impl + DefaultScorer +{ pub(super) fn new( - config: KvRouterConfig, + config: &KvRouterConfig, worker_label: &'static str, plain_decode: bool, ) -> Self { Self { - config, - worker_label, + overlap_score_credit: config.overlap_score_credit, + overlap_score_credit_decay: config.overlap_score_credit_decay, + host_cache_hit_weight: config.host_cache_hit_weight, + disk_cache_hit_weight: config.disk_cache_hit_weight, + shared_cache_multiplier: config.shared_cache_multiplier, + decode_active_request_weight: config.decode_active_request_weight, + prefill_load_scale: config.prefill_load_scale, + is_decode: worker_label == "decode", plain_decode, - min_prefill: 0, + prepared: PreparedRequest::default(), } } } -impl WorkerScorer for DefaultScorer { +impl WorkerScorer + for DefaultScorer +{ fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT } @@ -40,8 +134,23 @@ impl WorkerScorer for DefaultScorer { context: &WorkerSelectionContext<'_>, candidates: &[WorkerCandidate], ) -> Result<(), WorkerSelectionPolicyError> { - self.min_prefill = 0; - if context.tracks_prefill_tokens() && self.config.overlap_score_credit_decay > 0.0 { + // Plain disaggregated decode is load-only. Conditional decode retains cache credit. + let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { + 0.0 + } else { + self.overlap_score_credit + }; + self.prepared = PreparedRequest { + min_prefill: 0, + block_size: IntegerDivisor::new(u64::from(context.block_size())), + request_blocks: IntegerDivisor::new(context.request_blocks()), + overlap_credit, + use_decay: context.tracks_prefill_tokens() && self.overlap_score_credit_decay > 0.0, + subtract_from_decode: self.is_decode + && !context.tracks_prefill_tokens() + && overlap_credit > 0.0, + }; + if self.prepared.use_decay { let mut minimum = usize::MAX; for candidate in candidates { let load = candidate @@ -49,7 +158,7 @@ impl WorkerScorer for DefaultScorer { .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; minimum = minimum.min(load.active_prefill_tokens()); } - self.min_prefill = minimum; + self.prepared.min_prefill = minimum; } Ok(()) } @@ -71,40 +180,39 @@ impl WorkerScorer for DefaultScorer { } else { estimated_overlap }; - let shared = if context.has_tier_matches() { - cache.shared_beyond_device_blocks() - } else { - context - .shared_cache_hits() - .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)) - }; - let decay = - if context.tracks_prefill_tokens() && self.config.overlap_score_credit_decay > 0.0 { - let excess = load - .active_prefill_tokens() - .saturating_sub(self.min_prefill) as f64 - / context.block_size() as f64; - 1.0 / (1.0 - + self.config.overlap_score_credit_decay - * (excess / context.request_blocks() as f64)) + let shared_credit = if SHARED_CREDIT { + let shared = if context.has_tier_matches() { + cache.shared_beyond_device_blocks() } else { - 1.0 + context + .shared_cache_hits() + .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)) }; - // Plain disaggregated decode is load-only. Conditional decode retains cache credit. - let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { - 0.0 + self.shared_cache_multiplier * shared as f64 } else { - self.config.overlap_score_credit + // Preserve signed zero when the configured weight is -0.0. + self.shared_cache_multiplier }; - let credit = overlap_credit * decay * device - + self.config.host_cache_hit_weight * cache.host_overlap_blocks() - + self.config.disk_cache_hit_weight * cache.disk_overlap_blocks() - + self.config.shared_cache_multiplier * shared as f64; - let request_cost = self.config.decode_active_request_weight * load.active_requests() as f64; - let logit = if self.worker_label == "decode" - && !context.tracks_prefill_tokens() - && overlap_credit > 0.0 - { + let decay = if self.prepared.use_decay { + let excess = self.prepared.block_size.divide( + load.active_prefill_tokens() + .saturating_sub(self.prepared.min_prefill) as f64, + ); + 1.0 / (1.0 + + self.overlap_score_credit_decay * self.prepared.request_blocks.divide(excess)) + } else { + 1.0 + }; + let credit = self.prepared.overlap_credit * decay * device + + self.host_cache_hit_weight * cache.host_overlap_blocks() + + self.disk_cache_hit_weight * cache.disk_overlap_blocks() + + shared_credit; + let request_cost = if REQUEST_COST { + self.decode_active_request_weight * load.active_requests() as f64 + } else { + self.decode_active_request_weight + }; + let logit = if self.prepared.subtract_from_decode { (load.decode_cost_blocks() - credit).max(0.0) + request_cost } else { let raw_tokens = if !context.tracks_prefill_tokens() { @@ -115,10 +223,36 @@ impl WorkerScorer for DefaultScorer { } else { context.prompt_tokens() }; - let prefill = (raw_tokens as f64 / context.block_size() as f64 - credit).max(0.0); - self.config.prefill_load_scale * prefill + load.decode_cost_blocks() + request_cost + let prefill = (self.prepared.block_size.divide(raw_tokens as f64) - credit).max(0.0); + self.prefill_load_scale * prefill + load.decode_cost_blocks() + request_cost }; let cost = logit * candidate.preferred_taint_multiplier().unwrap_or(1.0); Ok(cost) } } + +#[cfg(test)] +mod tests { + use super::IntegerDivisor; + + #[test] + fn prepared_division_preserves_score_rounding() { + let counts = [0, 1, 7, 127, 2048, (1 << 53) - 1, 1 << 53, u64::MAX]; + for block_size in (0..32) + .map(|shift| 1u64 << shift) + .chain([3, 17, 63, u32::MAX as u64]) + { + for request_blocks in (0..64).map(|shift| 1u64 << shift).chain([3, 127, u64::MAX]) { + for count in counts { + let blocks = IntegerDivisor::new(block_size).divide(count as f64); + let expected = count as f64 / block_size as f64; + assert_eq!(blocks.to_bits(), expected.to_bits()); + assert_eq!( + IntegerDivisor::new(request_blocks).divide(blocks).to_bits(), + (expected / request_blocks as f64).to_bits(), + ); + } + } + } + } +} diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 5d18f3ab3f16..4152874aa903 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -14,16 +14,18 @@ use support::*; fn seeded_selection_matches_reference_across_cache_and_load_shapes() { for temperature in [0.0, 0.7] { for prompt in [1, 17, 127, 2048] { - for mode in 0..8 { + for mode in 0..32 { let (workers, mut request) = fixture(16, prompt); let config = KvRouterConfig { router_temperature: temperature, overlap_score_credit_decay: 0.6, host_cache_hit_weight: 0.25, disk_cache_hit_weight: 0.1, + decode_active_request_weight: if mode & 8 == 0 { 0.0 } else { 0.7 }, + shared_cache_multiplier: if mode & 16 == 0 { 0.0 } else { 0.6 }, ..Default::default() }; - if mode >= 4 { + if mode % 8 >= 4 { request.shared_cache_hits = Some(dynamo_kv_router::SharedCacheHits::from_ranges(vec![ 1..3, @@ -66,6 +68,96 @@ fn seeded_selection_matches_reference_across_cache_and_load_shapes() { } } +#[test] +fn unseeded_minimum_picker_only_selects_workers_tied_for_lowest_cost() { + let (workers, mut request) = fixture(8, 17); + let policy = default_policy( + KvRouterConfig { + overlap_score_credit: 0.0, + prefill_load_scale: 0.0, + decode_active_request_weight: 1.0, + ..Default::default() + }, + "prefill", + ); + for best in [[0, 1], [2, 5], [6, 7]] { + for (worker, load) in &mut request.worker_loads { + load.active_requests = if best.contains(&worker.worker_id) { + 0 + } else { + 10 + }; + load.active_decode_blocks = 0; + load.additional_active_blocks = 0; + } + for _ in 0..32 { + let selected = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + assert!(best.contains(&selected.worker.worker_id)); + } + } +} + +#[test] +fn prepared_values_follow_each_request_when_reusing_a_policy() { + for label in ["prefill", "decode"] { + for temperature in [0.0, 0.7] { + let config = KvRouterConfig { + router_temperature: temperature, + overlap_score_credit_decay: 0.6, + host_cache_hit_weight: 0.25, + disk_cache_hit_weight: 0.1, + ..Default::default() + }; + let reference = dynamo_kv_router::DefaultWorkerSelector::new_seeded( + Some(config.clone()), + label, + 42, + ); + let plugin = DefaultWorkerSelector::new_seeded(Some(config), label, 42); + for round in 0..64 { + let prompt = [1, 17, 127, 2048][round % 4]; + let block_size = [8, 16, 32, 17, 3][round % 5]; + let (workers, mut request) = fixture(16, prompt); + request.track_prefill_tokens = round % 2 == 0; + if round % 3 == 0 { + request.overlap.tier_overlap_blocks = Default::default(); + } + if round % 5 == 0 { + request.worker_loads.clear(); + } else { + for load in request.worker_loads.values_mut() { + load.active_prefill_tokens += round * 19; + } + } + let input = WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + block_size, + ); + let expected = reference.select_worker(input).unwrap(); + let actual = plugin.select_worker(input).unwrap(); + assert_eq!( + actual.worker, expected.worker, + "label={label} temperature={temperature} round={round}" + ); + assert_eq!(actual.cached_tokens, expected.cached_tokens); + assert_eq!( + actual.potential_decode_blocks, + expected.potential_decode_blocks + ); + } + } + } +} + #[test] fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { struct Inspect; diff --git a/lib/router-plugins/builtin/tests/scorer_prepare.rs b/lib/router-plugins/builtin/tests/scorer_prepare.rs index a7e760ce3816..258380a81505 100644 --- a/lib/router-plugins/builtin/tests/scorer_prepare.rs +++ b/lib/router-plugins/builtin/tests/scorer_prepare.rs @@ -263,3 +263,74 @@ fn rejects_nonfinite_contributions_and_overflow_before_picking() { assert!(error.to_string().contains("non-finite")); } } + +#[test] +fn picker_columns_and_costs_stay_aligned_after_a_scoring_error() { + struct FailSecondScoreOnce(usize); + impl WorkerScorer for FailSecondScoreOnce { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + candidate: &WorkerCandidate, + ) -> Result { + self.0 += 1; + if self.0 == 2 { + return Err(WorkerSelectionPolicyError::failed("score failed")); + } + Ok(candidate.load().unwrap().active_requests() as f64) + } + } + struct CheckColumns; + impl WorkerPicker for CheckColumns { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD + } + fn pick( + &mut self, + _: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + let cache = input.cache().unwrap(); + let load = input.load().unwrap(); + assert_eq!(cache.len(), input.candidates().len()); + assert_eq!(load.len(), input.candidates().len()); + for (row, candidate) in input.candidates().iter().enumerate() { + let worker = candidate.worker(); + let expected_overlap = (worker.worker_id * 7 + u64::from(worker.dp_rank)) % 9; + assert_eq!(cache[row].device_overlap_blocks(), expected_overlap as f64); + assert_eq!(load[row].active_requests(), worker.worker_id as usize % 5); + assert_eq!(candidate.cost(), load[row].active_requests() as f64); + } + Ok(0) + } + } + let (workers, mut request) = fixture(8, 17); + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + vec![Box::new(ExcludeWorkerOne)], + vec![Box::new(FailSecondScoreOnce(0))], + Box::new(CheckColumns), + ); + let select = |request: &dynamo_kv_router::scheduling::SchedulingRequest| { + policy.select_worker(WorkerSelectionInput::configured( + &workers, + request, + request.eligibility(), + 16, + )) + }; + assert!( + select(&request) + .unwrap_err() + .to_string() + .contains("score failed") + ); + request.allowed_worker_ids = Some([1, 3, 4].into_iter().collect()); + assert!([3, 4].contains(&select(&request).unwrap().worker.worker_id)); + request.allowed_worker_ids = None; + assert_ne!(select(&request).unwrap().worker.worker_id, 1); +} From c86d24a3272291a4ad5e2a4905637a27461d4d5f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 18 Sep 2026 23:53:29 +0000 Subject: [PATCH 08/31] fix(router): update queue snapshot cache map fixture Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/queue.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 73e30a4fc5cb..ed5f28e1fcb1 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -4064,7 +4064,7 @@ policy_classes: let (queue, _slots) = make_queue_with_providers(2, 16, 256, None, Some(provider)); let (mut request, _response) = make_request("adapter-request", 256); request.lora_name = Some("adapter".to_string()); - request.overlap.effective_cached_tokens = HashMap::from([ + request.overlap.effective_cached_tokens = FxHashMap::from_iter([ (WorkerWithDpRank::from_worker_id(0), 224), (WorkerWithDpRank::from_worker_id(1), 32), ]); From 5c6cb096491f4ec53cd2db30b1360e94a49a5a56 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 19 Sep 2026 00:28:19 +0000 Subject: [PATCH 09/31] fix(router): integrate builtin selection with unified plugins Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 2 + lib/kv-router/src/plugins/registry.rs | 6 +- lib/kv-router/src/scheduling/AGENTS.md | 2 +- lib/kv-router/src/scheduling/local.rs | 7 +- .../src/scheduling/selector/policy.rs | 11 ++- .../src/services/selection/service.rs | 4 +- lib/kv-router/tests/plugin_api_compat.rs | 2 +- lib/llm/src/kv_router/embedded.rs | 5 +- lib/llm/src/kv_router/plugins.rs | 5 +- lib/router-plugins/builtin/src/default/mod.rs | 5 +- .../builtin/src/default/parameters.rs | 10 ++- .../builtin/src/default/picker.rs | 2 +- .../builtin/src/default/scorer.rs | 6 +- .../builtin/src/default/selector.rs | 6 +- .../builtin/tests/plugin_registry.rs | 79 +++++++++++++++++++ 15 files changed, 121 insertions(+), 31 deletions(-) create mode 100644 lib/router-plugins/builtin/tests/plugin_registry.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 9bd5f486948e..07ffe9127f4e 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -628,6 +628,8 @@ The [example README](https://github.com/ai-dynamo/dynamo/blob/main/examples/rout The default worker selector lives in `dynamo-custom-policy-builtin`, beside the two-tier policy. Routing hosts install `default_registry()` even when they do not link a custom catalog. An omitted policy or the reserved `default` selection uses that factory. A Rust embedding that constructs `SelectionServiceBuilder` must supply a registry with a default factory or an explicit policy factory. +The default uses the same `RouterPluginRegistry` as custom worker policies and request classifiers. Installing it preserves the other registered providers. Resolving a classifier-only configuration still leaves worker selection on the builtin default; installing the default alone does not enable custom-plugin frontend restrictions. + To configure a separate default-policy instance, use the named provider: ```yaml diff --git a/lib/kv-router/src/plugins/registry.rs b/lib/kv-router/src/plugins/registry.rs index 90554fa0fdfa..91fc6a718ce3 100644 --- a/lib/kv-router/src/plugins/registry.rs +++ b/lib/kv-router/src/plugins/registry.rs @@ -118,7 +118,11 @@ impl RouterPluginRegistry { Ok(super::RouterPlugins { // The builtin default does not opt a frontend into custom-plugin restrictions. // Embedded routers still obtain it from the registry when constructing a pool. - worker_selection: if config.selected_worker_selection_policy_instance()?.is_some() { + worker_selection: if config + .selected_worker_selection_policy_instance() + .map_err(WorkerSelectionPolicyRegistryError::from)? + .is_some() + { self.resolve(config)? } else { None diff --git a/lib/kv-router/src/scheduling/AGENTS.md b/lib/kv-router/src/scheduling/AGENTS.md index 9186a7c38299..22d6d5289467 100644 --- a/lib/kv-router/src/scheduling/AGENTS.md +++ b/lib/kv-router/src/scheduling/AGENTS.md @@ -122,7 +122,7 @@ PolicyClassQueue("agents") ## Public Worker-Selection API -The `selector` module contains the public Rust contract for custom worker filters, scorers, and pickers. Treat each public item as a versioned external API. +The `plugins::worker_selection` module contains the public Rust contract for custom worker filters, scorers, and pickers. Treat each public item as a versioned external API. - Do not add a public field, accessor, input group, type, or re-export unless the task explicitly requires a new external policy capability. - An internal need in the default scorer, logging, tests, or SelectionService does not justify a public API addition. diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 00e8f82d8b3f..fdf8ef64e3bb 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -1143,12 +1143,13 @@ mod tests { .unwrap(); assert_eq!(calls.load(Ordering::Relaxed), 1); assert_eq!( - scheduler.get_potential_loads(None, 0, HashMap::new(), false)[0].active_requests, + scheduler.get_potential_loads(None, 0, FxHashMap::default(), false)[0].active_requests, 1 ); drop(booking.expect("admission must return an armed booking")); tokio::time::timeout(Duration::from_secs(1), async { - while scheduler.get_potential_loads(None, 0, HashMap::new(), false)[0].active_requests + while scheduler.get_potential_loads(None, 0, FxHashMap::default(), false)[0] + .active_requests != 0 { tokio::task::yield_now().await; @@ -1193,7 +1194,7 @@ mod tests { .await; assert!(matches!(result, Err(KvSchedulerError::DeadlineExceeded))); assert_eq!( - scheduler.get_potential_loads(None, 0, HashMap::new(), false)[0].active_requests, + scheduler.get_potential_loads(None, 0, FxHashMap::default(), false)[0].active_requests, 0 ); cancel_token.cancel(); diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index e48d12d3b568..e6ff9a82fcf1 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -10,14 +10,10 @@ use super::{ MaterializedSelectionInput, WorkerSelectionInput, WorkerSelector, select_worker_with_policy, }; -use crate::protocols::{ - WorkerConfigLike, WorkerId, WorkerSelectionResult, -}; +use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult}; use crate::scheduling::config::KvRouterConfig; use crate::scheduling::filter::RoutingEligibility; -use crate::scheduling::types::{ - KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError, -}; +use crate::scheduling::types::{KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError}; use crate::plugins::worker_selection::{ ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputs, @@ -369,6 +365,9 @@ impl WorkerSelector for WorkerSelectionPolicy { #[cfg(test)] mod tests { + use crate::plugins::worker_selection::WorkerInputView; + use crate::protocols::WorkerWithDpRank; + use crate::scheduling::SessionContext; use std::{ cell::Cell, collections::{HashMap, HashSet}, diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 05332d800a8f..0556ac7d8a6b 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -165,9 +165,7 @@ impl SelectionServiceBuilder { }; #[cfg(not(test))] if worker_selection_policy_factory.is_none() { - return Err( - super::policy_registry::WorkerSelectionPolicyRegistryError::MissingDefault.into(), - ); + return Err(crate::plugins::WorkerSelectionPolicyRegistryError::MissingDefault.into()); } let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); let indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; diff --git a/lib/kv-router/tests/plugin_api_compat.rs b/lib/kv-router/tests/plugin_api_compat.rs index 27964d9da8a2..99e28dae6b3b 100644 --- a/lib/kv-router/tests/plugin_api_compat.rs +++ b/lib/kv-router/tests/plugin_api_compat.rs @@ -110,7 +110,7 @@ fn legacy_plugins_resolve_through_the_common_registry() { .unwrap(); let registry: RouterPluginRegistry = registry; let policy = tempfile::NamedTempFile::new().unwrap(); - std::fs::write(policy.path(), "worker_selection:\n aggregated: legacy\n instances:\n - name: legacy\n type: legacy\n").unwrap(); + std::fs::write(policy.path(), "worker_selection:\n aggregated: legacy\n prefill: legacy\n decode: legacy\n encode: legacy\n instances:\n - name: legacy\n type: legacy\n").unwrap(); let config = KvRouterConfig { router_policy_config: Some(policy.path().display().to_string()), ..Default::default() diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index f3585409f832..583b5319042a 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -9,13 +9,14 @@ //! frontend retains transport, stream leases, and request-expiry ownership. use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::time::Duration; use anyhow::{Context, Result}; use dynamo_kv_router::WorkerType; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::identity::RoutingPartitionId; +use dynamo_kv_router::plugins::RouterPluginRegistry; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use dynamo_kv_router::scheduling::queue::DEFAULT_MAX_BATCHED_TOKENS; use dynamo_kv_router::scheduling::{ @@ -253,7 +254,7 @@ impl EmbeddedSelection { let service = SelectionServiceBuilder::new( args.kv_router_config.clone(), worker_type, - WorkerSelectionPolicyRegistry::default(), + RouterPluginRegistry::default(), ) .worker_selection_policy_factory(args.policy_factory) .host_manages_request_lifecycle() diff --git a/lib/llm/src/kv_router/plugins.rs b/lib/llm/src/kv_router/plugins.rs index 35f5b0aa802d..fb1fad7ed02b 100644 --- a/lib/llm/src/kv_router/plugins.rs +++ b/lib/llm/src/kv_router/plugins.rs @@ -21,7 +21,10 @@ pub fn install_router_plugin_registry(registry: RouterPluginRegistry) -> bool { /// The installed catalog, or the required builtin default policy. pub fn router_plugin_registry() -> RouterPluginRegistry { - INSTALLED_PLUGINS.get().cloned().unwrap_or_else(dynamo_custom_policy_builtin::default_registry) + INSTALLED_PLUGINS + .get() + .cloned() + .unwrap_or_else(dynamo_custom_policy_builtin::default_registry) } /// Carries configured plugins through the host's shared router construction. diff --git a/lib/router-plugins/builtin/src/default/mod.rs b/lib/router-plugins/builtin/src/default/mod.rs index 3d4268090f2f..121760a355f3 100644 --- a/lib/router-plugins/builtin/src/default/mod.rs +++ b/lib/router-plugins/builtin/src/default/mod.rs @@ -11,7 +11,10 @@ mod selector; pub(super) use parameters::register; pub use selector::DefaultWorkerSelector; -use dynamo_kv_router::{KvRouterConfig, WorkerSelectionPolicy, WorkerSelectionPolicyFactory}; +use dynamo_kv_router::KvRouterConfig; +use dynamo_kv_router::plugins::worker_selection::{ + WorkerSelectionPolicy, WorkerSelectionPolicyFactory, +}; use parking_lot::Mutex; use std::sync::Arc; diff --git a/lib/router-plugins/builtin/src/default/parameters.rs b/lib/router-plugins/builtin/src/default/parameters.rs index 6ac4fcbd8bd5..5c44d000fff1 100644 --- a/lib/router-plugins/builtin/src/default/parameters.rs +++ b/lib/router-plugins/builtin/src/default/parameters.rs @@ -4,6 +4,9 @@ //! Startup parameters and provider registration for the default policy. use dynamo_kv_router::KvRouterConfig; +use dynamo_kv_router::plugins::{ + RouterPluginRegistry, WorkerSelectionPolicyProviderError, WorkerSelectionPolicyRegistryError, +}; use std::sync::Arc; use super::policy_for_role; @@ -22,10 +25,9 @@ struct Parameters { } pub(crate) fn register( - registry: &mut dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistry, -) -> Result<(), dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistryError> { - use dynamo_kv_router::services::selection::WorkerSelectionPolicyProviderError; - registry.register( + registry: &mut RouterPluginRegistry, +) -> Result<(), WorkerSelectionPolicyRegistryError> { + registry.register_worker_selection( "dynamo-default-cost-fn", Arc::new(|parameters| { let parameters: Parameters = parameters.deserialize()?; diff --git a/lib/router-plugins/builtin/src/default/picker.rs b/lib/router-plugins/builtin/src/default/picker.rs index ce0e3bfa98c7..3b38f77e0413 100644 --- a/lib/router-plugins/builtin/src/default/picker.rs +++ b/lib/router-plugins/builtin/src/default/picker.rs @@ -3,7 +3,7 @@ //! Minimum-cost selection, tie-breaking, and temperature sampling of scored candidates. -use dynamo_kv_router::{ +use dynamo_kv_router::plugins::worker_selection::{ WorkerInputView, WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicyError, }; use parking_lot::Mutex; diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index c640ce5951d4..0dd0b3cef8da 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -3,9 +3,9 @@ //! Cache and load cost calculation. The host owns snapshots and validates finite scores. -use dynamo_kv_router::{ - KvRouterConfig, WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, - WorkerSelectionPolicyError, +use dynamo_kv_router::KvRouterConfig; +use dynamo_kv_router::plugins::worker_selection::{ + WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, }; /// Resolve optional terms once so disabled weights add neither branches nor conversions to diff --git a/lib/router-plugins/builtin/src/default/selector.rs b/lib/router-plugins/builtin/src/default/selector.rs index 3742aa4fd85a..36bd8e0f677c 100644 --- a/lib/router-plugins/builtin/src/default/selector.rs +++ b/lib/router-plugins/builtin/src/default/selector.rs @@ -3,11 +3,9 @@ //! Synchronous adapter for direct selection and deterministic replay. +use dynamo_kv_router::plugins::worker_selection::{WorkerInputs, WorkerSelectionPolicy}; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerSelectionResult}; -use dynamo_kv_router::{ - KvRouterConfig, KvSchedulerError, WorkerInputs, WorkerSelectionInput, WorkerSelectionPolicy, - WorkerSelector, -}; +use dynamo_kv_router::{KvRouterConfig, KvSchedulerError, WorkerSelectionInput, WorkerSelector}; use parking_lot::Mutex; use std::sync::Arc; diff --git a/lib/router-plugins/builtin/tests/plugin_registry.rs b/lib/router-plugins/builtin/tests/plugin_registry.rs new file mode 100644 index 000000000000..2e802e8bef7d --- /dev/null +++ b/lib/router-plugins/builtin/tests/plugin_registry.rs @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use dynamo_kv_router::plugins::RouterPluginRegistry; +use dynamo_kv_router::plugins::request_classifier::{ + ClassifyFuture, ClassifyRequest, RequestClassifier, +}; +use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerType}; + +struct PassThrough; + +impl RequestClassifier for PassThrough { + fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + Box::pin(async move { Ok(request) }) + } +} + +fn config(yaml: &str) -> (tempfile::NamedTempFile, KvRouterConfig) { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(file.path(), yaml).unwrap(); + let config = KvRouterConfig { + router_policy_config: Some(file.path().display().to_string()), + ..Default::default() + }; + (file, config) +} + +#[test] +fn builtin_default_does_not_opt_into_custom_frontend_restrictions() { + let mut registry = dynamo_custom_policy_builtin::default_registry(); + dynamo_custom_policy_builtin::register(&mut registry).unwrap(); + let (_file, explicit_default) = config("worker_selection:\n aggregated: default\n"); + for config in [KvRouterConfig::default(), explicit_default] { + assert!(registry.resolve(&config).unwrap().is_some()); + assert!(registry.resolve_plugins(&config).unwrap().is_empty()); + } +} + +#[test] +fn builtin_registration_preserves_classifiers_and_resolves_mixed_pools() { + for selection in ["default", "named-default"] { + let mut registry = RouterPluginRegistry::default(); + registry + .register_request_classifier( + "pass-through", + Arc::new(|_| Ok(Arc::new(|_| Box::new(PassThrough)))), + ) + .unwrap(); + // Installing the default must preserve classifier providers already in the catalog. + dynamo_custom_policy_builtin::register(&mut registry).unwrap(); + let (_file, config) = config(&format!( + r#" +request_classifier: + type: pass-through +worker_selection: + prefill: {selection} + instances: + - name: named-default + type: dynamo-default-cost-fn +"# + )); + let plugins = registry.resolve_plugins(&config).unwrap(); + assert!(!plugins.is_empty()); + assert!(plugins.request_classifier().is_some()); + assert_eq!(plugins.worker_selection().is_some(), selection != "default"); + // With only a classifier selected, the host still resolves the builtin policy. + let factory = registry.resolve(&config).unwrap().unwrap(); + for role in [ + WorkerType::Aggregated, + WorkerType::Prefill, + WorkerType::Decode, + WorkerType::Encode, + ] { + factory(&config, role, RoutingPartitionRef::new("model", "default")); + } + } +} From d39fa9fe837a77ff9aa60ee0cbe6003745ad3d2a Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 19 Sep 2026 00:43:23 +0000 Subject: [PATCH 10/31] fix(router): update stacked plugin fixtures and feature gates Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/lib.rs | 2 +- lib/llm/src/discovery/watcher.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index dae7cddac38c..539210a0e79a 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -40,7 +40,7 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, }; -#[cfg(any(feature = "custom-policy", feature = "select-service"))] +#[cfg(feature = "select-service")] use dynamo_kv_router::plugins::RouterPluginRegistry; use dynamo_kv_router::{KvRouterConfig, plugins::RouterPlugins}; use dynamo_llm::entrypoint::RouterConfig; diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 7b192d9d0efa..ce51c0df343b 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -2139,7 +2139,7 @@ request_classifier: } }); let selectors = Arc::new(parking_lot::Mutex::new(Vec::new())); - let mut registry = dynamo_kv_router::plugins::RouterPluginRegistry::default(); + let mut registry = dynamo_custom_policy_builtin::default_registry(); registry .register_request_classifier("test", Arc::new(move |_| Ok(factory.clone()))) .unwrap(); @@ -2152,7 +2152,7 @@ request_classifier: let selectors = selectors.clone(); Ok(Arc::new(move |config, role, partition| { selectors.lock().push((role, partition.into_owned())); - dynamo_kv_router::WorkerSelectionPolicy::default( + dynamo_custom_policy_builtin::default_policy( config.clone(), role.default_selector_label(), ) From 711bc8861eca7ceb61c7b8be37b771ba20cafd2f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 18:02:44 +0000 Subject: [PATCH 11/31] refactor(router): resolve builtin policy at startup Signed-off-by: Ishan Dhanani --- .../inference-gateway/ext-proc/src/runner.rs | 6 +- .../router/custom-worker-selection.mdx | 7 +- lib/bindings/python/rust/lib.rs | 20 ++--- lib/bindings/python/rust/llm/entrypoint.rs | 68 +++++++++++---- lib/bindings/python/rust/llm/kv.rs | 8 +- lib/kv-router/benches/selection_core.rs | 2 +- lib/kv-router/benches/worker_selection.rs | 2 +- lib/kv-router/src/plugins.rs | 13 +++ lib/kv-router/src/plugins/registry.rs | 83 ++++++++----------- lib/kv-router/src/scheduling/selector/mod.rs | 16 ++-- .../src/scheduling/selector/policy.rs | 56 ++++++------- .../src/scheduling/selector/reference.rs | 3 +- .../src/services/selection/core/mod.rs | 6 +- .../src/services/selection/core/tests.rs | 11 ++- .../src/services/selection/core/workers.rs | 20 ++--- .../src/services/selection/membership.rs | 2 +- .../src/services/selection/service.rs | 64 ++++++++------ lib/kv-router/src/services/selection/tests.rs | 18 +++- lib/llm/src/discovery/watcher.rs | 38 +++++---- lib/llm/src/entrypoint/input/http.rs | 4 +- lib/llm/src/kv_router.rs | 22 ++--- lib/llm/src/kv_router/plugins.rs | 4 + .../builtin/benches/default_lifecycle.rs | 2 +- lib/router-plugins/builtin/src/default/mod.rs | 22 +++-- .../builtin/src/default/parameters.rs | 73 ++++++++++------ .../builtin/src/default/scorer.rs | 6 +- .../builtin/src/default/selector.rs | 1 + lib/router-plugins/builtin/src/lib.rs | 5 +- .../builtin/tests/plugin_registry.rs | 39 +++++++-- 29 files changed, 370 insertions(+), 251 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/runner.rs b/deploy/inference-gateway/ext-proc/src/runner.rs index 2c8abc97af84..102e6ef9bb90 100644 --- a/deploy/inference-gateway/ext-proc/src/runner.rs +++ b/deploy/inference-gateway/ext-proc/src/runner.rs @@ -134,11 +134,7 @@ pub async fn run(policy_registry: Option) -> Resu { anyhow::bail!("linked worker-selection policies require DYN_EPP_MODE=standalone") } - run_inner( - mode, - policy_registry.unwrap_or_else(dynamo_custom_policy_builtin::default_registry), - ) - .await + run_inner(mode, policy_registry.unwrap_or_default()).await } fn init_tracing() { diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 07ffe9127f4e..a28ccbc368ee 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -44,6 +44,7 @@ Write your own policy when no shipped policy expresses the rule you need. The re An external policy owns its filters, scorers, and picker. The default scorer and picker are private components of `dynamo-custom-policy-builtin`. They implement the same public traits as external policies. + ## Public Plugin API and Compatibility Use `dynamo_kv_router::plugins` for plugin registration and these modules for plugin contracts: @@ -626,9 +627,11 @@ The [example README](https://github.com/ai-dynamo/dynamo/blob/main/examples/rout ## Default Policy Plugin -The default worker selector lives in `dynamo-custom-policy-builtin`, beside the two-tier policy. Routing hosts install `default_registry()` even when they do not link a custom catalog. An omitted policy or the reserved `default` selection uses that factory. A Rust embedding that constructs `SelectionServiceBuilder` must supply a registry with a default factory or an explicit policy factory. +The default worker selector lives in `dynamo-custom-policy-builtin`, beside the two-tier policy. Routing hosts install its default factory once, even without a custom catalog. An omitted policy or the reserved `default` selection uses that factory. Optional catalogs register named providers without replacing the default or existing classifiers. + +A Rust embedding that constructs `SelectionServiceBuilder` must supply a registry with a default factory, such as `default_registry()`, or an explicit policy factory. A missing factory fails construction before workers or synchronization resources start. Each routing partition receives its own policy instance. -The default uses the same `RouterPluginRegistry` as custom worker policies and request classifiers. Installing it preserves the other registered providers. Resolving a classifier-only configuration still leaves worker selection on the builtin default; installing the default alone does not enable custom-plugin frontend restrictions. +Resolving `RouterPlugins` from the host registry includes the default factory even when no named policy is selected. `RouterPlugins::default()` remains an empty builder input. Use `has_custom_worker_selection()` to detect an explicit named or programmatic policy, and `has_custom_plugins()` to detect that policy or a classifier. Factory presence alone does not imply a custom selection. These checks preserve stock frontend modes and avoid waiting for a typed model card solely because the builtin factory is present. A classifier-only configuration still uses the builtin default for worker selection. To configure a separate default-policy instance, use the named provider: diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 539210a0e79a..fe3412f0c8c4 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -426,35 +426,29 @@ pub(crate) fn router_plugins(config: &KvRouterConfig) -> anyhow::Result RouterPluginRegistry { - #[cfg(feature = "custom-policy")] - { - dynamo_llm::kv_router::plugins::router_plugin_registry() - } - - #[cfg(not(feature = "custom-policy"))] - { - dynamo_custom_policy_builtin::default_registry() - } + dynamo_llm::kv_router::plugins::router_plugin_registry() } #[cfg(feature = "custom-policy")] fn register_core_with_router_plugins(m: &Bound<'_, PyModule>) -> PyResult<()> { - let mut registry = dynamo_custom_policy_builtin::default_registry(); + let mut registry = dynamo_kv_router::plugins::RouterPluginRegistry::default(); // The policies Dynamo ships register first, so a replaced catalog that reuses one of their // type names fails here instead of silently overriding it. dynamo_custom_policy_builtin::register(&mut registry) diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 256f3b1e5ffc..6dab6ba64939 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -1026,26 +1026,18 @@ pub fn run_input<'p>( .kv_router_config, ) .map_err(to_pyerr)?; - let has_router_plugin = !plugins.is_empty(); - if has_router_plugin - && !engine_config + validate_router_plugin_input( + &plugins, + &input_enum, + engine_config .inner .local_model() .router_config() .router_mode - .is_kv_routing() - { - return Err(PyValueError::new_err( - "linked router plugins require --router-mode kv", - )); - } - if has_router_plugin && !matches!(&input_enum, Input::Http) { - return Err(PyValueError::new_err( - "linked router plugins require HTTP frontend input", - )); - } + .is_kv_routing(), + )?; crate::future_into_py(py, async move { - if has_router_plugin { + if matches!(&input_enum, Input::Http) { HttpFrontend::default() .frontend_route_extensions(frontend_route_extensions) .plugins(plugins) @@ -1066,9 +1058,55 @@ pub fn run_input<'p>( }) } +fn validate_router_plugin_input( + plugins: &dynamo_kv_router::plugins::RouterPlugins, + input: &Input, + is_kv_routing: bool, +) -> PyResult<()> { + if plugins.has_custom_plugins() { + if !is_kv_routing { + return Err(PyValueError::new_err( + "linked router plugins require --router-mode kv", + )); + } + if !matches!(input, Input::Http) { + return Err(PyValueError::new_err( + "linked router plugins require HTTP frontend input", + )); + } + } + Ok(()) +} + pub fn to_pyerr(err: E) -> PyErr where E: Display, { PyException::new_err(format!("{}", err)) } + +#[cfg(test)] +mod plugin_input_tests { + use super::*; + use dynamo_kv_router::plugins::RouterPlugins; + + #[test] + fn resolved_builtin_supports_stock_inputs_and_routing_modes() { + let plugins = crate::router_plugins(&RsKvRouterConfig::default()).unwrap(); + assert!(plugins.worker_selection().is_some()); + for input in [Input::Http, Input::Grpc, Input::Stdin, Input::Text] { + for is_kv in [false, true] { + assert!(validate_router_plugin_input(&plugins, &input, is_kv).is_ok()); + } + } + } + + #[test] + fn explicit_policy_retains_frontend_restrictions() { + let plugins = RouterPlugins::default() + .with_worker_selection(dynamo_custom_policy_builtin::default_factory()); + assert!(validate_router_plugin_input(&plugins, &Input::Http, true).is_ok()); + assert!(validate_router_plugin_input(&plugins, &Input::Http, false).is_err()); + assert!(validate_router_plugin_input(&plugins, &Input::Grpc, true).is_err()); + } +} diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 10602d3dc5f4..f67ebf288eb7 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -1934,7 +1934,7 @@ async fn create_kv_router_from_endpoint( .as_ref() .map(|cfg| cfg.use_remote_indexer || cfg.serve_indexer) .unwrap_or(false); - let needs_policy_role = plugins.worker_selection().is_some(); + let needs_policy_role = plugins.has_custom_worker_selection(); let (model_name, policy_model_name, enable_eagle, worker_role, policy_worker_role, load_source) = { let maybe_card = if needs_model_name || needs_policy_role { let wait_secs: u64 = std::env::var("DYN_ROUTER_MODEL_CARD_WAIT_SECS") @@ -2032,7 +2032,11 @@ async fn create_kv_router_from_endpoint( // Preserve the model card's role and display name, which can differ // from the metric role and the routing partition name. - let plugins = if let Some(factory) = plugins.worker_selection().cloned() { + let plugins = if let Some(factory) = plugins + .worker_selection() + .filter(|_| needs_policy_role) + .cloned() + { let policy_worker_role = policy_worker_role .expect("a configured worker-selection policy waits for a typed model card above"); let policy_model_name = policy_model_name.unwrap_or_default(); diff --git a/lib/kv-router/benches/selection_core.rs b/lib/kv-router/benches/selection_core.rs index ade4d3f6c197..cfe3bd52158e 100644 --- a/lib/kv-router/benches/selection_core.rs +++ b/lib/kv-router/benches/selection_core.rs @@ -46,7 +46,7 @@ fn core_with_workers( CancellationToken::new(), SelectionCacheConfig::default(), std::sync::Arc::new(|config, role, _| { - dynamo_kv_router::WorkerSelectionPolicy::default( + dynamo_kv_router::WorkerSelectionPolicy::reference( config.clone(), role.default_selector_label(), ) diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index 0ed136eedaf4..e753cc6952e5 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -373,7 +373,7 @@ fn default_policy_wrapper(c: &mut Criterion) { for worker_count in [2, 32, 1_024] { let (workers, request) = fixture(worker_count); let direct = DefaultWorkerSelector::new(Some(config.clone()), "prefill"); - let policy = WorkerSelectionPolicy::default(config.clone(), "prefill"); + let policy = WorkerSelectionPolicy::reference(config.clone(), "prefill"); group.throughput(Throughput::Elements(worker_count as u64)); group.bench_with_input( diff --git a/lib/kv-router/src/plugins.rs b/lib/kv-router/src/plugins.rs index ecc205c80937..f1eac4b52fc3 100644 --- a/lib/kv-router/src/plugins.rs +++ b/lib/kv-router/src/plugins.rs @@ -27,12 +27,14 @@ use worker_selection::WorkerSelectionPolicyFactory; #[derive(Clone, Default)] pub struct RouterPlugins { worker_selection: Option, + custom_worker_selection: bool, request_classifier: Option, } impl RouterPlugins { pub fn with_worker_selection(mut self, factory: WorkerSelectionPolicyFactory) -> Self { self.worker_selection = Some(factory); + self.custom_worker_selection = true; self } @@ -45,6 +47,17 @@ impl RouterPlugins { self.worker_selection.is_none() && self.request_classifier.is_none() } + /// Whether worker selection was explicitly supplied, rather than filled by the host default. + /// Hosts use this to require typed worker discovery without delaying stock router startup. + pub fn has_custom_worker_selection(&self) -> bool { + self.custom_worker_selection + } + + /// Whether explicit policy or classifier choices require custom-plugin frontend support. + pub fn has_custom_plugins(&self) -> bool { + self.custom_worker_selection || self.request_classifier.is_some() + } + pub fn worker_selection(&self) -> Option<&WorkerSelectionPolicyFactory> { self.worker_selection.as_ref() } diff --git a/lib/kv-router/src/plugins/registry.rs b/lib/kv-router/src/plugins/registry.rs index 91fc6a718ce3..f8d54c097b8b 100644 --- a/lib/kv-router/src/plugins/registry.rs +++ b/lib/kv-router/src/plugins/registry.rs @@ -13,8 +13,6 @@ use super::request_classifier::RequestClassifierRegistry; use super::request_classifier::{ RequestClassifierFactory, RequestClassifierProvider, RequestClassifierRegistryError, }; -#[cfg(test)] -use super::worker_selection::WorkerSelectionPolicy; use super::worker_selection::WorkerSelectionPolicyFactory; use crate::WorkerType; use crate::config::KvRouterConfig; @@ -116,17 +114,11 @@ impl RouterPluginRegistry { config: &KvRouterConfig, ) -> Result { Ok(super::RouterPlugins { - // The builtin default does not opt a frontend into custom-plugin restrictions. - // Embedded routers still obtain it from the registry when constructing a pool. - worker_selection: if config + worker_selection: self.resolve(config)?, + custom_worker_selection: config .selected_worker_selection_policy_instance() .map_err(WorkerSelectionPolicyRegistryError::from)? - .is_some() - { - self.resolve(config)? - } else { - None - }, + .is_some(), request_classifier: self.resolve_request_classifier(config)?, }) } @@ -239,45 +231,25 @@ impl RouterPluginRegistry { return Ok(self.default_factory.clone()); } - #[cfg(not(test))] - if self.default_factory.is_none() - && [ - aggregated.is_none(), - prefill.is_none(), - decode.is_none(), - encode.is_none(), - ] - .contains(&true) - { - return Err(WorkerSelectionPolicyRegistryError::MissingDefault); - } - let default_factory = self.default_factory.clone(); + let [aggregated, prefill, decode, encode] = [aggregated, prefill, decode, encode] + .map(|factory| factory.or_else(|| self.default_factory.clone())); + let required = |factory: Option| { + factory.ok_or(WorkerSelectionPolicyRegistryError::MissingDefault) + }; + let (aggregated, prefill, decode, encode) = ( + required(aggregated)?, + required(prefill)?, + required(decode)?, + required(encode)?, + ); Ok(Some(Arc::new(move |config, worker_type, partition| { let selected = match worker_type { - WorkerType::Aggregated => aggregated.as_ref(), - WorkerType::Prefill => prefill.as_ref(), - WorkerType::Decode => decode.as_ref(), - WorkerType::Encode => encode.as_ref(), + WorkerType::Aggregated => &aggregated, + WorkerType::Prefill => &prefill, + WorkerType::Decode => &decode, + WorkerType::Encode => &encode, }; - match selected { - Some(factory) => factory(config, worker_type, partition), - None => { - #[cfg(test)] - if default_factory.is_none() { - return WorkerSelectionPolicy::default( - config.clone(), - worker_type.default_selector_label(), - ); - } - default_factory - .as_ref() - .expect("default factory validated at resolution")( - config, - worker_type, - partition, - ) - } - } + selected(config, worker_type, partition) }))) } @@ -573,6 +545,23 @@ worker_selection: ); } + #[test] + fn partial_role_selection_requires_an_explicit_fallback() { + let mut registry = WorkerSelectionPolicyRegistry::default(); + registry.register("alpha", Arc::new(provider)).unwrap(); + let config = config(); + let selections = WorkerSelectionPolicySelections { + aggregated: Some("first".to_string()), + prefill: None, + decode: None, + encode: None, + }; + assert!(matches!( + registry.resolve_selections(Some(&config), selections), + Err(WorkerSelectionPolicyRegistryError::MissingDefault) + )); + } + #[test] fn role_scoped_resolution_ignores_unlinked_policy_types_for_other_roles() { let policy_file = tempfile::NamedTempFile::new().unwrap(); diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 3c40b0e4f1db..7a893913dfa7 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -19,13 +19,10 @@ pub use crate::plugins::worker_selection::{ use reference::{DefaultWorkerPicker, DefaultWorkerScorer}; pub use policy::WorkerSelectionPolicy; -use policy::{ - CustomWorkerSelectionState, WorkerSelectionPolicyStateRef, collect_custom_candidates, -}; +use policy::{ComposedPolicyState, WorkerSelectionPolicyStateRef, collect_policy_candidates}; #[cfg(any(test, feature = "bench"))] use reference::pick_default_worker; -use super::config::KvRouterConfig; use super::filter::{RoutingEligibility, WorkerEligibilityError}; use super::types::{KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError}; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank}; @@ -342,7 +339,6 @@ fn log_selection( // 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( - kv_router_config: &KvRouterConfig, worker_type: &'static str, state: WorkerSelectionPolicyStateRef<'_>, workers: &HashMap, @@ -365,24 +361,22 @@ fn select_worker_with_policy( } } - #[cfg(not(any(test, feature = "bench")))] - let _ = kv_router_config; let mut input = MaterializedSelectionInput::new(request, block_size); input.context.pinned_worker = eligibility.pinned_worker(); let selected = match state { #[cfg(any(test, feature = "bench"))] - WorkerSelectionPolicyStateRef::Default(picker) => { + WorkerSelectionPolicyStateRef::Reference(kv_router_config, picker) => { let scorer = DefaultWorkerScorer { kv_router_config, worker_type, }; pick_default_worker(&scorer, picker, &input, workers, request, eligibility) } - WorkerSelectionPolicyStateRef::Custom(state) => { + WorkerSelectionPolicyStateRef::Composed(state) => { let mut state = state.borrow_mut(); let has_eligible_worker = - collect_custom_candidates(&mut state, &input, workers, request, eligibility)?; - let CustomWorkerSelectionState { + collect_policy_candidates(&mut state, &input, workers, request, eligibility)?; + let ComposedPolicyState { picker, picker_inputs, candidates, diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index e6ff9a82fcf1..d55f7f9ca112 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -23,18 +23,18 @@ use crate::plugins::worker_selection::{ #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] pub(super) enum WorkerSelectionPolicyState { #[cfg(any(test, feature = "bench"))] - Default(DefaultWorkerPicker), + Reference(Box, DefaultWorkerPicker), /// Policy-local state owned and called serially by one scheduler queue actor. - Custom(RefCell), + Composed(RefCell), } pub(super) enum WorkerSelectionPolicyStateRef<'a> { #[cfg(any(test, feature = "bench"))] - Default(&'a DefaultWorkerPicker), - Custom(&'a RefCell), + Reference(&'a KvRouterConfig, &'a DefaultWorkerPicker), + Composed(&'a RefCell), } -pub(super) struct CustomWorkerSelectionState { +pub(super) struct ComposedPolicyState { pub(super) filters: Vec>, pub(super) scorers: Vec>, pub(super) picker: Box, @@ -52,17 +52,17 @@ pub(super) struct CustomWorkerSelectionState { /// Routing hosts supply a policy factory. Both builtin and external policies compose /// scorers and a picker through [`Self::new`], with optional filters through [`Self::new_with_filters`]. pub struct WorkerSelectionPolicy { - kv_router_config: KvRouterConfig, worker_label: &'static str, state: WorkerSelectionPolicyState, exclusive_affinity: bool, } impl WorkerSelectionPolicy { - /// Build a custom policy with no filters. + /// Build a policy with no filters. /// /// `worker_label` identifies the worker pool in routing logs. A typed policy factory normally - /// passes [`crate::WorkerType::as_str`]. + /// passes [`crate::WorkerType::as_str`]. The config argument is retained for API compatibility; + /// policy parameters belong to the supplied scorers and picker and are not retained here. pub fn new( kv_router_config: KvRouterConfig, worker_label: &'static str, @@ -72,12 +72,13 @@ impl WorkerSelectionPolicy { Self::new_with_filters(kv_router_config, worker_label, Vec::new(), scorers, picker) } - /// Build a custom policy from ordered filters, additive scorers, and one picker. + /// Build a policy from ordered filters, additive scorers, and one picker. /// /// `worker_label` identifies the worker pool in routing logs. A typed policy factory normally - /// passes [`crate::WorkerType::as_str`]. + /// passes [`crate::WorkerType::as_str`]. The config argument is retained for API compatibility; + /// policy parameters belong to the supplied scorers and picker and are not retained here. pub fn new_with_filters( - kv_router_config: KvRouterConfig, + _kv_router_config: KvRouterConfig, worker_label: &'static str, filters: Vec>, scorers: Vec>, @@ -91,10 +92,9 @@ impl WorkerSelectionPolicy { inputs | scorer.required_worker_inputs() }); Self { - kv_router_config, worker_label, exclusive_affinity: false, - state: WorkerSelectionPolicyState::Custom(RefCell::new(CustomWorkerSelectionState { + state: WorkerSelectionPolicyState::Composed(RefCell::new(ComposedPolicyState { filters, scorers, picker, @@ -116,18 +116,17 @@ impl WorkerSelectionPolicy { self } - /// Wrap Dynamo's built-in selector for a host that uses the policy selector type. + /// Construct the native reference implementation for parity tests and benchmarks. /// /// `worker_label` selects the built-in scoring and logging contract. Typed hosts use /// [`crate::WorkerType::default_selector_label`] to preserve Dynamo's historical behavior. #[cfg(any(test, feature = "bench"))] - pub fn default(kv_router_config: KvRouterConfig, worker_label: &'static str) -> Self { + pub fn reference(kv_router_config: KvRouterConfig, worker_label: &'static str) -> Self { let picker = DefaultWorkerPicker::new(); Self { - kv_router_config, worker_label, exclusive_affinity: false, - state: WorkerSelectionPolicyState::Default(picker), + state: WorkerSelectionPolicyState::Reference(Box::new(kv_router_config), picker), } } } @@ -154,7 +153,7 @@ fn push_picker_candidate( } } -impl CustomWorkerSelectionState { +impl ComposedPolicyState { // Keep row construction and storage together to avoid passing a full row through a call. #[inline(always)] fn push_candidate(&mut self, candidate: WorkerCandidate) { @@ -225,8 +224,8 @@ impl CustomWorkerSelectionState { } } -pub(super) fn collect_custom_candidates( - state: &mut CustomWorkerSelectionState, +pub(super) fn collect_policy_candidates( + state: &mut ComposedPolicyState, input: &MaterializedSelectionInput<'_>, workers: &HashMap, request: &SchedulingRequest, @@ -319,7 +318,7 @@ pub(super) fn collect_custom_candidates( impl WorkerSelector for WorkerSelectionPolicy { fn uses_exclusive_affinity_target(&self) -> bool { #[cfg(any(test, feature = "bench"))] - if matches!(&self.state, WorkerSelectionPolicyState::Default(_)) { + if matches!(&self.state, WorkerSelectionPolicyState::Reference(..)) { return true; } self.exclusive_affinity @@ -328,8 +327,8 @@ impl WorkerSelector for WorkerSelectionPolicy { fn required_worker_inputs(&self) -> WorkerInputs { match &self.state { #[cfg(any(test, feature = "bench"))] - WorkerSelectionPolicyState::Default(_) => WorkerInputs::CACHE | WorkerInputs::LOAD, - WorkerSelectionPolicyState::Custom(state) => { + WorkerSelectionPolicyState::Reference(..) => WorkerInputs::CACHE | WorkerInputs::LOAD, + WorkerSelectionPolicyState::Composed(state) => { let state = state.borrow(); state.filter_inputs | state.scorer_picker_inputs } @@ -344,15 +343,14 @@ impl WorkerSelector for WorkerSelectionPolicy { let (workers, request, eligibility, block_size) = input.into_configured()?; let state = match &self.state { #[cfg(any(test, feature = "bench"))] - WorkerSelectionPolicyState::Default(picker) => { - WorkerSelectionPolicyStateRef::Default(picker) + WorkerSelectionPolicyState::Reference(config, picker) => { + WorkerSelectionPolicyStateRef::Reference(config, picker) } - WorkerSelectionPolicyState::Custom(state) => { - WorkerSelectionPolicyStateRef::Custom(state) + WorkerSelectionPolicyState::Composed(state) => { + WorkerSelectionPolicyStateRef::Composed(state) } }; select_worker_with_policy( - &self.kv_router_config, self.worker_label, state, workers, @@ -420,7 +418,7 @@ mod tests { 16, )) .unwrap(); - let policy = WorkerSelectionPolicy::default(config, "test"); + let policy = WorkerSelectionPolicy::reference(config, "test"); assert!(uses_exclusive_affinity(&policy)); let actual = policy .select_worker(WorkerSelectionInput::configured( diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index ea1d1d8f7aa9..16a98e354ec7 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -582,9 +582,8 @@ impl WorkerSelector for DefaultWorkerSelector { ) -> Result { let (workers, request, eligibility, block_size) = input.into_configured()?; select_worker_with_policy( - &self.kv_router_config, self.worker_type, - WorkerSelectionPolicyStateRef::Default(&self.picker), + WorkerSelectionPolicyStateRef::Reference(&self.kv_router_config, &self.picker), workers, request, eligibility, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 05aa9e7b88d4..404203b87640 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -276,7 +276,7 @@ pub struct SelectionCore { listens_for_kv_events: bool, indexer_registry: Arc, kv_router_config: crate::config::KvRouterConfig, - worker_selection_policy_factory: Option, + worker_selection_policy_factory: WorkerSelectionPolicyFactory, host: SelectionHost, worker_type: WorkerType, cancel_token: CancellationToken, @@ -339,7 +339,7 @@ impl SelectionCore { indexer_threads, cancel_token, None, - Some(policy_factory), + policy_factory, SelectionHost::default(), WorkerType::Aggregated, true, @@ -360,7 +360,7 @@ impl SelectionCore { indexer_threads: usize, cancel_token: CancellationToken, replica_config: Option, - worker_selection_policy_factory: Option, + worker_selection_policy_factory: WorkerSelectionPolicyFactory, host: SelectionHost, worker_type: WorkerType, signal_indexer_ready: bool, diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index b42c9d119623..06c694845c13 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -40,7 +40,7 @@ fn local_core_with( cancel_token, SelectionCacheConfig::default(), std::sync::Arc::new(|config, role, _| { - crate::WorkerSelectionPolicy::default(config.clone(), role.default_selector_label()) + crate::WorkerSelectionPolicy::reference(config.clone(), role.default_selector_label()) }), ) .expect("valid test config") @@ -63,7 +63,14 @@ fn core_with( 1, CancellationToken::new(), None, - policy_factory, + policy_factory.unwrap_or_else(|| { + Arc::new(|config, role, _| { + crate::WorkerSelectionPolicy::reference( + config.clone(), + role.default_selector_label(), + ) + }) + }), host, worker_type, true, diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 3488ebbfd91f..2423d8845bf7 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -316,21 +316,11 @@ impl SelectionCore { block_size, )) }); - let selector = match self.worker_selection_policy_factory.as_ref() { - Some(factory) => { - factory(&self.kv_router_config, self.worker_type, key.as_ref()) - } - #[cfg(test)] - None => { - WorkerSelectionPolicy::default(self.kv_router_config.clone(), worker_label) - } - #[cfg(not(test))] - None => { - return Err(SelectionError::BadRequest( - "routing host must supply a worker-selection policy factory".into(), - )); - } - }; + let selector = (self.worker_selection_policy_factory)( + &self.kv_router_config, + self.worker_type, + key.as_ref(), + ); let profile = self .kv_router_config .policy_profile(Some(&key.model_name)) diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index e1cb98f3758c..a5ff153f696e 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -219,7 +219,7 @@ mod tests { CancellationToken::new(), SelectionCacheConfig::default(), std::sync::Arc::new(|config, role, _| { - crate::WorkerSelectionPolicy::default( + crate::WorkerSelectionPolicy::reference( config.clone(), role.default_selector_label(), ) diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 0556ac7d8a6b..fedcbcc60eb5 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -158,15 +158,12 @@ impl SelectionServiceBuilder { anyhow::bail!("standalone selection does not support request_classifier plugins"); } let worker_selection_policy_factory = match self.worker_selection_policy_factory { - Some(factory) => Some(factory), + Some(factory) => factory, None => self .plugin_registry - .resolve_for_worker_type(&self.kv_router_config, self.worker_type)?, + .resolve_for_worker_type(&self.kv_router_config, self.worker_type)? + .ok_or(crate::plugins::WorkerSelectionPolicyRegistryError::MissingDefault)?, }; - #[cfg(not(test))] - if worker_selection_policy_factory.is_none() { - return Err(crate::plugins::WorkerSelectionPolicyRegistryError::MissingDefault.into()); - } let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); let indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; let recover_from_peers = !self.indexer_peers.is_empty(); @@ -322,7 +319,7 @@ impl SelectionService { cancel_token.clone(), SelectionCacheConfig::default(), std::sync::Arc::new(|config, role, _| { - crate::WorkerSelectionPolicy::default( + crate::WorkerSelectionPolicy::reference( config.clone(), role.default_selector_label(), ) @@ -532,6 +529,30 @@ mod tests { } } + fn test_registry() -> RouterPluginRegistry { + RouterPluginRegistry::default().with_default_factory(Arc::new(|config, role, _| { + crate::WorkerSelectionPolicy::reference(config.clone(), role.default_selector_label()) + })) + } + + #[tokio::test] + async fn missing_default_is_rejected_at_construction() { + let result = SelectionServiceBuilder::new( + test_config(), + WorkerType::Aggregated, + RouterPluginRegistry::default(), + ) + .build() + .await; + let Err(error) = result else { + panic!("missing default accepted") + }; + assert!(matches!( + error.downcast_ref::(), + Some(crate::plugins::WorkerSelectionPolicyRegistryError::MissingDefault) + )); + } + fn reserve_tcp_port() -> u16 { StdTcpListener::bind("127.0.0.1:0") .unwrap() @@ -610,7 +631,7 @@ worker_selection: match SelectionServiceBuilder::new( test_config(), WorkerType::Aggregated, - RouterPluginRegistry::default(), + test_registry(), ) .indexer_threads(1) .replica_sync(port, Vec::new()) @@ -629,15 +650,12 @@ worker_selection: #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn startup_and_shutdown_release_replica_resources() { let port = reserve_tcp_port(); - let failed = SelectionServiceBuilder::new( - test_config(), - WorkerType::Aggregated, - RouterPluginRegistry::default(), - ) - .indexer_threads(1) - .replica_sync(port, vec!["invalid".to_string()]) - .build() - .await; + let failed = + SelectionServiceBuilder::new(test_config(), WorkerType::Aggregated, test_registry()) + .indexer_threads(1) + .replica_sync(port, vec!["invalid".to_string()]) + .build() + .await; assert!(failed.is_err()); let service = build_on_port(port).await; @@ -684,14 +702,10 @@ worker_selection: let server = tokio::spawn(async move { axum::serve(listener, app).await }); let build = tokio::spawn( - SelectionServiceBuilder::new( - test_config(), - WorkerType::Aggregated, - RouterPluginRegistry::default(), - ) - .indexer_threads(1) - .indexer_peers(vec![peer_url]) - .build(), + SelectionServiceBuilder::new(test_config(), WorkerType::Aggregated, test_registry()) + .indexer_threads(1) + .indexer_peers(vec![peer_url]) + .build(), ); tokio::time::timeout(Duration::from_secs(3), gate.requested.notified()) .await diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index 78b9b62179ac..99bbd8f5a203 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -1783,7 +1783,14 @@ async fn selector_replica_sync_propagates_request_lifecycle() { config_a .service_builder( crate::WorkerType::Aggregated, - WorkerSelectionPolicyRegistry::default(), + WorkerSelectionPolicyRegistry::default().with_default_factory(Arc::new( + |config, role, _| { + crate::WorkerSelectionPolicy::reference( + config.clone(), + role.default_selector_label(), + ) + }, + )), ) .build() .await @@ -1793,7 +1800,14 @@ async fn selector_replica_sync_propagates_request_lifecycle() { SelectionServiceBuilder::new( test_config(), crate::WorkerType::Aggregated, - WorkerSelectionPolicyRegistry::default(), + WorkerSelectionPolicyRegistry::default().with_default_factory(Arc::new( + |config, role, _| { + crate::WorkerSelectionPolicy::reference( + config.clone(), + role.default_selector_label(), + ) + }, + )), ) .indexer_threads(1) .replica_sync(port_b, Vec::new()) diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index ce51c0df343b..0dd7ca110217 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -29,9 +29,7 @@ use crate::{ entrypoint::{self, ChatEngineFactoryCallback, RouterConfig}, http::service::metrics::Metrics, kv_router::plugins::RouterPluginBuilder, - kv_router::{ - EncoderRouter, PrefillRouter, RouterLoadSource, RoutingLoadContext, SelectionPolicySource, - }, + kv_router::{EncoderRouter, PrefillRouter, RouterLoadSource, RoutingLoadContext}, local_model::runtime_config::{ TokenizerBackend, VLLM_INFERENCE_V1_GENERATE_CAPABILITY, VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, @@ -414,7 +412,7 @@ impl ModelWatcher { card.download_config(self.local_model_path.as_deref()) .await?; - validate_policy_worker_role(card, &self.plugins.selection_policy())?; + validate_policy_worker_role(card, &self.plugins)?; // Prepare without exact video routing unless the cohort agreed on a contract. if spec.video_contract.is_none() @@ -1367,13 +1365,9 @@ fn effective_router_config<'a>( /// decode or aggregated, so it requires an explicit `worker_type`. fn validate_policy_worker_role( card: &ModelDeploymentCard, - policy: &SelectionPolicySource, + plugins: &RouterPluginBuilder, ) -> anyhow::Result<()> { - if matches!( - policy, - SelectionPolicySource::Factory(_) | SelectionPolicySource::Prepared(_) - ) && card.worker_type.is_none() - { + if plugins.has_custom_worker_selection() && card.worker_type.is_none() { anyhow::bail!( "custom worker-selection policies require model cards with an explicit worker_type" ); @@ -2959,17 +2953,25 @@ request_classifier: } #[test] - fn custom_selector_requires_explicit_worker_type() { - let registry = SelectionPolicySource::Registry; - let factory = SelectionPolicySource::Factory(Arc::new(|_, _, _| { - unreachable!("role validation never constructs the policy") - })); + fn only_explicit_policies_require_typed_model_cards() { + let unresolved = RouterPluginBuilder::default(); + let resolved = dynamo_custom_policy_builtin::default_registry() + .resolve_plugins(&dynamo_kv_router::KvRouterConfig::default()) + .unwrap(); + assert!(resolved.worker_selection().is_some()); + let builtin = RouterPluginBuilder::new(resolved); + let custom = RouterPluginBuilder::new( + dynamo_kv_router::plugins::RouterPlugins::default().with_worker_selection(Arc::new( + |_, _, _| unreachable!("role validation never constructs the policy"), + )), + ); let mut card = ModelDeploymentCard::with_name_only("model"); - assert!(validate_policy_worker_role(&card, ®istry).is_ok()); - assert!(validate_policy_worker_role(&card, &factory).is_err()); + assert!(validate_policy_worker_role(&card, &unresolved).is_ok()); + assert!(validate_policy_worker_role(&card, &builtin).is_ok()); + assert!(validate_policy_worker_role(&card, &custom).is_err()); card.worker_type = Some(WorkerType::Decode); - assert!(validate_policy_worker_role(&card, &factory).is_ok()); + assert!(validate_policy_worker_role(&card, &custom).is_ok()); } #[test] diff --git a/lib/llm/src/entrypoint/input/http.rs b/lib/llm/src/entrypoint/input/http.rs index d01e28b5dfd4..b0f88f1f2b3f 100644 --- a/lib/llm/src/entrypoint/input/http.rs +++ b/lib/llm/src/entrypoint/input/http.rs @@ -102,7 +102,9 @@ impl HttpFrontend { { anyhow::bail!("request classifiers require --router-mode kv"); } - if !self.plugins.is_empty() && !matches!(&engine_config, EngineConfig::Dynamic { .. }) { + if self.plugins.has_custom_plugins() + && !matches!(&engine_config, EngineConfig::Dynamic { .. }) + { anyhow::bail!("custom router plugins require a dynamic engine"); } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index e2812919a790..0eba202a0505 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -108,27 +108,23 @@ pub enum SelectionPolicySource { } impl SelectionPolicySource { - /// Resolve to the factory the partition will call. `label` is the worker - /// pool name the default policy logs under. + /// Resolve to the factory the partition will call. The label argument is retained for + /// compatibility; the factory receives the typed worker role. pub fn resolve( &self, config: &KvRouterConfig, worker_type: WorkerType, - label: &'static str, + _label: &'static str, ) -> Result { match self { Self::Factory(factory) => Ok(factory.clone()), Self::Prepared(prepared) => Ok(prepared.factory.clone()), - Self::Registry => Ok( - match worker_selection_policy_registry() - .resolve_for_worker_type(config, worker_type)? - { - Some(factory) => factory, - None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { - dynamo_custom_policy_builtin::default_policy(config.clone(), label) - }), - }, - ), + Self::Registry => worker_selection_policy_registry() + .resolve_for_worker_type(config, worker_type)? + .ok_or_else(|| { + dynamo_kv_router::plugins::WorkerSelectionPolicyRegistryError::MissingDefault + .into() + }), } } diff --git a/lib/llm/src/kv_router/plugins.rs b/lib/llm/src/kv_router/plugins.rs index fb1fad7ed02b..08a2c6027b30 100644 --- a/lib/llm/src/kv_router/plugins.rs +++ b/lib/llm/src/kv_router/plugins.rs @@ -52,6 +52,10 @@ impl RouterPluginBuilder { Ok(()) } + pub(crate) fn has_custom_worker_selection(&self) -> bool { + self.plugins.has_custom_worker_selection() + } + pub(crate) fn selection_policy(&self) -> SelectionPolicySource { self.plugins.worker_selection().cloned().map_or( SelectionPolicySource::Registry, diff --git a/lib/router-plugins/builtin/benches/default_lifecycle.rs b/lib/router-plugins/builtin/benches/default_lifecycle.rs index cf7daac88783..91bcbe91447e 100644 --- a/lib/router-plugins/builtin/benches/default_lifecycle.rs +++ b/lib/router-plugins/builtin/benches/default_lifecycle.rs @@ -25,7 +25,7 @@ fn bench(c: &mut Criterion) { .nresamples(1_000); for count in [8, 64, 512] { let reference: WorkerSelectionPolicyFactory = Arc::new(|config, role, _| { - WorkerSelectionPolicy::default(config.clone(), role.default_selector_label()) + WorkerSelectionPolicy::reference(config.clone(), role.default_selector_label()) }); for (label, factory) in [("reference", reference), ("plugin", default_factory())] { let core = SelectionCore::try_new_local( diff --git a/lib/router-plugins/builtin/src/default/mod.rs b/lib/router-plugins/builtin/src/default/mod.rs index 121760a355f3..0657694637e5 100644 --- a/lib/router-plugins/builtin/src/default/mod.rs +++ b/lib/router-plugins/builtin/src/default/mod.rs @@ -8,6 +8,7 @@ mod picker; mod scorer; mod selector; +use parameters::PolicyParameters; pub(super) use parameters::register; pub use selector::DefaultWorkerSelector; @@ -21,31 +22,42 @@ use std::sync::Arc; /// Construct the builtin default from configured policy parameters. /// Per-request score overrides are not used. Request load-tracking remains host-owned. pub fn default_policy(config: KvRouterConfig, worker_label: &'static str) -> WorkerSelectionPolicy { - policy_with_rng(config, worker_label, None, false) + let parameters = PolicyParameters::from(&config); + policy_with_rng(config, parameters, worker_label, None, false) } fn policy_with_rng( config: KvRouterConfig, + parameters: PolicyParameters, worker_label: &'static str, rng: Option>>, plain_decode: bool, ) -> WorkerSelectionPolicy { - let scorer = scorer::build(&config, worker_label, plain_decode); - let picker = picker::DefaultPicker::new(config.router_temperature, rng); + let scorer = scorer::build(¶meters, worker_label, plain_decode); + let picker = picker::DefaultPicker::new(parameters.router_temperature, rng); WorkerSelectionPolicy::new(config, worker_label, vec![scorer], Box::new(picker)) .with_exclusive_affinity(true) } /// Factory installed by routing hosts, including hosts without a custom catalog. pub fn default_factory() -> WorkerSelectionPolicyFactory { - Arc::new(|config, role, _partition| policy_for_role(config.clone(), role)) + Arc::new(|config, role, _partition| { + policy_for_role(config.clone(), role, PolicyParameters::from(config)) + }) } fn policy_for_role( config: KvRouterConfig, role: dynamo_kv_router::WorkerType, + parameters: PolicyParameters, ) -> WorkerSelectionPolicy { let plain_decode = role == dynamo_kv_router::WorkerType::Decode && !config.conditional_disagg_enabled; - policy_with_rng(config, role.default_selector_label(), None, plain_decode) + policy_with_rng( + config, + parameters, + role.default_selector_label(), + None, + plain_decode, + ) } diff --git a/lib/router-plugins/builtin/src/default/parameters.rs b/lib/router-plugins/builtin/src/default/parameters.rs index 5c44d000fff1..070cafc58740 100644 --- a/lib/router-plugins/builtin/src/default/parameters.rs +++ b/lib/router-plugins/builtin/src/default/parameters.rs @@ -24,6 +24,52 @@ struct Parameters { router_temperature: Option, } +/// Only the startup values consumed by the default scorer and picker. +#[derive(Clone, Copy)] +pub(super) struct PolicyParameters { + pub(super) overlap_score_credit: f64, + pub(super) overlap_score_credit_decay: f64, + pub(super) prefill_load_scale: f64, + pub(super) decode_active_request_weight: f64, + pub(super) host_cache_hit_weight: f64, + pub(super) disk_cache_hit_weight: f64, + pub(super) shared_cache_multiplier: f64, + pub(super) router_temperature: f64, +} + +impl From<&KvRouterConfig> for PolicyParameters { + fn from(config: &KvRouterConfig) -> Self { + Parameters::default().resolve(config) + } +} + +impl Parameters { + fn resolve(&self, config: &KvRouterConfig) -> PolicyParameters { + PolicyParameters { + overlap_score_credit: self + .overlap_score_credit + .unwrap_or(config.overlap_score_credit), + overlap_score_credit_decay: self + .overlap_score_credit_decay + .unwrap_or(config.overlap_score_credit_decay), + prefill_load_scale: self.prefill_load_scale.unwrap_or(config.prefill_load_scale), + decode_active_request_weight: self + .decode_active_request_weight + .unwrap_or(config.decode_active_request_weight), + host_cache_hit_weight: self + .host_cache_hit_weight + .unwrap_or(config.host_cache_hit_weight), + disk_cache_hit_weight: self + .disk_cache_hit_weight + .unwrap_or(config.disk_cache_hit_weight), + shared_cache_multiplier: self + .shared_cache_multiplier + .unwrap_or(config.shared_cache_multiplier), + router_temperature: self.router_temperature.unwrap_or(config.router_temperature), + } + } +} + pub(crate) fn register( registry: &mut RouterPluginRegistry, ) -> Result<(), WorkerSelectionPolicyRegistryError> { @@ -58,32 +104,7 @@ pub(crate) fn register( } Ok(Arc::new( move |config: &KvRouterConfig, role, _partition| { - let mut config = config.clone(); - if let Some(value) = parameters.overlap_score_credit { - config.overlap_score_credit = value; - } - if let Some(value) = parameters.overlap_score_credit_decay { - config.overlap_score_credit_decay = value; - } - if let Some(value) = parameters.prefill_load_scale { - config.prefill_load_scale = value; - } - if let Some(value) = parameters.decode_active_request_weight { - config.decode_active_request_weight = value; - } - if let Some(value) = parameters.host_cache_hit_weight { - config.host_cache_hit_weight = value; - } - if let Some(value) = parameters.disk_cache_hit_weight { - config.disk_cache_hit_weight = value; - } - if let Some(value) = parameters.shared_cache_multiplier { - config.shared_cache_multiplier = value; - } - if let Some(value) = parameters.router_temperature { - config.router_temperature = value; - } - policy_for_role(config, role) + policy_for_role(config.clone(), role, parameters.resolve(config)) }, )) }), diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 0dd0b3cef8da..b04f58ac2297 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -3,7 +3,7 @@ //! Cache and load cost calculation. The host owns snapshots and validates finite scores. -use dynamo_kv_router::KvRouterConfig; +use super::parameters::PolicyParameters; use dynamo_kv_router::plugins::worker_selection::{ WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, }; @@ -11,7 +11,7 @@ use dynamo_kv_router::plugins::worker_selection::{ /// Resolve optional terms once so disabled weights add neither branches nor conversions to /// each worker score. The shared-cache specialization also omits range traversal code. pub(super) fn build( - config: &KvRouterConfig, + config: &PolicyParameters, worker_label: &'static str, plain_decode: bool, ) -> Box { @@ -103,7 +103,7 @@ impl DefaultScorer { pub(super) fn new( - config: &KvRouterConfig, + config: &PolicyParameters, worker_label: &'static str, plain_decode: bool, ) -> Self { diff --git a/lib/router-plugins/builtin/src/default/selector.rs b/lib/router-plugins/builtin/src/default/selector.rs index 36bd8e0f677c..96782ae00dae 100644 --- a/lib/router-plugins/builtin/src/default/selector.rs +++ b/lib/router-plugins/builtin/src/default/selector.rs @@ -43,6 +43,7 @@ impl DefaultWorkerSelector { Self { policy: Mutex::new(policy_with_rng( config.clone(), + super::PolicyParameters::from(&config), worker_type, rng.clone(), false, diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index ced2c8edcfbd..84bb30b28e56 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -19,7 +19,7 @@ pub fn default_registry() -> RouterPluginRegistry { use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegistryError}; -/// Register every policy Dynamo ships. +/// Register the named providers Dynamo ships, without changing the host's default factory. /// /// `default` is reserved by the registry for Dynamo's built-in worker selector, so no policy here /// can shadow it. A later catalog that reuses one of these type names fails registration rather @@ -27,7 +27,6 @@ use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegis pub fn register( registry: &mut RouterPluginRegistry, ) -> Result<(), WorkerSelectionPolicyRegistryError> { - *registry = std::mem::take(registry).with_default_factory(default_factory()); default::register(registry)?; two_tier_cost_fn::register(registry) } @@ -54,7 +53,7 @@ mod tests { ..Default::default() }; - let mut registry = RouterPluginRegistry::default(); + let mut registry = default_registry(); register(&mut registry).unwrap(); let resolved = registry.resolve(&config); (config, resolved) diff --git a/lib/router-plugins/builtin/tests/plugin_registry.rs b/lib/router-plugins/builtin/tests/plugin_registry.rs index 2e802e8bef7d..b63d13842bd8 100644 --- a/lib/router-plugins/builtin/tests/plugin_registry.rs +++ b/lib/router-plugins/builtin/tests/plugin_registry.rs @@ -3,7 +3,7 @@ use std::sync::Arc; -use dynamo_kv_router::plugins::RouterPluginRegistry; +use dynamo_kv_router::plugins::RouterPlugins; use dynamo_kv_router::plugins::request_classifier::{ ClassifyFuture, ClassifyRequest, RequestClassifier, }; @@ -34,21 +34,24 @@ fn builtin_default_does_not_opt_into_custom_frontend_restrictions() { let (_file, explicit_default) = config("worker_selection:\n aggregated: default\n"); for config in [KvRouterConfig::default(), explicit_default] { assert!(registry.resolve(&config).unwrap().is_some()); - assert!(registry.resolve_plugins(&config).unwrap().is_empty()); + let plugins = registry.resolve_plugins(&config).unwrap(); + assert!(plugins.worker_selection().is_some()); + assert!(!plugins.has_custom_worker_selection()); + assert!(!plugins.has_custom_plugins()); } } #[test] fn builtin_registration_preserves_classifiers_and_resolves_mixed_pools() { for selection in ["default", "named-default"] { - let mut registry = RouterPluginRegistry::default(); + let mut registry = dynamo_custom_policy_builtin::default_registry(); registry .register_request_classifier( "pass-through", Arc::new(|_| Ok(Arc::new(|_| Box::new(PassThrough)))), ) .unwrap(); - // Installing the default must preserve classifier providers already in the catalog. + // Registering named providers must preserve the fallback and classifier providers. dynamo_custom_policy_builtin::register(&mut registry).unwrap(); let (_file, config) = config(&format!( r#" @@ -64,7 +67,12 @@ worker_selection: let plugins = registry.resolve_plugins(&config).unwrap(); assert!(!plugins.is_empty()); assert!(plugins.request_classifier().is_some()); - assert_eq!(plugins.worker_selection().is_some(), selection != "default"); + assert!(plugins.worker_selection().is_some()); + assert!(plugins.has_custom_plugins()); + assert_eq!( + plugins.has_custom_worker_selection(), + selection != "default" + ); // With only a classifier selected, the host still resolves the builtin policy. let factory = registry.resolve(&config).unwrap().unwrap(); for role in [ @@ -77,3 +85,24 @@ worker_selection: } } } + +#[test] +fn programmatic_factory_is_an_explicit_custom_selection() { + let plugins = RouterPlugins::default() + .with_worker_selection(dynamo_custom_policy_builtin::default_factory()); + assert!(plugins.has_custom_worker_selection()); + assert!(plugins.has_custom_plugins()); +} + +#[test] +fn registering_named_policies_does_not_replace_the_hosts_default() { + let fallback = dynamo_custom_policy_builtin::default_factory(); + let mut registry = dynamo_kv_router::plugins::RouterPluginRegistry::default() + .with_default_factory(fallback.clone()); + dynamo_custom_policy_builtin::register(&mut registry).unwrap(); + let resolved = registry + .resolve(&KvRouterConfig::default()) + .unwrap() + .unwrap(); + assert!(Arc::ptr_eq(&fallback, &resolved)); +} From 79bab85853a0754c097d85028bfd5ae4b8907895 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 18:53:33 +0000 Subject: [PATCH 12/31] bench(router): compare scoring API shapes at scale Signed-off-by: Ishan Dhanani --- .../builtin/benches/default_selection.rs | 79 ++++++++++++++++++- 1 file changed, 77 insertions(+), 2 deletions(-) diff --git a/lib/router-plugins/builtin/benches/default_selection.rs b/lib/router-plugins/builtin/benches/default_selection.rs index 0564effb25a1..db4fb30fe31d 100644 --- a/lib/router-plugins/builtin/benches/default_selection.rs +++ b/lib/router-plugins/builtin/benches/default_selection.rs @@ -54,7 +54,7 @@ fn bench(c: &mut Criterion) { .measurement_time(Duration::from_millis(500)) .sample_size(30) .nresamples(1_000); - for count in [8, 64, 256, 1024] { + for count in [8, 64, 256, 1024, 4096] { let (workers, request) = support::fixture(count, 2048); let config = KvRouterConfig { router_temperature: temperature, @@ -123,5 +123,80 @@ fn bench(c: &mut Criterion) { group.finish(); } } -criterion_group!(benches, bench); + +struct LoadScorer(f64); +impl dynamo_kv_router::WorkerScorer for LoadScorer { + fn required_worker_inputs(&self) -> dynamo_kv_router::WorkerInputs { + dynamo_kv_router::WorkerInputs::LOAD + } + fn score( + &mut self, + _: &dynamo_kv_router::WorkerSelectionContext<'_>, + candidate: &dynamo_kv_router::WorkerCandidate, + ) -> Result { + Ok(candidate.load().unwrap().active_requests() as f64 * self.0) + } +} + +struct MinimumPicker; +impl dynamo_kv_router::WorkerPicker for MinimumPicker { + fn pick( + &mut self, + _: &dynamo_kv_router::WorkerSelectionContext<'_>, + input: dynamo_kv_router::WorkerInputView<'_>, + ) -> Result { + Ok(input + .candidates() + .iter() + .enumerate() + .min_by(|(_, a), (_, b)| a.cost().total_cmp(&b.cost())) + .unwrap() + .0) + } +} + +fn stacked(c: &mut Criterion) { + for scorer_count in [1, 2, 4] { + let mut group = c.benchmark_group(format!("stacked_scoring/s{scorer_count}")); + group + .warm_up_time(Duration::from_millis(200)) + .measurement_time(Duration::from_millis(500)) + .sample_size(30) + .nresamples(1_000); + for count in [8, 64, 1024, 4096] { + let (workers, request) = support::fixture(count, 2048); + let scorers = (0..scorer_count) + .map(|index| { + Box::new(LoadScorer((index + 1) as f64)) + as Box + }) + .collect(); + let policy = dynamo_kv_router::WorkerSelectionPolicy::new( + KvRouterConfig::default(), + "prefill", + scorers, + Box::new(MinimumPicker), + ); + let mut select = || { + black_box( + policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(), + ); + }; + let count_allocs = allocations(&mut select); + eprintln!( + "stacked allocations per 100 warm selections: scorers={scorer_count} workers={count} allocations={count_allocs}" + ); + group.bench_function(BenchmarkId::new("plugin", count), |b| b.iter(&mut select)); + } + group.finish(); + } +} +criterion_group!(benches, bench, stacked); criterion_main!(benches); From 8da3f69d60f61fd414f0c3b6535a8df86fd2f941 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 19:01:27 +0000 Subject: [PATCH 13/31] bench(router): prototype batch scoring with private preparation Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 56 ++++--- .../router/custom-policy-example/README.md | 2 +- .../disagg-filter-score-pick/src/scorer.rs | 46 +++--- .../simple-filter-score-pick/src/scorer.rs | 37 +++-- .../src/scorer/active_requests.rs | 16 +- .../src/scorer/uncached_blocks.rs | 16 +- lib/kv-router/benches/worker_selection.rs | 65 ++++---- lib/kv-router/src/plugins/worker_selection.rs | 26 ++- .../src/scheduling/selector/policy.rs | 65 ++++---- lib/kv-router/src/services/selection/tests.rs | 20 ++- lib/kv-router/tests/plugin_api_compat.rs | 19 ++- .../builtin/benches/default_selection.rs | 10 +- .../builtin/src/default/scorer.rs | 31 +++- .../builtin/tests/default_lifecycle.rs | 18 ++- .../builtin/tests/scorer_prepare.rs | 148 ++++++++++++++---- 15 files changed, 364 insertions(+), 211 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index a28ccbc368ee..43a7e11777f7 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -11,14 +11,13 @@ subtitle: Build Rust filters, scorers, and pickers for the Dynamo frontend or EP ## How It Works -This feature replaces the worker-ranking part of Dynamo's routing pipeline. A `WorkerFilter` can exclude a host-eligible worker, a `WorkerScorer` assigns a finite cost to each remaining worker, Dynamo adds costs from all configured scorers, and a `WorkerPicker` chooses one row from the scored candidates. Filters run in declaration order. Each scorer can then prepare an aggregate from all surviving workers before scoring each worker. Rejecting every worker returns an error. Lower costs rank first by convention, but the picker can implement deterministic selection, sampling, tie-breaking, or policy-local state. Dynamo continues to own discovery, eligibility, score validation, accounting, and reservation. +This feature replaces the worker-ranking part of Dynamo's routing pipeline. A `WorkerFilter` can exclude a host-eligible worker, a `WorkerScorer` assigns a finite cost to each remaining worker, Dynamo adds costs from all configured scorers, and a `WorkerPicker` chooses one row from the scored candidates. Filters run in declaration order. Each scorer receives all surviving workers in one call and writes one cost per worker. Rejecting every worker returns an error. Lower costs rank first by convention, but the picker can implement deterministic selection, sampling, tie-breaking, or policy-local state. Dynamo continues to own discovery, eligibility, score validation, accounting, and reservation. ```mermaid flowchart LR Request["Incoming request"] --> Eligibility["Dynamo eligibility filters"] Eligibility --> Filters["WorkerFilter(s)
keep or reject each worker"] - Filters --> Prepare["WorkerScorer.prepare
all surviving workers"] - Prepare --> Scorers["WorkerScorer.score
one cost per worker"] + Filters --> Scorers["WorkerScorer.score
all surviving workers, one cost per worker"] Scorers --> Sum["Dynamo sums scorer costs"] Sum --> Picker["WorkerPicker
one candidate row"] Picker --> Validate["Dynamo validates the row"] @@ -38,7 +37,7 @@ Write your own policy when no shipped policy expresses the rule you need. The re | Stage | Input | Output | Use this stage for | |---|---|---|---| | `WorkerFilter` | Request context and one host-eligible worker | Keep or reject | The rule is a hard requirement, not a ranking preference | -| `WorkerScorer` | Request context; all surviving workers in `prepare`, one worker in `score` | One finite cost per worker | The rule ranks workers or adds a penalty | +| `WorkerScorer` | Request context, all surviving workers, and an aligned output slice | One finite cost per worker | The rule ranks workers or adds a penalty | | `WorkerPicker` | All eligible workers and their total costs | One row index | The rule samples, breaks ties, tracks policy state, or ignores total cost | | Policy factory | Router configuration, worker type, and routing partition | One `WorkerSelectionPolicy` | Prefill, decode, models, or routing groups need different components | @@ -143,28 +142,21 @@ Pass filters to `WorkerSelectionPolicy::new_with_filters`. If the policy has no ### Score Workers -A scorer expresses a preference without excluding a worker. It returns one finite cost for one worker. Lower total cost is better by convention. +A scorer expresses a preference without excluding a worker. It writes one finite cost for each worker into a host-owned output slice. Lower total cost is better by convention. This scorer first finds the lowest active-request count among surviving workers, then scores each worker by the number of requests above that minimum. Subtracting the same minimum from each cost preserves the ranking: ```rust use dynamo_kv_router::plugins::worker_selection::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, - WorkerSelectionPolicyError, + WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, }; +/// Scores active requests above the least-loaded surviving candidate. #[derive(Default)] struct ActiveRequestsScorer { minimum: usize, } - -impl WorkerScorer for ActiveRequestsScorer { - /// Requests load inputs for the active-request count. - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::LOAD - } - - /// Prepare a request-local minimum after all filters have run. +impl ActiveRequestsScorer { fn prepare( &mut self, _context: &WorkerSelectionContext<'_>, @@ -179,22 +171,34 @@ impl WorkerScorer for ActiveRequestsScorer { } Ok(()) } +} + +impl WorkerScorer for ActiveRequestsScorer { + /// Requests load inputs for the active-request count. + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } - /// Return a relative load cost. Subtracting the minimum preserves the ranking. + /// Prepare the batch minimum and write one relative load cost per worker. fn score( &mut self, - _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - Ok((load.active_requests() - self.minimum) as f64) + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.prepare(context, candidates)?; + for (candidate, cost) in candidates.iter().zip(costs) { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + *cost = (load.active_requests() - self.minimum) as f64; + } + Ok(()) } } ``` -`prepare` is optional and defaults to no work. Dynamo calls it once per scorer after filtering, before any `score` call. It receives a borrowed slice with the same declared inputs and snapshot used for scoring. Recompute request-local aggregates on each call. The slice order is unspecified, and the slice cannot be retained. Dynamo skips preparation when no candidates survive; a preparation error stops selection before picking or booking. +Dynamo calls `score` once per scorer after filtering. The candidate slice borrows a fixed request snapshot; its order is unspecified and it cannot be retained. The output slice has the same length and starts with NaN in every slot. Write every cost; an unwritten slot or a non-finite contribution or total stops selection before picking or booking. Empty sets skip scoring. Any preparation helper belongs to the scorer implementation and runs inside its `score` call. A policy can stack multiple scorers. Dynamo calls them in declaration order and adds their costs. Dynamo rejects a non-finite contribution or total. @@ -618,7 +622,7 @@ Holding a classification future pauses admission before dispatch; it does not in - Return a valid picker row. - Treat candidate order as unspecified. - Request only the signal groups that the component reads. -- Keep blocking I/O and panics out of `keep`, `prepare`, `score`, and `pick`. +- Keep blocking I/O and panics out of `keep`, `score`, and `pick`. - Keep policy state local to the factory-created policy unless cross-partition sharing is a deliberate requirement. - Build the policy against the same Dynamo revision as the frontend or EPP. - If a signal adds work, storage, allocation, or another scan, run the worker-selection benchmark. @@ -651,7 +655,7 @@ The provider also accepts `overlap_score_credit_decay`, `decode_active_request_w `WorkerCacheInput::accounting_cache_estimate()` returns the host's weighted block estimate and rounded token estimate. Missing estimates are zero. These estimates are not clamped to prompt length and do not measure physical GPU residency. The plugin can use them without changing the load that the host records. `WorkerLoadInput::is_available()` distinguishes a missing load observation from an observed idle worker. Both values describe the current selection snapshot. -The default scorer owns the cache and load formula, including preferred-taint costs. Its `prepare` method finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. +The default scorer owns the cache and load formula, including preferred-taint costs. Its private `prepare` helper finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. Direct Rust callers and replay import `DefaultWorkerSelector` from `dynamo_custom_policy_builtin`. Scheduler actors use `default_policy()` directly. The core's old selector is available only as a test and benchmark reference, not in production builds. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index decef548ce5c..8bc3a9892154 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -31,7 +31,7 @@ Preferred routing taints are optional candidate metadata. A filter, scorer, or p | `disagg-filter-score-pick` | Prefill and decode workers each need the complete policy flow | | `simple-stacked-score-pick` | Multiple scorer costs compose before one picker runs | -The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap, computes the least-loaded surviving worker in `WorkerScorer::prepare`, and scores active requests above that minimum. Preparation uses the same request snapshot as scoring and resets the minimum for every selection. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. +The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap, receives all surviving workers in one `score` call, and scores active requests above the least-loaded worker. A private preparation helper computes the minimum from that batch. The scorer writes one finite contribution per row into the host-owned output slice. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. The [`soft-pin-repin` policy](soft-pin-repin/README.md) documents its load threshold, soft-binding behavior, and two-Mocker `A -> B -> B` walkthrough. diff --git a/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs b/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs index ee666d84a31f..fcee27216ac6 100644 --- a/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs +++ b/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs @@ -18,19 +18,23 @@ impl WorkerScorer for PrefillLoadScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - let cache = candidate - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; - let active_prefill_blocks = - load.active_prefill_tokens() as f64 / context.block_size() as f64; - let uncached_request_blocks = - (context.request_blocks() as f64 - cache.device_overlap_blocks()).max(0.0); - Ok(active_prefill_blocks + uncached_request_blocks) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + let cache = candidate + .cache() + .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; + let active_prefill_blocks = + load.active_prefill_tokens() as f64 / context.block_size() as f64; + let uncached_request_blocks = + (context.request_blocks() as f64 - cache.device_overlap_blocks()).max(0.0); + *cost = active_prefill_blocks + uncached_request_blocks; + } + Ok(()) } } @@ -47,11 +51,15 @@ impl WorkerScorer for DecodeLoadScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - Ok(load.decode_cost_blocks()) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + *cost = load.decode_cost_blocks(); + } + Ok(()) } } diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs index 50840db7a340..4b375c9fe871 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs @@ -12,14 +12,7 @@ use dynamo_kv_router::plugins::worker_selection::{ pub(crate) struct ActiveRequestsScorer { minimum: usize, } - -impl WorkerScorer for ActiveRequestsScorer { - /// Requests load inputs for the active-request count. - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::LOAD - } - - /// Prepare a request-local minimum after all filters have run. +impl ActiveRequestsScorer { fn prepare( &mut self, _context: &WorkerSelectionContext<'_>, @@ -34,16 +27,28 @@ impl WorkerScorer for ActiveRequestsScorer { } Ok(()) } +} + +impl WorkerScorer for ActiveRequestsScorer { + /// Requests load inputs for the active-request count. + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } - /// Return a relative load cost. Subtracting the minimum preserves the ranking. + /// Prepare the batch minimum and write one relative load cost per worker. fn score( &mut self, - _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - Ok((load.active_requests() - self.minimum) as f64) + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.prepare(context, candidates)?; + for (candidate, cost) in candidates.iter().zip(costs) { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + *cost = (load.active_requests() - self.minimum) as f64; + } + Ok(()) } } diff --git a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs index e57d17cafa0b..fa47fabaa0f0 100644 --- a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs +++ b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs @@ -17,11 +17,15 @@ impl WorkerScorer for ActiveRequestsScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - Ok(load.active_requests() as f64) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + *cost = load.active_requests() as f64; + } + Ok(()) } } diff --git a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs index 3c0ca010a4f0..c3f2d817bd05 100644 --- a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs +++ b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs @@ -17,11 +17,15 @@ impl WorkerScorer for UncachedBlocksScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let cache = candidate - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; - Ok((context.request_blocks() as f64 - cache.device_overlap_blocks()).max(0.0)) + candidates: &[WorkerCandidate], + 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"))?; + *cost = (context.request_blocks() as f64 - cache.device_overlap_blocks()).max(0.0); + } + Ok(()) } } diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index e753cc6952e5..8e42031302bf 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -70,20 +70,25 @@ impl WorkerScorer for BenchScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let cache = candidate - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load 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; - Ok((uncached_blocks + load_blocks) * candidate.preferred_taint_multiplier().unwrap_or(1.0)) + candidates: &[WorkerCandidate], + 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"))?; + 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) + * candidate.preferred_taint_multiplier().unwrap_or(1.0); + } + Ok(()) } } @@ -99,20 +104,24 @@ impl WorkerScorer for IgnoringPreferenceScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let cache = candidate - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load 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; - Ok(uncached_blocks + load_blocks) + candidates: &[WorkerCandidate], + 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"))?; + 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; + } + Ok(()) } } diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index 3955b936d4a1..c2f67540e0a9 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -128,27 +128,19 @@ pub trait WorkerScorer: Send { WorkerInputs::NONE } - /// Prepare once for this selection using every candidate that survived host eligibility - /// and policy filters. All scorers prepare before any candidate is scored. This is not - /// called for an empty candidate set; returning an error stops selection before picking. + /// Score every candidate that survived host eligibility and policy filters in one call. + /// The borrowed snapshot and its row order remain fixed for this selection. Empty candidate + /// sets skip scoring. Scorers run in declaration order; an error stops before picking. /// - /// The slice borrows the same request snapshot and declared worker inputs used by `score`. - /// Its order is unspecified, and it cannot be retained after this call. Reset request-local - /// aggregates here rather than carrying them between selections. The default does no work. - fn prepare( - &mut self, - _context: &WorkerSelectionContext<'_>, - _candidates: &[WorkerCandidate], - ) -> Result<(), WorkerSelectionPolicyError> { - Ok(()) - } - - /// Return one finite, lower-is-better cost contribution for an eligible worker row. + /// `costs` has the same length as `candidates` and is initialized to NaN before each call. + /// Write one finite, lower-is-better contribution to every slot. The host validates and adds + /// each contribution to that row's total. Neither slice may be retained after this call. fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result; + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError>; } /// Filters run in declaration order for each candidate. Callback order across different diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index d55f7f9ca112..f4c83e13a889 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -42,6 +42,7 @@ pub(super) struct ComposedPolicyState { pub(super) scorer_picker_inputs: WorkerInputs, pub(super) picker_inputs: WorkerInputs, unscored_candidates: Vec, + score_contributions: Vec, pub(super) candidates: Vec, pub(super) cache_inputs: Vec, pub(super) load_inputs: Vec, @@ -102,6 +103,7 @@ impl WorkerSelectionPolicy { scorer_picker_inputs, picker_inputs, unscored_candidates: Vec::new(), + score_contributions: Vec::new(), candidates: Vec::new(), cache_inputs: Vec::new(), load_inputs: Vec::new(), @@ -180,6 +182,7 @@ impl ComposedPolicyState { let Self { scorers, unscored_candidates, + score_contributions, candidates, .. } = self; @@ -187,38 +190,23 @@ impl ComposedPolicyState { return Ok(()); } debug_assert_eq!(unscored_candidates.len(), candidates.len()); - for scorer in scorers.iter_mut() { - scorer.prepare(context, unscored_candidates)?; - } - if let [scorer] = scorers.as_mut_slice() { - // One scorer needs neither a per-row scorer loop nor separate contribution/total - // checks. Keep the initial addition to preserve the sign of a zero total. - for (row, (candidate, scored)) in unscored_candidates.iter().zip(candidates).enumerate() + score_contributions.resize(candidates.len(), f64::NAN); + for (scorer_index, scorer) in scorers.iter_mut().enumerate() { + score_contributions.fill(f64::NAN); + scorer.score(context, unscored_candidates, score_contributions)?; + for (row, (contribution, scored)) in score_contributions + .iter() + .zip(candidates.iter_mut()) + .enumerate() { - let cost = scorer.score(context, candidate)?; - if !cost.is_finite() { - return Err(WorkerSelectionPolicyError::NonFiniteCost { - scorer_index: 0, - row, - } - .into()); - } - scored.cost = 0.0 + cost; - } - return Ok(()); - } - for (row, (candidate, scored)) in unscored_candidates.iter().zip(candidates).enumerate() { - let mut cost = 0.0; - for (scorer_index, scorer) in scorers.iter_mut().enumerate() { - let contribution = scorer.score(context, candidate)?; - cost += contribution; + let cost = scored.cost + contribution; if !contribution.is_finite() || !cost.is_finite() { return Err( WorkerSelectionPolicyError::NonFiniteCost { scorer_index, row }.into(), ); } + scored.cost = cost; } - scored.cost = cost; } Ok(()) } @@ -509,10 +497,14 @@ mod tests { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - assert!(candidate.preferred_taint_multiplier().is_some()); - Ok(0.0) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + assert!(candidate.preferred_taint_multiplier().is_some()); + *cost = 0.0; + } + Ok(()) } } @@ -814,8 +806,9 @@ mod tests { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, - ) -> Result { + _candidates: &[WorkerCandidate], + _costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { unreachable!("rejected worker must not reach scorers") } } @@ -867,9 +860,13 @@ mod tests { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, - ) -> Result { - Ok(0.0) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (_candidate, cost) in candidates.iter().zip(costs) { + *cost = 0.0; + } + Ok(()) } } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index 99bbd8f5a203..d07eece55881 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -51,9 +51,13 @@ impl WorkerScorer for WorkerIdScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - Ok(candidate.worker().worker_id as f64) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + *cost = candidate.worker().worker_id as f64; + } + Ok(()) } } @@ -81,9 +85,13 @@ impl WorkerScorer for NonFiniteScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, - ) -> Result { - Ok(f64::NAN) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (_candidate, cost) in candidates.iter().zip(costs) { + *cost = f64::NAN; + } + Ok(()) } } diff --git a/lib/kv-router/tests/plugin_api_compat.rs b/lib/kv-router/tests/plugin_api_compat.rs index 99e28dae6b3b..fbc64640e451 100644 --- a/lib/kv-router/tests/plugin_api_compat.rs +++ b/lib/kv-router/tests/plugin_api_compat.rs @@ -1,7 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Compile old plugin implementations against the canonical registry and contracts. +//! Compile legacy plugin import paths against the canonical registry and contracts. +//! Scorer implementations use the current batch signature; old scalar signatures must migrate. //! TODO(v1.7): Remove the legacy import coverage when the compatibility exports are removed. use std::sync::Arc; @@ -41,12 +42,16 @@ impl WorkerScorer for LegacyPolicy { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - let _: u64 = context.request_blocks(); - let _: u32 = context.block_size(); - let load = candidate.load().expect("requested load inputs"); - Ok(load.active_requests() as f64 + load.decode_cost_blocks()) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + let _: u64 = context.request_blocks(); + let _: u32 = context.block_size(); + let load = candidate.load().expect("requested load inputs"); + *cost = load.active_requests() as f64 + load.decode_cost_blocks(); + } + Ok(()) } } diff --git a/lib/router-plugins/builtin/benches/default_selection.rs b/lib/router-plugins/builtin/benches/default_selection.rs index db4fb30fe31d..7c512a8a3251 100644 --- a/lib/router-plugins/builtin/benches/default_selection.rs +++ b/lib/router-plugins/builtin/benches/default_selection.rs @@ -132,9 +132,13 @@ impl dynamo_kv_router::WorkerScorer for LoadScorer { fn score( &mut self, _: &dynamo_kv_router::WorkerSelectionContext<'_>, - candidate: &dynamo_kv_router::WorkerCandidate, - ) -> Result { - Ok(candidate.load().unwrap().active_requests() as f64 * self.0) + candidates: &[dynamo_kv_router::WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), dynamo_kv_router::WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + *cost = candidate.load().unwrap().active_requests() as f64 * self.0; + } + Ok(()) } } diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index b04f58ac2297..727deed5065e 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -122,13 +122,9 @@ impl } } -impl WorkerScorer - for DefaultScorer +impl + DefaultScorer { - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT - } - fn prepare( &mut self, context: &WorkerSelectionContext<'_>, @@ -163,7 +159,7 @@ impl WorkerScorer Ok(()) } - fn score( + fn score_worker( &mut self, context: &WorkerSelectionContext<'_>, candidate: &WorkerCandidate, @@ -231,6 +227,27 @@ impl WorkerScorer } } +impl WorkerScorer + for DefaultScorer +{ + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT + } + + fn score( + &mut self, + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.prepare(context, candidates)?; + for (candidate, cost) in candidates.iter().zip(costs) { + *cost = self.score_worker(context, candidate)?; + } + Ok(()) + } +} + #[cfg(test)] mod tests { use super::IntegerDivisor; diff --git a/lib/router-plugins/builtin/tests/default_lifecycle.rs b/lib/router-plugins/builtin/tests/default_lifecycle.rs index 144d076c9abd..2ac05945629f 100644 --- a/lib/router-plugins/builtin/tests/default_lifecycle.rs +++ b/lib/router-plugins/builtin/tests/default_lifecycle.rs @@ -96,7 +96,7 @@ async fn preparation_error_does_not_book_and_allows_retry() { use std::sync::Arc; struct FailOnce(bool); - impl WorkerScorer for FailOnce { + impl FailOnce { fn prepare( &mut self, _: &WorkerSelectionContext<'_>, @@ -107,12 +107,20 @@ async fn preparation_error_does_not_book_and_allows_retry() { } Ok(()) } + } + + impl WorkerScorer for FailOnce { fn score( &mut self, - _: &WorkerSelectionContext<'_>, - _: &WorkerCandidate, - ) -> Result { - Ok(0.0) + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.prepare(context, candidates)?; + for (_candidate, cost) in candidates.iter().zip(costs) { + *cost = 0.0; + } + Ok(()) } } struct First; diff --git a/lib/router-plugins/builtin/tests/scorer_prepare.rs b/lib/router-plugins/builtin/tests/scorer_prepare.rs index 258380a81505..fe91168a1c61 100644 --- a/lib/router-plugins/builtin/tests/scorer_prepare.rs +++ b/lib/router-plugins/builtin/tests/scorer_prepare.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Exercise scorer preparation from an external crate using only public plugin inputs. +//! Exercise batch scoring and private preparation using only public plugin inputs. mod support; use std::sync::{Arc, Mutex}; @@ -25,12 +25,7 @@ struct RelativeLoadScorer { minimum: usize, calls: Arc>>, } - -impl WorkerScorer for RelativeLoadScorer { - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::LOAD - } - +impl RelativeLoadScorer { fn prepare( &mut self, context: &WorkerSelectionContext<'_>, @@ -51,14 +46,25 @@ impl WorkerScorer for RelativeLoadScorer { .unwrap(); Ok(()) } +} + +impl WorkerScorer for RelativeLoadScorer { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } fn score( &mut self, - _: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - self.calls.lock().unwrap().push(Call::Score(self.index)); - Ok((candidate.load().unwrap().active_requests() - self.minimum) as f64) + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.prepare(context, candidates)?; + for (candidate, cost) in candidates.iter().zip(costs) { + self.calls.lock().unwrap().push(Call::Score(self.index)); + *cost = (candidate.load().unwrap().active_requests() - self.minimum) as f64; + } + Ok(()) } } @@ -101,7 +107,7 @@ impl WorkerPicker for InspectCosts { } #[test] -fn prepares_all_scorers_from_surviving_workers_and_resets_each_selection() { +fn scores_batches_from_surviving_workers_and_resets_each_selection() { let (workers, mut request) = fixture(4, 17); request.allowed_worker_ids = Some([1, 2, 3].into_iter().collect()); let calls = Arc::new(Mutex::new(Vec::new())); @@ -147,11 +153,13 @@ fn prepares_all_scorers_from_surviving_workers_and_resets_each_selection() { } else { vec![2, 2, 3, 3] }; - assert_eq!(calls[0], Call::Prepare(0, expected.clone())); - assert_eq!(calls[1], Call::Prepare(1, expected.clone())); assert_eq!(calls.len(), 2 + expected.len() * 2 + 1); - for pair in calls[2..calls.len() - 1].chunks_exact(2) { - assert_eq!(pair, [Call::Score(0), Call::Score(1)]); + for index in 0..2 { + let start = index * (expected.len() + 1); + assert_eq!(calls[start], Call::Prepare(index, expected.clone())); + for call in &calls[start + 1..start + 1 + expected.len()] { + assert_eq!(*call, Call::Score(index)); + } } assert_eq!(calls.last(), Some(&Call::Pick)); calls.clear(); @@ -159,7 +167,7 @@ fn prepares_all_scorers_from_surviving_workers_and_resets_each_selection() { } struct FailPreparation(Arc>); -impl WorkerScorer for FailPreparation { +impl FailPreparation { fn prepare( &mut self, _: &WorkerSelectionContext<'_>, @@ -168,11 +176,16 @@ impl WorkerScorer for FailPreparation { *self.0.lock().unwrap() += 1; Err(WorkerSelectionPolicyError::failed("prepare failed")) } +} + +impl WorkerScorer for FailPreparation { fn score( &mut self, - _: &WorkerSelectionContext<'_>, - _: &WorkerCandidate, - ) -> Result { + context: &WorkerSelectionContext<'_>, + candidates: &[WorkerCandidate], + _costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.prepare(context, candidates)?; panic!("preparation failure must stop scoring") } } @@ -232,9 +245,13 @@ fn rejects_nonfinite_contributions_and_overflow_before_picking() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - _: &WorkerCandidate, - ) -> Result { - Ok(self.0) + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (_candidate, cost) in candidates.iter().zip(costs) { + *cost = self.0; + } + Ok(()) } } let (workers, request) = fixture(1, 17); @@ -274,13 +291,17 @@ fn picker_columns_and_costs_stay_aligned_after_a_scoring_error() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, - ) -> Result { - self.0 += 1; - if self.0 == 2 { - return Err(WorkerSelectionPolicyError::failed("score failed")); + candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + self.0 += 1; + if self.0 == 2 { + return Err(WorkerSelectionPolicyError::failed("score failed")); + } + *cost = candidate.load().unwrap().active_requests() as f64; } - Ok(candidate.load().unwrap().active_requests() as f64) + Ok(()) } } struct CheckColumns; @@ -334,3 +355,70 @@ fn picker_columns_and_costs_stay_aligned_after_a_scoring_error() { request.allowed_worker_ids = None; assert_ne!(select(&request).unwrap().worker.worker_id, 1); } + +#[test] +fn unwritten_costs_cannot_reuse_a_previous_selection_or_scorer() { + struct Fill; + impl WorkerScorer for Fill { + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + costs.fill(7.0); + Ok(()) + } + } + struct OmitOnSecondCall(bool); + impl WorkerScorer for OmitOnSecondCall { + fn score( + &mut self, + _: &WorkerSelectionContext<'_>, + _: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + if std::mem::replace(&mut self.0, false) { + costs.fill(1.0); + } + Ok(()) + } + } + struct First; + impl WorkerPicker for First { + fn pick( + &mut self, + _: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + assert!(input.candidates().iter().all(|c| c.cost() == 8.0)); + Ok(0) + } + } + let (workers, request) = fixture(2, 17); + let policy = WorkerSelectionPolicy::new( + KvRouterConfig::default(), + "test", + vec![Box::new(Fill), Box::new(OmitOnSecondCall(true))], + Box::new(First), + ); + let select = || { + policy.select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + }; + select().unwrap(); + let error = select().unwrap_err(); + assert!(matches!( + error, + dynamo_kv_router::KvSchedulerError::WorkerSelectionPolicy( + WorkerSelectionPolicyError::NonFiniteCost { + scorer_index: 1, + row: 0 + } + ) + )); +} From c36b7aadab88bfef1575c13de84f8b6522077f90 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 19:54:49 +0000 Subject: [PATCH 14/31] fix(router): gate request cache context by component inputs Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 8 +- .../router/custom-policy-example/README.md | 2 + lib/kv-router/src/plugins/worker_selection.rs | 33 +++- lib/kv-router/src/scheduling/selector/mod.rs | 4 +- .../src/scheduling/selector/policy.rs | 37 +++-- .../builtin/src/default/scorer.rs | 8 +- .../builtin/tests/cache_context.rs | 149 ++++++++++++++++++ .../builtin/tests/default_policy.rs | 2 +- 8 files changed, 219 insertions(+), 24 deletions(-) create mode 100644 lib/router-plugins/builtin/tests/cache_context.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 43a7e11777f7..04ede4df70a3 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -431,7 +431,7 @@ The disaggregated prefill hop and query-only probes do not run classifiers. Stat ## Available Signals -Request context and worker identity are always available. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. +Worker identity and ordinary request metadata are always available. Request-level cache context requires `WorkerInputs::CACHE`. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. ### Request Context @@ -439,8 +439,8 @@ Request context and worker identity are always available. Dynamo calculates opti |---|---| | `prompt_tokens()` | Exact incoming prompt token count, borrowed from the request without allocation | | `request_blocks()` | Incoming prompt size in KV blocks, rounded up | -| `has_tier_matches()` | Whether the current request lookup contains any tier-specific matches, before policy filters | -| `shared_cache_hits()` | Optional unweighted shared-cache ranges from the current lookup | +| `has_tier_matches()` | With `CACHE`, returns `Some(true)` if the lookup contains any tier-specific matches before filtering, or `Some(false)` otherwise; without `CACHE`, returns `None` | +| `shared_cache_hits()` | With `CACHE`, returns the optional unweighted shared-cache ranges from the current lookup; without `CACHE`, returns `None` | | `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 | @@ -450,6 +450,8 @@ Request context and worker identity are always available. Dynamo calculates opti | `strict_priority()` | Strict integer priority. The queue orders larger values first | | `router_temperature_override()` | Optional per-request router temperature override | +Each filter, scorer, and picker must declare `WorkerInputs::CACHE` in its own `required_worker_inputs()` to read the two cache context accessors. Another component requesting cache does not enable them. The host borrows the existing lookup snapshot; reading these accessors adds no cache lookup or range copy. + ### Session Context `session_context()` returns `None` when the request has no session metadata. This policy-facing view contains selected session metadata; it is not Dynamo's internal request envelope. When present, it provides: diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 8bc3a9892154..052a46d3530e 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -377,4 +377,6 @@ Send the same `curl` request from a fourth terminal. The frontend log records se Filters, scorers, and pickers can read `context.prompt_tokens()` for the exact input token count. This returns the existing request value without a prompt copy or allocation. Use `context.request_blocks()` when the policy needs rounded KV blocks instead. For example, at 16 tokens per block, a 17-token prompt has 2 blocks. +Request-level cache information also requires `WorkerInputs::CACHE`: `context.has_tier_matches()` returns `None` without that declaration and `Some(false)` when declared but no tier matches were found. `context.shared_cache_hits()` returns `None` when cache was not requested or the request has no shared-cache ranges. Each filter, scorer, and picker declares its own inputs; another component requesting cache does not grant access. + The builtin default now uses the public policy API too. Rust hosts install `dynamo_custom_policy_builtin::default_registry()` before adding custom providers. This supplies default selection for roles that the custom configuration does not replace. diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index c2f67540e0a9..471387c36298 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -37,6 +37,7 @@ pub struct WorkerSelectionContext<'a> { pub(crate) block_size: u32, pub(crate) track_prefill_tokens: bool, pub(crate) has_tier_matches: bool, + pub(crate) inputs: WorkerInputs, pub(crate) pinned_worker: Option, pub(crate) router_temperature_override: Option, } @@ -65,7 +66,7 @@ pub struct WorkerInputs(u8); impl WorkerInputs { /// Request no optional worker inputs. pub const NONE: Self = Self(0); - /// Request KV-cache overlap inputs. + /// Request worker KV-cache overlap inputs and request-level cache context. pub const CACHE: Self = Self(1 << 0); /// Request active-load inputs. pub const LOAD: Self = Self(1 << 1); @@ -191,17 +192,39 @@ impl WorkerSelectionContext<'_> { } /// Shared-cache ranges from this request's lookup snapshot, if present. + /// Returns None unless this component declared [`WorkerInputs::CACHE`]. /// The ranges are unweighted block offsets. The host owns their lifetime and /// policy inspection does not perform a lookup or change accounting. pub fn shared_cache_hits(&self) -> Option<&crate::SharedCacheHits> { - self.request.shared_cache_hits.as_ref() + if self.inputs.contains(WorkerInputs::CACHE) { + self.request.shared_cache_hits.as_ref() + } else { + None + } } /// Whether this request has any tier-specific cache matches, before worker filtering. - /// False means the host only supplied its accounting estimate (or no cache data). + /// Returns None unless this component declared [`WorkerInputs::CACHE`]. Some(false) + /// means the host only supplied its accounting estimate (or no cache data). /// The value describes the current lookup snapshot, not worker cache capacity. - pub fn has_tier_matches(&self) -> bool { - self.has_tier_matches + pub fn has_tier_matches(&self) -> Option { + self.inputs + .contains(WorkerInputs::CACHE) + .then_some(self.has_tier_matches) + } + + /// Restrict request-level signals to this component's startup input declaration. + pub(crate) fn with_inputs(&self, inputs: WorkerInputs) -> Self { + Self { + request: self.request, + request_blocks: self.request_blocks, + block_size: self.block_size, + track_prefill_tokens: self.track_prefill_tokens, + has_tier_matches: self.has_tier_matches, + inputs, + pinned_worker: self.pinned_worker, + router_temperature_override: self.router_temperature_override, + } } /// Return the incoming prompt size in KV blocks. diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 7a893913dfa7..811102e381b3 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -140,6 +140,7 @@ impl<'a> MaterializedSelectionInput<'a> { request_blocks: request.request_blocks(block_size), block_size, track_prefill_tokens: request.track_prefill_tokens, + inputs: WorkerInputs::NONE, pinned_worker: request.pinned_worker, has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() @@ -407,7 +408,8 @@ fn select_worker_with_policy( .contains(WorkerInputs::LOAD) .then_some(load_inputs.as_slice()), }; - let row = picker.pick(&input.context, picker_input)?; + let context = input.context.with_inputs(*picker_inputs); + let row = picker.pick(&context, picker_input)?; let Some(candidate) = candidates.get(row) else { return Err(WorkerSelectionPolicyError::InvalidPickerRow { row, diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index f4c83e13a889..543c9b90954b 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -35,8 +35,8 @@ pub(super) enum WorkerSelectionPolicyStateRef<'a> { } pub(super) struct ComposedPolicyState { - pub(super) filters: Vec>, - pub(super) scorers: Vec>, + pub(super) filters: Vec<(WorkerInputs, Box)>, + pub(super) scorers: Vec<(WorkerInputs, Box)>, pub(super) picker: Box, pub(super) filter_inputs: WorkerInputs, pub(super) scorer_picker_inputs: WorkerInputs, @@ -86,12 +86,23 @@ impl WorkerSelectionPolicy { picker: Box, ) -> Self { let picker_inputs = picker.required_worker_inputs(); - let filter_inputs = filters.iter().fold(WorkerInputs::NONE, |inputs, filter| { - inputs | filter.required_worker_inputs() - }); - let scorer_picker_inputs = scorers.iter().fold(picker_inputs, |inputs, scorer| { - inputs | scorer.required_worker_inputs() - }); + // Freeze each declaration once; callbacks must not inherit another component's access. + let filters: Vec<_> = filters + .into_iter() + .map(|filter| (filter.required_worker_inputs(), filter)) + .collect(); + let scorers: Vec<_> = scorers + .into_iter() + .map(|scorer| (scorer.required_worker_inputs(), scorer)) + .collect(); + let filter_inputs = filters + .iter() + .fold(WorkerInputs::NONE, |inputs, (required, _)| { + inputs | *required + }); + let scorer_picker_inputs = scorers + .iter() + .fold(picker_inputs, |inputs, (required, _)| inputs | *required); Self { worker_label, exclusive_affinity: false, @@ -191,9 +202,10 @@ impl ComposedPolicyState { } debug_assert_eq!(unscored_candidates.len(), candidates.len()); score_contributions.resize(candidates.len(), f64::NAN); - for (scorer_index, scorer) in scorers.iter_mut().enumerate() { + for (scorer_index, (inputs, scorer)) in scorers.iter_mut().enumerate() { score_contributions.fill(f64::NAN); - scorer.score(context, unscored_candidates, score_contributions)?; + let context = context.with_inputs(*inputs); + scorer.score(&context, unscored_candidates, score_contributions)?; for (row, (contribution, scored)) in score_contributions .iter() .zip(candidates.iter_mut()) @@ -269,8 +281,9 @@ pub(super) fn collect_policy_candidates( filter_preferred_taint_multiplier, state.filter_inputs, ); - for filter in &mut state.filters { - match filter.keep(&input.context, &filter_candidate) { + for (inputs, filter) in &mut state.filters { + let context = input.context.with_inputs(*inputs); + match filter.keep(&context, &filter_candidate) { Ok(true) => {} Ok(false) => return false, Err(policy_error) => { diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 727deed5065e..4a802832d1f1 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -60,6 +60,7 @@ struct DefaultScorer { #[derive(Default)] struct PreparedRequest { min_prefill: usize, + has_tier_matches: bool, block_size: IntegerDivisor, request_blocks: IntegerDivisor, overlap_credit: f64, @@ -138,6 +139,9 @@ impl }; self.prepared = PreparedRequest { min_prefill: 0, + has_tier_matches: context + .has_tier_matches() + .ok_or_else(|| WorkerSelectionPolicyError::failed("cache context unavailable"))?, block_size: IntegerDivisor::new(u64::from(context.block_size())), request_blocks: IntegerDivisor::new(context.request_blocks()), overlap_credit, @@ -171,13 +175,13 @@ impl .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); - let device = if context.has_tier_matches() { + let device = if self.prepared.has_tier_matches { cache.device_overlap_blocks() } else { estimated_overlap }; let shared_credit = if SHARED_CREDIT { - let shared = if context.has_tier_matches() { + let shared = if self.prepared.has_tier_matches { cache.shared_beyond_device_blocks() } else { context diff --git a/lib/router-plugins/builtin/tests/cache_context.rs b/lib/router-plugins/builtin/tests/cache_context.rs new file mode 100644 index 000000000000..a3ba47ef33af --- /dev/null +++ b/lib/router-plugins/builtin/tests/cache_context.rs @@ -0,0 +1,149 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +mod support; + +use dynamo_kv_router::{ + KvRouterConfig, SharedCacheHits, WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, + WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, + WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, +}; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +struct CacheProbe { + inputs: WorkerInputs, + tier_matches: bool, + shared_hits: bool, + declarations: Arc, + calls: Arc, +} + +impl CacheProbe { + fn declare(&self) -> WorkerInputs { + self.declarations.fetch_add(1, Ordering::Relaxed); + self.inputs + } + + fn check(&self, context: &WorkerSelectionContext<'_>) { + self.calls.fetch_add(1, Ordering::Relaxed); + // Ordinary request context stays available without CACHE. + assert_eq!(context.prompt_tokens(), 17); + if self.inputs.contains(WorkerInputs::CACHE) { + assert_eq!(context.has_tier_matches(), Some(self.tier_matches)); + assert_eq!(context.shared_cache_hits().is_some(), self.shared_hits); + if let Some(hits) = context.shared_cache_hits() { + assert_eq!(hits.hits_beyond(0), 4); + } + } else { + assert_eq!(context.has_tier_matches(), None); + assert!(context.shared_cache_hits().is_none()); + } + } +} + +impl WorkerFilter for CacheProbe { + fn required_worker_inputs(&self) -> WorkerInputs { + self.declare() + } + + fn keep( + &mut self, + context: &WorkerSelectionContext<'_>, + _candidate: &WorkerCandidate, + ) -> Result { + self.check(context); + Ok(true) + } +} + +impl WorkerScorer for CacheProbe { + fn required_worker_inputs(&self) -> WorkerInputs { + self.declare() + } + + fn score( + &mut self, + context: &WorkerSelectionContext<'_>, + _candidates: &[WorkerCandidate], + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + self.check(context); + costs.fill(0.0); + Ok(()) + } +} + +impl WorkerPicker for CacheProbe { + fn required_worker_inputs(&self) -> WorkerInputs { + self.declare() + } + + fn pick( + &mut self, + context: &WorkerSelectionContext<'_>, + _input: WorkerInputView<'_>, + ) -> Result { + self.check(context); + Ok(0) + } +} + +#[test] +fn cache_context_requires_each_components_own_declaration() { + for mask in 0..8 { + for unrequested in [WorkerInputs::NONE, WorkerInputs::LOAD] { + for tier_matches in [false, true] { + for shared_hits in [false, true] { + let (workers, mut request) = support::fixture(2, 17); + if !tier_matches { + request.overlap.tier_overlap_blocks = Default::default(); + } + if shared_hits { + request.shared_cache_hits = + Some(SharedCacheHits::from_ranges(vec![1..3, 5..7])); + } + let declarations = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let probe = |inputs| CacheProbe { + inputs, + tier_matches, + shared_hits, + declarations: declarations.clone(), + calls: calls.clone(), + }; + let inputs = |bit| { + if mask & bit == 0 { + unrequested + } else { + WorkerInputs::CACHE + } + }; + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + vec![Box::new(probe(inputs(1))), Box::new(probe(unrequested))], + vec![Box::new(probe(inputs(2))), Box::new(probe(unrequested))], + Box::new(probe(inputs(4))), + ); + assert_eq!(declarations.load(Ordering::Relaxed), 5); + for _ in 0..2 { + policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + } + // Four worker/ranks pass two filters, then two scorers and one picker run. + assert_eq!(calls.load(Ordering::Relaxed), 22); + assert_eq!(declarations.load(Ordering::Relaxed), 5); + } + } + } + } +} diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 4152874aa903..d76b9ab0fbb9 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -172,7 +172,7 @@ fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { ) -> Result { assert_eq!(context.prompt_tokens(), 17); assert_eq!(context.request_blocks(), 2); - assert!(context.has_tier_matches()); + assert_eq!(context.has_tier_matches(), Some(true)); assert!(input.load().unwrap().iter().all(|load| load.is_available())); for cache in input.cache().unwrap() { let (blocks, tokens) = cache.accounting_cache_estimate(); From 0e37754a366f1208ab94895415ee4054e8e11e88 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 20:09:15 +0000 Subject: [PATCH 15/31] perf(router): reuse cache-gated filter context across workers Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/selector/policy.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 543c9b90954b..ab148cbb5cc8 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -265,6 +265,8 @@ pub(super) fn collect_policy_candidates( && state.filter_inputs.contains(WorkerInputs::PREFERRED_TAINT); let materialize_additional_preferred_taint = eligibility.pinned_worker().is_none() && additional_inputs.contains(WorkerInputs::PREFERRED_TAINT); + // Reuse the borrowed context across rows; only the current filter's access mask changes. + let mut filter_context = input.context.with_inputs(WorkerInputs::NONE); let mut has_eligible_worker = false; let mut error = None; eligibility.any_eligible_worker_rank(workers, |worker, config| { @@ -282,8 +284,8 @@ pub(super) fn collect_policy_candidates( state.filter_inputs, ); for (inputs, filter) in &mut state.filters { - let context = input.context.with_inputs(*inputs); - match filter.keep(&context, &filter_candidate) { + filter_context.inputs = *inputs; + match filter.keep(&filter_context, &filter_candidate) { Ok(true) => {} Ok(false) => return false, Err(policy_error) => { From 1dca8375be40936984a56becdccaca7841749cf3 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 21:16:10 +0000 Subject: [PATCH 16/31] refactor(router): scope plugin inputs through borrowed views Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 33 +- .../router/custom-policy-example/README.md | 6 +- .../disagg-filter-score-pick/src/filter.rs | 2 +- .../disagg-filter-score-pick/src/scorer.rs | 7 +- .../simple-filter-score-pick/src/filter.rs | 2 +- .../simple-filter-score-pick/src/scorer.rs | 9 +- .../src/scorer/active_requests.rs | 5 +- .../src/scorer/uncached_blocks.rs | 5 +- lib/kv-router/benches/worker_selection.rs | 10 +- lib/kv-router/src/lib.rs | 6 +- lib/kv-router/src/plugins/worker_selection.rs | 363 +----------------- .../src/plugins/worker_selection/context.rs | 132 +++++++ .../src/plugins/worker_selection/inputs.rs | 283 ++++++++++++++ lib/kv-router/src/scheduling/selector/mod.rs | 23 +- .../src/scheduling/selector/policy.rs | 38 +- .../src/scheduling/selector/reference.rs | 34 +- .../src/services/selection/core/tests.rs | 9 +- lib/kv-router/src/services/selection/tests.rs | 10 +- lib/kv-router/tests/plugin_api_compat.rs | 10 +- lib/llm/src/kv_router.rs | 14 +- .../builtin/benches/default_selection.rs | 60 ++- .../builtin/src/default/scorer.rs | 29 +- .../builtin/tests/cache_context.rs | 207 +++++++--- .../builtin/tests/default_lifecycle.rs | 6 +- .../builtin/tests/default_policy.rs | 2 +- .../builtin/tests/scorer_prepare.rs | 24 +- 26 files changed, 800 insertions(+), 529 deletions(-) create mode 100644 lib/kv-router/src/plugins/worker_selection/context.rs create mode 100644 lib/kv-router/src/plugins/worker_selection/inputs.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 04ede4df70a3..b711ba2f3528 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -126,7 +126,7 @@ impl WorkerFilter for MinimumDeviceOverlapFilter { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { let cache = candidate .cache() @@ -148,7 +148,7 @@ This scorer first finds the lowest active-request count among surviving workers, ```rust use dynamo_kv_router::plugins::worker_selection::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, + WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, }; /// Scores active requests above the least-loaded surviving candidate. @@ -160,7 +160,7 @@ impl ActiveRequestsScorer { fn prepare( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { self.minimum = usize::MAX; for candidate in candidates { @@ -183,7 +183,7 @@ impl WorkerScorer for ActiveRequestsScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.prepare(context, candidates)?; @@ -198,7 +198,7 @@ impl WorkerScorer for ActiveRequestsScorer { } ``` -Dynamo calls `score` once per scorer after filtering. The candidate slice borrows a fixed request snapshot; its order is unspecified and it cannot be retained. The output slice has the same length and starts with NaN in every slot. Write every cost; an unwritten slot or a non-finite contribution or total stops selection before picking or booking. Empty sets skip scoring. Any preparation helper belongs to the scorer implementation and runs inside its `score` call. +Dynamo calls `score` once per scorer after filtering. The `WorkerCandidates` view borrows a fixed request snapshot; its order is unspecified and it cannot be retained. Use `iter()` to visit workers, `get(row)` to read one worker, and `len()` or `is_empty()` to inspect the batch. Each returned `WorkerCandidate` borrows the stored row and exposes only that scorer's declared inputs. Creating the view does not copy the worker table. The output slice has the same length and starts with NaN in every slot. Write every cost; an unwritten slot or a non-finite contribution or total stops selection before picking or booking. Empty sets skip scoring. Any preparation helper belongs to the scorer implementation and runs inside its `score` call. A policy can stack multiple scorers. Dynamo calls them in declaration order and adds their costs. Dynamo rejects a non-finite contribution or total. @@ -439,8 +439,7 @@ Worker identity and ordinary request metadata are always available. Request-leve |---|---| | `prompt_tokens()` | Exact incoming prompt token count, borrowed from the request without allocation | | `request_blocks()` | Incoming prompt size in KV blocks, rounded up | -| `has_tier_matches()` | With `CACHE`, returns `Some(true)` if the lookup contains any tier-specific matches before filtering, or `Some(false)` otherwise; without `CACHE`, returns `None` | -| `shared_cache_hits()` | With `CACHE`, returns the optional unweighted shared-cache ranges from the current lookup; without `CACHE`, returns `None` | +| `cache()` | With `CACHE`, returns a borrowed `RequestCacheInput`; otherwise returns `None` | | `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 | @@ -450,7 +449,16 @@ Worker identity and ordinary request metadata are always available. Request-leve | `strict_priority()` | Strict integer priority. The queue orders larger values first | | `router_temperature_override()` | Optional per-request router temperature override | -Each filter, scorer, and picker must declare `WorkerInputs::CACHE` in its own `required_worker_inputs()` to read the two cache context accessors. Another component requesting cache does not enable them. The host borrows the existing lookup snapshot; reading these accessors adds no cache lookup or range copy. +Each filter, scorer, and picker must declare `WorkerInputs::CACHE` in its own `required_worker_inputs()` to access `context.cache()`. Another component requesting cache does not enable it. The returned view exposes: + +| Accessor | Meaning | +|---|---| +| `has_tier_matches()` | Whether the lookup contains any GPU, CPU, or disk matches before filtering; `false` means only accounting estimates or no cache data were supplied | +| `shared_hits()` | Optional unweighted shared-cache ranges from the current lookup, with `total_hits` and `hits_beyond(prefix)` | + +A declared cache view exists even when no shared-cache result was supplied; in that case, `shared_hits()` returns `None`. The host borrows the existing lookup snapshot. Reading the view adds no cache lookup, range copy, or allocation. + +Calculate shared-cache credit inside the policy from `shared_hits()` and the worker's cache data. For example, call `hits_beyond(device_prefix)` to count shared blocks beyond the GPU prefix. The default policy uses its accounting estimate as the prefix when tier matches are absent. It performs this calculation only when shared-cache credit is enabled. ### Session Context @@ -500,17 +508,18 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `CACHE` | `device_overlap_blocks()` | Device-resident prefix overlap in blocks | | `CACHE` | `host_overlap_blocks()` | Host-pinned prefix overlap in blocks | | `CACHE` | `disk_overlap_blocks()` | Disk prefix overlap in blocks | -| `CACHE` | `shared_beyond_device_blocks()` | Shared-cache hits beyond the device-resident prefix | +| `CACHE` | `accounting_cache_estimate()` | Host accounting estimate as weighted KV blocks and rounded tokens; not clamped to prompt length | +| `LOAD` | `is_available()` | Whether the host supplied a load observation; distinguishes missing data from an observed idle worker | | `LOAD` | `active_prefill_tokens()` | Tokens currently active in the worker's prefill stage | | `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 | -The cache 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. +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 requested groups through `WorkerCandidate::cache()` or `load()`. A picker reads index-aligned arrays through `WorkerInputView`. Each component must request every optional group that it reads. +A filter or scorer reads requested groups through `WorkerCandidate::cache()` or `load()`. A picker reads index-aligned arrays through `WorkerInputView`. Each component must request every optional group that it reads. Candidate and context views enforce that declaration independently for every filter, scorer, and picker, including preferred-taint access. The host materializes the combined requirements once and shares borrowed views; it does not copy or rewrite every row for each scorer. -`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, Dynamo does not materialize the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. +`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. ## Link the Policy Into Dynamo diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 052a46d3530e..55de546142c0 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -20,7 +20,7 @@ policy crate -> catalog crate -> router-policy YAML -> frontend or EPP binary Dynamo owns discovery, eligibility, queueing, validation, reservations, accounting, and metrics. A policy sees only eligible workers and returns one candidate row. -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, Dynamo does not materialize the multiplier. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules. +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. ## Pick a Starting Point @@ -377,6 +377,8 @@ Send the same `curl` request from a fourth terminal. The frontend log records se Filters, scorers, and pickers can read `context.prompt_tokens()` for the exact input token count. This returns the existing request value without a prompt copy or allocation. Use `context.request_blocks()` when the policy needs rounded KV blocks instead. For example, at 16 tokens per block, a 17-token prompt has 2 blocks. -Request-level cache information also requires `WorkerInputs::CACHE`: `context.has_tier_matches()` returns `None` without that declaration and `Some(false)` when declared but no tier matches were found. `context.shared_cache_hits()` returns `None` when cache was not requested or the request has no shared-cache ranges. Each filter, scorer, and picker declares its own inputs; another component requesting cache does not grant access. +Request-level cache information also requires `WorkerInputs::CACHE`: `context.cache()` returns `None` without that declaration. With it, the borrowed view exposes `has_tier_matches()` and `shared_hits()`. The former is false when the snapshot has no tier matches; the latter is `None` when no shared-cache result was supplied. A policy can call `shared_hits().map(|hits| hits.hits_beyond(prefix))` to calculate shared hits after its chosen prefix, only when needed. + +Scorers receive a `WorkerCandidates` view with `iter()`, `get(row)`, `len()`, and `is_empty()`. Filters receive one borrowed `WorkerCandidate`. Both expose `cache()`, `load()`, and `preferred_taint_multiplier()` only when that component declares the corresponding `CACHE`, `LOAD`, or `PREFERRED_TAINT` group. The same rule applies to picker inputs. Another component requesting a group does not grant access. Views share the host's stored rows without allocating or copying a candidate table. The builtin default now uses the public policy API too. Rust hosts install `dynamo_custom_policy_builtin::default_registry()` before adding custom providers. This supplies default selection for roles that the custom configuration does not replace. diff --git a/examples/router/custom-policy-example/disagg-filter-score-pick/src/filter.rs b/examples/router/custom-policy-example/disagg-filter-score-pick/src/filter.rs index b2e83f2ad7d9..313b4d1dca76 100644 --- a/examples/router/custom-policy-example/disagg-filter-score-pick/src/filter.rs +++ b/examples/router/custom-policy-example/disagg-filter-score-pick/src/filter.rs @@ -19,7 +19,7 @@ impl WorkerFilter for MinimumDeviceOverlapFilter { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { let cache = candidate .cache() diff --git a/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs b/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs index fcee27216ac6..52a7f703a0b5 100644 --- a/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs +++ b/examples/router/custom-policy-example/disagg-filter-score-pick/src/scorer.rs @@ -4,7 +4,8 @@ //! Load scorers for prefill and decode routing partitions. use dynamo_kv_router::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, + WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicyError, }; /// Scores prefill workers by active work and uncached request blocks. @@ -18,7 +19,7 @@ impl WorkerScorer for PrefillLoadScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { @@ -51,7 +52,7 @@ impl WorkerScorer for DecodeLoadScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/filter.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/filter.rs index 6fb790a8c5ec..78fd732503e9 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/filter.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/filter.rs @@ -22,7 +22,7 @@ impl WorkerFilter for MinimumDeviceOverlapFilter { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { let cache = candidate .cache() diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs index 4b375c9fe871..a7f4a551cc6b 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs @@ -4,7 +4,8 @@ //! Active-request scorer for the `simple-filter-score-pick` policy. use dynamo_kv_router::plugins::worker_selection::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, + WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicyError, }; /// Scores active requests above the least-loaded surviving candidate. @@ -16,10 +17,10 @@ impl ActiveRequestsScorer { fn prepare( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { self.minimum = usize::MAX; - for candidate in candidates { + for candidate in candidates.iter() { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; @@ -39,7 +40,7 @@ impl WorkerScorer for ActiveRequestsScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.prepare(context, candidates)?; diff --git a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs index fa47fabaa0f0..32f0dc766ace 100644 --- a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs +++ b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/active_requests.rs @@ -4,7 +4,8 @@ //! Active-request cost for the stacked scorer stage. use dynamo_kv_router::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, + WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicyError, }; pub(crate) struct ActiveRequestsScorer; @@ -17,7 +18,7 @@ impl WorkerScorer for ActiveRequestsScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { diff --git a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs index c3f2d817bd05..63759d767a2b 100644 --- a/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs +++ b/examples/router/custom-policy-example/simple-stacked-score-pick/src/scorer/uncached_blocks.rs @@ -4,7 +4,8 @@ //! Uncached-block cost for the stacked scorer stage. use dynamo_kv_router::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, + WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicyError, }; pub(crate) struct UncachedBlocksScorer; @@ -17,7 +18,7 @@ impl WorkerScorer for UncachedBlocksScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index 8e42031302bf..c1af90cf007c 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -14,8 +14,8 @@ use dynamo_kv_router::protocols::{ }; use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode}; use dynamo_kv_router::{ - DefaultWorkerSelector, KvRouterConfig, SchedulingRequest, WorkerCandidate, WorkerFilter, - WorkerInputView, WorkerInputs, WorkerLoadProjection, WorkerPicker, WorkerScorer, + DefaultWorkerSelector, KvRouterConfig, SchedulingRequest, WorkerCandidate, WorkerCandidates, + WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadProjection, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, }; @@ -54,7 +54,7 @@ impl WorkerFilter for KeepAllFilter { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, + _candidate: WorkerCandidate<'_>, ) -> Result { Ok(true) } @@ -70,7 +70,7 @@ impl WorkerScorer for BenchScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { @@ -104,7 +104,7 @@ impl WorkerScorer for IgnoringPreferenceScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { diff --git a/lib/kv-router/src/lib.rs b/lib/kv-router/src/lib.rs index bf2df89bdac0..921049a584a2 100644 --- a/lib/kv-router/src/lib.rs +++ b/lib/kv-router/src/lib.rs @@ -81,9 +81,9 @@ pub use selector::DefaultWorkerSelector; pub use selector::{WorkerSelectionInput, WorkerSelector}; // TODO(v1.7): Remove these compatibility re-exports; use crate::plugins instead. pub use plugins::worker_selection::{ - ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputView, - WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, - WorkerSelectionPolicy, + RequestCacheInput, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, + WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, + WorkerSelectionContext, WorkerSelectionPolicy, }; pub use session_prefix_index::{ LogicalNode, NodeId, SessionId, SessionPrefixIndexError, SessionPrefixIndexer, diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index 471387c36298..22b9a0a23a7f 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -4,6 +4,15 @@ //! Public filter, scorer, and picker contracts and their input signals. mod config; +mod context; +mod inputs; + +pub use context::{RequestCacheInput, WorkerSelectionContext}; +pub(crate) use inputs::CandidateData; +pub use inputs::{ + ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, WorkerInputView, + WorkerInputs, WorkerLoadInput, +}; pub use super::registry::{ WorkerSelectionPolicyParameters, WorkerSelectionPolicyProvider, @@ -16,12 +25,8 @@ pub use crate::scheduling::{ pub(crate) use config::RawWorkerSelectionConfig; pub use config::{WorkerSelectionConfig, WorkerSelectionInstance}; -use std::ops::BitOr; -use std::sync::Arc; - -use crate::protocols::{WorkerAffinityTarget, WorkerWithDpRank}; -use crate::scheduling::SchedulingRequest; use crate::{KvRouterConfig, RoutingPartitionRef, WorkerType}; +use std::sync::Arc; /// Factory that creates one worker-selection policy per routing partition. pub type WorkerSelectionPolicyFactory = Arc< @@ -30,98 +35,6 @@ pub type WorkerSelectionPolicyFactory = Arc< + Sync, >; -/// Request-level values available to custom filters, scorers, and pickers. -pub struct WorkerSelectionContext<'a> { - pub(crate) request: &'a SchedulingRequest, - pub(crate) request_blocks: u64, - pub(crate) block_size: u32, - pub(crate) track_prefill_tokens: bool, - pub(crate) has_tier_matches: bool, - pub(crate) inputs: WorkerInputs, - pub(crate) pinned_worker: Option, - pub(crate) router_temperature_override: Option, -} - -/// One eligible worker and the optional inputs requested by a filter or scorer. -pub struct WorkerCandidate { - pub(crate) worker: WorkerWithDpRank, - pub(crate) inputs: WorkerInputs, - pub(crate) cache: WorkerCacheInput, - pub(crate) load: WorkerLoadInput, - pub(crate) preferred_taint_multiplier: Option, -} - -/// One eligible worker and its total cost after all scorers run. -#[derive(Clone, Copy)] -pub struct ScoredWorkerCandidate { - pub(crate) worker: WorkerWithDpRank, - pub(crate) cost: f64, - pub(crate) preferred_taint_multiplier: Option, -} - -/// Optional worker-signal groups requested by scorers and pickers. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub struct WorkerInputs(u8); - -impl WorkerInputs { - /// Request no optional worker inputs. - pub const NONE: Self = Self(0); - /// Request worker KV-cache overlap inputs and request-level cache context. - pub const CACHE: Self = Self(1 << 0); - /// Request active-load inputs. - pub const LOAD: Self = Self(1 << 1); - /// Request preferred-taint routing metadata. - pub const PREFERRED_TAINT: Self = Self(1 << 2); - /// Request host-owned active-request counts. - pub const OCCUPANCY: Self = Self(1 << 5); - #[cfg(any(test, feature = "bench"))] - pub(crate) const ALL: Self = Self(Self::CACHE.0 | Self::LOAD.0 | Self::PREFERRED_TAINT.0); - - pub const fn contains(self, other: Self) -> bool { - self.0 & other.0 == other.0 - } - - pub(crate) fn without(self, other: Self) -> Self { - Self(self.0 & !other.0) - } -} - -impl BitOr for WorkerInputs { - type Output = Self; - - fn bitor(self, rhs: Self) -> Self::Output { - Self(self.0 | rhs.0) - } -} - -/// KV-cache overlap values for one worker. -#[derive(Clone, Copy, Default)] -pub struct WorkerCacheInput { - pub(crate) effective_overlap_blocks: f64, - pub(crate) estimated_cached_tokens: usize, - pub(crate) device_overlap_blocks: f64, - pub(crate) host_overlap_blocks: f64, - pub(crate) disk_overlap_blocks: f64, - pub(crate) shared_beyond_device_blocks: u32, -} - -/// Active-load values for one worker. -#[derive(Clone, Copy, Default)] -pub struct WorkerLoadInput { - pub(crate) available: bool, - pub(crate) active_prefill_tokens: usize, - pub(crate) decode_cost_blocks: f64, - pub(crate) active_requests: usize, -} - -/// Borrowed, index-aligned view of one custom picker's requested worker inputs. -#[derive(Clone, Copy)] -pub struct WorkerInputView<'a> { - pub(crate) candidates: &'a [ScoredWorkerCandidate], - pub(crate) cache: Option<&'a [WorkerCacheInput]>, - pub(crate) load: Option<&'a [WorkerLoadInput]>, -} - /// Adds one finite cost contribution to each eligible worker. pub trait WorkerScorer: Send { /// Declare the worker-signal groups needed by this scorer. @@ -130,16 +43,17 @@ pub trait WorkerScorer: Send { } /// Score every candidate that survived host eligibility and policy filters in one call. - /// The borrowed snapshot and its row order remain fixed for this selection. Empty candidate + /// The borrowed snapshot and its row order remain fixed for this selection. Each scorer + /// can access only its own declared optional inputs, even when other components request more. Empty candidate /// sets skip scoring. Scorers run in declaration order; an error stops before picking. /// /// `costs` has the same length as `candidates` and is initialized to NaN before each call. /// Write one finite, lower-is-better contribution to every slot. The host validates and adds - /// each contribution to that row's total. Neither slice may be retained after this call. + /// each contribution to that row's total. The candidate view and cost slice may not be retained after this call. fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError>; } @@ -157,7 +71,7 @@ pub trait WorkerFilter: Send { fn keep( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result; } @@ -176,250 +90,3 @@ pub trait WorkerPicker: Send { input: WorkerInputView<'_>, ) -> Result; } - -impl WorkerSelectionContext<'_> { - /// The exact worker/rank imposed by the host for this selection, if any. - /// Includes explicit pins and eligible exclusive-affinity targets. This is - /// read-only routing metadata, not permission to change eligibility. - pub fn pinned_worker(&self) -> Option { - self.pinned_worker - } - - /// Exact incoming prompt length in tokens. Borrowed from this request; no rounding, - /// cache weighting, or additional storage is involved. - pub fn prompt_tokens(&self) -> usize { - self.request.isl_tokens - } - - /// Shared-cache ranges from this request's lookup snapshot, if present. - /// Returns None unless this component declared [`WorkerInputs::CACHE`]. - /// The ranges are unweighted block offsets. The host owns their lifetime and - /// policy inspection does not perform a lookup or change accounting. - pub fn shared_cache_hits(&self) -> Option<&crate::SharedCacheHits> { - if self.inputs.contains(WorkerInputs::CACHE) { - self.request.shared_cache_hits.as_ref() - } else { - None - } - } - - /// Whether this request has any tier-specific cache matches, before worker filtering. - /// Returns None unless this component declared [`WorkerInputs::CACHE`]. Some(false) - /// means the host only supplied its accounting estimate (or no cache data). - /// The value describes the current lookup snapshot, not worker cache capacity. - pub fn has_tier_matches(&self) -> Option { - self.inputs - .contains(WorkerInputs::CACHE) - .then_some(self.has_tier_matches) - } - - /// Restrict request-level signals to this component's startup input declaration. - pub(crate) fn with_inputs(&self, inputs: WorkerInputs) -> Self { - Self { - request: self.request, - request_blocks: self.request_blocks, - block_size: self.block_size, - track_prefill_tokens: self.track_prefill_tokens, - has_tier_matches: self.has_tier_matches, - inputs, - pinned_worker: self.pinned_worker, - router_temperature_override: self.router_temperature_override, - } - } - - /// Return the incoming prompt size in KV blocks. - pub fn request_blocks(&self) -> u64 { - self.request_blocks - } - - /// Return the number of tokens in one KV block. - pub fn block_size(&self) -> u32 { - self.block_size - } - - /// Return whether this request contributes to prefill-load tracking. - pub fn tracks_prefill_tokens(&self) -> bool { - self.track_prefill_tokens - } - - /// Return the session metadata available to worker selection. - pub fn session_context(&self) -> Option<&SessionContext> { - self.request.session_context.as_ref() - } - - /// Return the session-affinity target resolved by the request host. - /// - /// The default selector treats an eligible target as exclusive. Custom policies receive it as - /// advisory context; it may be absent from their candidate set when unavailable or filtered. - pub fn affinity_target(&self) -> Option { - self.request.affinity_target - } - - /// Return the expected output length, if the request supplies one. - pub fn expected_output_tokens(&self) -> Option { - self.request.expected_output_tokens - } - - /// Return the request's scheduler priority boost. - pub fn priority_jump(&self) -> f64 { - self.request.priority_jump - } - - /// Return the request's strict integer priority. - pub fn strict_priority(&self) -> u32 { - self.request.strict_priority - } - - /// Return the request-level router temperature override, if present. - pub fn router_temperature_override(&self) -> Option { - self.router_temperature_override - } -} - -impl WorkerCandidate { - /// Return this candidate's worker ID and data-parallel rank. - pub fn worker(&self) -> WorkerWithDpRank { - self.worker - } - - /// Return KV-cache inputs when the component requested [`WorkerInputs::CACHE`]. - pub fn cache(&self) -> Option<&WorkerCacheInput> { - self.inputs - .contains(WorkerInputs::CACHE) - .then_some(&self.cache) - } - - /// Return active-load inputs when the component requested [`WorkerInputs::LOAD`]. - pub fn load(&self) -> Option<&WorkerLoadInput> { - self.inputs - .contains(WorkerInputs::LOAD) - .then_some(&self.load) - } - - /// Return the optional cost multiplier from preferred routing constraints when the component - /// requested [`WorkerInputs::PREFERRED_TAINT`]. - /// - /// Required routing constraints are enforced by host eligibility. This preferred value is - /// ordinary candidate metadata and is only materialized for components that declare the - /// capability. - pub fn preferred_taint_multiplier(&self) -> Option { - self.preferred_taint_multiplier - } - - pub(crate) fn with_inputs_from(&self, additional: &Self, inputs: WorkerInputs) -> Self { - debug_assert_eq!(self.worker, additional.worker); - Self { - worker: self.worker, - inputs, - cache: if inputs.contains(WorkerInputs::CACHE) { - if self.inputs.contains(WorkerInputs::CACHE) { - self.cache - } else { - additional.cache - } - } else { - WorkerCacheInput::default() - }, - load: if inputs.contains(WorkerInputs::LOAD) { - if self.inputs.contains(WorkerInputs::LOAD) { - self.load - } else { - additional.load - } - } else { - WorkerLoadInput::default() - }, - preferred_taint_multiplier: self - .preferred_taint_multiplier - .or(additional.preferred_taint_multiplier), - } - } -} - -impl ScoredWorkerCandidate { - /// Return this candidate's worker ID and data-parallel rank. - pub fn worker(&self) -> WorkerWithDpRank { - self.worker - } - - /// Return the sum of all scorer contributions for this candidate. - pub fn cost(&self) -> f64 { - self.cost - } - - /// Return the optional cost multiplier from preferred routing constraints when the picker - /// requested [`WorkerInputs::PREFERRED_TAINT`]. - pub fn preferred_taint_multiplier(&self) -> Option { - self.preferred_taint_multiplier - } -} - -impl WorkerCacheInput { - /// Host accounting estimate for this worker, in weighted KV blocks and rounded tokens. - /// Lower-tier matches use the host's cache weights. Missing estimates are zero; - /// neither value is clamped to prompt length. This is the current lookup snapshot, - /// not a count of physically resident GPU tokens. Policy scores do not alter it. - pub fn accounting_cache_estimate(&self) -> (f64, usize) { - (self.effective_overlap_blocks, self.estimated_cached_tokens) - } - - /// Return device-resident prefix overlap in KV blocks. - pub fn device_overlap_blocks(&self) -> f64 { - self.device_overlap_blocks - } - - /// Return host-pinned prefix overlap in KV blocks. - pub fn host_overlap_blocks(&self) -> f64 { - self.host_overlap_blocks - } - - /// Return disk prefix overlap in KV blocks. - pub fn disk_overlap_blocks(&self) -> f64 { - self.disk_overlap_blocks - } - - /// Return shared-cache hits beyond the device-resident prefix. - pub fn shared_beyond_device_blocks(&self) -> u32 { - self.shared_beyond_device_blocks - } -} - -impl WorkerLoadInput { - /// Whether the host supplied a load projection for this worker in this selection. - /// False distinguishes a missing observation from an observed idle worker. - pub fn is_available(&self) -> bool { - self.available - } - - /// Return the tokens active in this worker's prefill stage. - pub fn active_prefill_tokens(&self) -> usize { - self.active_prefill_tokens - } - - /// Return the projected active decode footprint in KV blocks. - pub fn decode_cost_blocks(&self) -> f64 { - self.decode_cost_blocks - } - - /// Return this worker's active request count. - pub fn active_requests(&self) -> usize { - self.active_requests - } -} - -impl<'a> WorkerInputView<'a> { - /// Return the eligible candidates and their total costs. - pub fn candidates(self) -> &'a [ScoredWorkerCandidate] { - self.candidates - } - - /// Return index-aligned KV-cache inputs when the picker requested them. - pub fn cache(self) -> Option<&'a [WorkerCacheInput]> { - self.cache - } - - /// Return index-aligned active-load inputs when the picker requested them. - pub fn load(self) -> Option<&'a [WorkerLoadInput]> { - self.load - } -} diff --git a/lib/kv-router/src/plugins/worker_selection/context.rs b/lib/kv-router/src/plugins/worker_selection/context.rs new file mode 100644 index 000000000000..6c50dbae98d0 --- /dev/null +++ b/lib/kv-router/src/plugins/worker_selection/context.rs @@ -0,0 +1,132 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Request metadata and the borrowed, CACHE-gated lookup snapshot. + +use super::{SessionContext, WorkerInputs}; +use crate::protocols::{SharedCacheHits, WorkerAffinityTarget, WorkerWithDpRank}; +use crate::scheduling::SchedulingRequest; + +/// Request-level values available to custom filters, scorers, and pickers. +pub struct WorkerSelectionContext<'a> { + pub(crate) request: &'a SchedulingRequest, + pub(crate) request_blocks: u64, + pub(crate) block_size: u32, + pub(crate) track_prefill_tokens: bool, + pub(crate) has_tier_matches: bool, + pub(crate) inputs: WorkerInputs, + pub(crate) pinned_worker: Option, + pub(crate) router_temperature_override: Option, +} + +/// Request-wide cache facts from the host's current lookup snapshot. +/// Borrowed for the selection callback; no worker data is copied or looked up here. +#[derive(Clone, Copy)] +pub struct RequestCacheInput<'a> { + shared_hits: Option<&'a SharedCacheHits>, + has_tier_matches: bool, +} + +impl<'a> RequestCacheInput<'a> { + /// Unweighted shared-cache ranges in KV block positions, or None if no result was supplied. + /// The host owns this snapshot. It does not change during the callback, and may lag engine state. + /// Use `hits_beyond(prefix)` to exclude hits already covered by the policy's chosen prefix. + pub fn shared_hits(self) -> Option<&'a SharedCacheHits> { + self.shared_hits + } + + /// Whether the request snapshot contains any tier-specific matches before worker filtering. + /// False means only accounting estimates (or no cache data) were supplied. This describes + /// observed matches, not worker cache capacity or whether a particular worker has a match. + pub fn has_tier_matches(self) -> bool { + self.has_tier_matches + } +} + +impl WorkerSelectionContext<'_> { + /// The exact worker/rank imposed by the host for this selection, if any. + /// Includes explicit pins and eligible exclusive-affinity targets. This is + /// read-only routing metadata, not permission to change eligibility. + pub fn pinned_worker(&self) -> Option { + self.pinned_worker + } + + /// Exact incoming prompt length in tokens. Borrowed from this request; no rounding, + /// cache weighting, or additional storage is involved. + pub fn prompt_tokens(&self) -> usize { + self.request.isl_tokens + } + + /// Borrow request-wide cache facts when this component declared [`WorkerInputs::CACHE`]. + /// This is the existing lookup snapshot; accessing it adds no lookup or allocation. + pub fn cache(&self) -> Option> { + self.inputs + .contains(WorkerInputs::CACHE) + .then_some(RequestCacheInput { + shared_hits: self.request.shared_cache_hits.as_ref(), + has_tier_matches: self.has_tier_matches, + }) + } + + /// Restrict request-level signals to this component's startup input declaration. + pub(crate) fn with_inputs(&self, inputs: WorkerInputs) -> Self { + Self { + request: self.request, + request_blocks: self.request_blocks, + block_size: self.block_size, + track_prefill_tokens: self.track_prefill_tokens, + has_tier_matches: self.has_tier_matches, + inputs, + pinned_worker: self.pinned_worker, + router_temperature_override: self.router_temperature_override, + } + } + + /// Return the incoming prompt size in KV blocks. + pub fn request_blocks(&self) -> u64 { + self.request_blocks + } + + /// Return the number of tokens in one KV block. + pub fn block_size(&self) -> u32 { + self.block_size + } + + /// Return whether this request contributes to prefill-load tracking. + pub fn tracks_prefill_tokens(&self) -> bool { + self.track_prefill_tokens + } + + /// Return the session metadata available to worker selection. + pub fn session_context(&self) -> Option<&SessionContext> { + self.request.session_context.as_ref() + } + + /// Return the session-affinity target resolved by the request host. + /// + /// The default selector treats an eligible target as exclusive. Custom policies receive it as + /// advisory context; it may be absent from their candidate set when unavailable or filtered. + pub fn affinity_target(&self) -> Option { + self.request.affinity_target + } + + /// Return the expected output length, if the request supplies one. + pub fn expected_output_tokens(&self) -> Option { + self.request.expected_output_tokens + } + + /// Return the request's scheduler priority boost. + pub fn priority_jump(&self) -> f64 { + self.request.priority_jump + } + + /// Return the request's strict integer priority. + pub fn strict_priority(&self) -> u32 { + self.request.strict_priority + } + + /// Return the request-level router temperature override, if present. + pub fn router_temperature_override(&self) -> Option { + self.router_temperature_override + } +} diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs new file mode 100644 index 000000000000..6466997449ea --- /dev/null +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -0,0 +1,283 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Worker input groups, stored snapshots, and component-scoped borrowed views. + +use crate::protocols::WorkerWithDpRank; +use std::ops::BitOr; + +/// Host-owned row materialized once for the union of component input requirements. +pub(crate) struct CandidateData { + pub(crate) worker: WorkerWithDpRank, + pub(crate) inputs: WorkerInputs, + pub(crate) cache: WorkerCacheInput, + pub(crate) load: WorkerLoadInput, + pub(crate) preferred_taint_multiplier: Option, +} + +/// A borrowed view of one worker, restricted to this component's declared inputs. +#[derive(Clone, Copy)] +pub struct WorkerCandidate<'a> { + data: &'a CandidateData, + inputs: WorkerInputs, +} + +impl<'a> WorkerCandidate<'a> { + pub(crate) fn new(data: &'a CandidateData, inputs: WorkerInputs) -> Self { + Self { data, inputs } + } + + /// Worker identity is always available. + pub fn worker(self) -> WorkerWithDpRank { + self.data.worker + } + + /// Read this worker's cache snapshot only when this component declared CACHE. + pub fn cache(self) -> Option<&'a WorkerCacheInput> { + self.inputs + .contains(WorkerInputs::CACHE) + .then_some(&self.data.cache) + } + + /// Read this worker's load snapshot only when this component declared LOAD. + pub fn load(self) -> Option<&'a WorkerLoadInput> { + self.inputs + .contains(WorkerInputs::LOAD) + .then_some(&self.data.load) + } + + /// Preferred-taint cost multiplier, only when this component declared PREFERRED_TAINT. + /// None means no preference was supplied, access was not requested, or the host pinned the worker. + pub fn preferred_taint_multiplier(self) -> Option { + if self.inputs.contains(WorkerInputs::PREFERRED_TAINT) { + self.data.preferred_taint_multiplier + } else { + None + } + } +} + +/// Borrowed candidate batch restricted to one scorer's declared inputs. +/// All scorers see the same surviving workers in the same order. Creating or iterating this +/// view does not copy rows, allocate, or change another component's access. +#[derive(Clone, Copy)] +pub struct WorkerCandidates<'a> { + rows: &'a [CandidateData], + inputs: WorkerInputs, +} + +impl<'a> WorkerCandidates<'a> { + pub(crate) fn new(rows: &'a [CandidateData], inputs: WorkerInputs) -> Self { + Self { rows, inputs } + } + + /// Number of surviving workers, equal to the scorer's cost-buffer length. + pub fn len(self) -> usize { + self.rows.len() + } + + /// Whether the batch contains no workers. + pub fn is_empty(self) -> bool { + self.rows.is_empty() + } + + /// Borrow one row with this scorer's input permissions, or None for an out-of-range index. + pub fn get(self, row: usize) -> Option> { + self.rows + .get(row) + .map(|data| WorkerCandidate::new(data, self.inputs)) + } + + /// Iterate over the surviving workers without copying their data. + pub fn iter( + self, + ) -> impl ExactSizeIterator> + DoubleEndedIterator + Clone { + self.rows + .iter() + .map(move |data| WorkerCandidate::new(data, self.inputs)) + } +} + +/// One eligible worker and its total cost after all scorers run. +#[derive(Clone, Copy)] +pub struct ScoredWorkerCandidate { + pub(crate) worker: WorkerWithDpRank, + pub(crate) cost: f64, + pub(crate) preferred_taint_multiplier: Option, +} + +/// Optional worker-signal groups requested by scorers and pickers. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct WorkerInputs(u8); + +impl WorkerInputs { + /// Request no optional worker inputs. + pub const NONE: Self = Self(0); + /// Request worker KV-cache overlap inputs and request-level cache context. + pub const CACHE: Self = Self(1 << 0); + /// Request active-load inputs. + pub const LOAD: Self = Self(1 << 1); + /// Request preferred-taint routing metadata. + pub const PREFERRED_TAINT: Self = Self(1 << 2); + /// Request host-owned active-request counts. + pub const OCCUPANCY: Self = Self(1 << 5); + #[cfg(any(test, feature = "bench"))] + pub(crate) const ALL: Self = Self(Self::CACHE.0 | Self::LOAD.0 | Self::PREFERRED_TAINT.0); + + pub const fn contains(self, other: Self) -> bool { + self.0 & other.0 == other.0 + } + + pub(crate) fn without(self, other: Self) -> Self { + Self(self.0 & !other.0) + } +} + +impl BitOr for WorkerInputs { + type Output = Self; + + fn bitor(self, rhs: Self) -> Self::Output { + Self(self.0 | rhs.0) + } +} + +/// KV-cache overlap values for one worker. +#[derive(Clone, Copy, Default)] +pub struct WorkerCacheInput { + pub(crate) effective_overlap_blocks: f64, + pub(crate) estimated_cached_tokens: usize, + pub(crate) device_overlap_blocks: f64, + pub(crate) host_overlap_blocks: f64, + pub(crate) disk_overlap_blocks: f64, +} + +/// Active-load values for one worker. +#[derive(Clone, Copy, Default)] +pub struct WorkerLoadInput { + pub(crate) available: bool, + pub(crate) active_prefill_tokens: usize, + pub(crate) decode_cost_blocks: f64, + pub(crate) active_requests: usize, +} + +/// Borrowed, index-aligned view of one custom picker's requested worker inputs. +#[derive(Clone, Copy)] +pub struct WorkerInputView<'a> { + pub(crate) candidates: &'a [ScoredWorkerCandidate], + pub(crate) cache: Option<&'a [WorkerCacheInput]>, + pub(crate) load: Option<&'a [WorkerLoadInput]>, +} + +impl CandidateData { + pub(crate) fn with_inputs_from(&self, additional: &Self, inputs: WorkerInputs) -> Self { + debug_assert_eq!(self.worker, additional.worker); + Self { + worker: self.worker, + inputs, + cache: if inputs.contains(WorkerInputs::CACHE) { + if self.inputs.contains(WorkerInputs::CACHE) { + self.cache + } else { + additional.cache + } + } else { + WorkerCacheInput::default() + }, + load: if inputs.contains(WorkerInputs::LOAD) { + if self.inputs.contains(WorkerInputs::LOAD) { + self.load + } else { + additional.load + } + } else { + WorkerLoadInput::default() + }, + preferred_taint_multiplier: self + .preferred_taint_multiplier + .or(additional.preferred_taint_multiplier), + } + } +} + +impl ScoredWorkerCandidate { + /// Return this candidate's worker ID and data-parallel rank. + pub fn worker(&self) -> WorkerWithDpRank { + self.worker + } + + /// Return the sum of all scorer contributions for this candidate. + pub fn cost(&self) -> f64 { + self.cost + } + + /// Return the optional cost multiplier from preferred routing constraints when the picker + /// requested [`WorkerInputs::PREFERRED_TAINT`]. + pub fn preferred_taint_multiplier(&self) -> Option { + self.preferred_taint_multiplier + } +} + +impl WorkerCacheInput { + /// Host accounting estimate for this worker, in weighted KV blocks and rounded tokens. + /// Lower-tier matches use the host's cache weights. Missing estimates are zero; + /// neither value is clamped to prompt length. This is the current lookup snapshot, + /// not a count of physically resident GPU tokens. Policy scores do not alter it. + pub fn accounting_cache_estimate(&self) -> (f64, usize) { + (self.effective_overlap_blocks, self.estimated_cached_tokens) + } + + /// Return device-resident prefix overlap in KV blocks. + pub fn device_overlap_blocks(&self) -> f64 { + self.device_overlap_blocks + } + + /// Return host-pinned prefix overlap in KV blocks. + pub fn host_overlap_blocks(&self) -> f64 { + self.host_overlap_blocks + } + + /// Return disk prefix overlap in KV blocks. + pub fn disk_overlap_blocks(&self) -> f64 { + self.disk_overlap_blocks + } +} + +impl WorkerLoadInput { + /// Whether the host supplied a load projection for this worker in this selection. + /// False distinguishes a missing observation from an observed idle worker. + pub fn is_available(&self) -> bool { + self.available + } + + /// Return the tokens active in this worker's prefill stage. + pub fn active_prefill_tokens(&self) -> usize { + self.active_prefill_tokens + } + + /// Return the projected active decode footprint in KV blocks. + pub fn decode_cost_blocks(&self) -> f64 { + self.decode_cost_blocks + } + + /// Return this worker's active request count. + pub fn active_requests(&self) -> usize { + self.active_requests + } +} + +impl<'a> WorkerInputView<'a> { + /// Return the eligible candidates and their total costs. + pub fn candidates(self) -> &'a [ScoredWorkerCandidate] { + self.candidates + } + + /// Return index-aligned KV-cache inputs when the picker requested them. + pub fn cache(self) -> Option<&'a [WorkerCacheInput]> { + self.cache + } + + /// Return index-aligned active-load inputs when the picker requested them. + pub fn load(self) -> Option<&'a [WorkerLoadInput]> { + self.load + } +} diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 811102e381b3..403bc92e629f 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -12,12 +12,14 @@ pub use reference::DefaultWorkerSelector; // TODO(v1.7): Remove these compatibility re-exports; use crate::plugins instead. pub use crate::plugins::worker_selection::{ - ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputView, - WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, + RequestCacheInput, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, + WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, + WorkerSelectionContext, }; #[cfg(any(test, feature = "bench"))] use reference::{DefaultWorkerPicker, DefaultWorkerScorer}; +use crate::plugins::worker_selection::CandidateData; pub use policy::WorkerSelectionPolicy; use policy::{ComposedPolicyState, WorkerSelectionPolicyStateRef, collect_policy_candidates}; #[cfg(any(test, feature = "bench"))] @@ -159,7 +161,7 @@ impl<'a> MaterializedSelectionInput<'a> { worker: WorkerWithDpRank, preferred_taint_multiplier: Option, inputs: WorkerInputs, - ) -> WorkerCandidate { + ) -> CandidateData { self.row_with_device_overlap( worker, preferred_taint_multiplier, @@ -175,10 +177,8 @@ impl<'a> MaterializedSelectionInput<'a> { preferred_taint_multiplier: Option, inputs: WorkerInputs, select_device_overlap: impl FnOnce(f64, f64) -> f64, - ) -> WorkerCandidate { - let cached_tokens = if inputs.contains(WorkerInputs::CACHE) - || (inputs.contains(WorkerInputs::LOAD) && self.request.track_prefill_tokens) - { + ) -> CandidateData { + let cached_tokens = if inputs.contains(WorkerInputs::CACHE) { self.request.effective_cached_tokens_for(worker) } else { 0 @@ -201,12 +201,6 @@ impl<'a> MaterializedSelectionInput<'a> { .unwrap_or(0.0); let device_overlap_blocks = select_device_overlap(effective_overlap_blocks, reported_device_overlap_blocks); - let shared_beyond = |device_blocks: f64| { - self.request.shared_cache_hits.as_ref().map_or(0, |hits| { - // `hits_beyond` expects the unweighted device prefix depth. - hits.hits_beyond(device_blocks.round().max(0.0) as u32) - }) - }; WorkerCacheInput { effective_overlap_blocks, estimated_cached_tokens: cached_tokens, @@ -227,7 +221,6 @@ impl<'a> MaterializedSelectionInput<'a> { .get(&worker) .copied() .unwrap_or(0) as f64, - shared_beyond_device_blocks: shared_beyond(device_overlap_blocks), } } else { WorkerCacheInput::default() @@ -245,7 +238,7 @@ impl<'a> MaterializedSelectionInput<'a> { WorkerLoadInput::default() }; - WorkerCandidate { + CandidateData { worker, inputs, cache, diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index ab148cbb5cc8..e0e7a06ddb8a 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -16,8 +16,9 @@ use crate::scheduling::filter::RoutingEligibility; use crate::scheduling::types::{KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError}; use crate::plugins::worker_selection::{ - ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputs, - WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, + CandidateData, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, + WorkerFilter, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, + WorkerSelectionContext, }; #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] @@ -41,7 +42,7 @@ pub(super) struct ComposedPolicyState { pub(super) filter_inputs: WorkerInputs, pub(super) scorer_picker_inputs: WorkerInputs, pub(super) picker_inputs: WorkerInputs, - unscored_candidates: Vec, + unscored_candidates: Vec, score_contributions: Vec, pub(super) candidates: Vec, pub(super) cache_inputs: Vec, @@ -146,7 +147,7 @@ impl WorkerSelectionPolicy { #[inline(always)] fn push_picker_candidate( - candidate: &WorkerCandidate, + candidate: &CandidateData, cost: f64, picker_inputs: WorkerInputs, candidates: &mut Vec, @@ -156,7 +157,11 @@ fn push_picker_candidate( candidates.push(ScoredWorkerCandidate { worker: candidate.worker, cost, - preferred_taint_multiplier: candidate.preferred_taint_multiplier, + preferred_taint_multiplier: if picker_inputs.contains(WorkerInputs::PREFERRED_TAINT) { + candidate.preferred_taint_multiplier + } else { + None + }, }); if picker_inputs.contains(WorkerInputs::CACHE) { cache_inputs.push(candidate.cache); @@ -169,7 +174,7 @@ fn push_picker_candidate( impl ComposedPolicyState { // Keep row construction and storage together to avoid passing a full row through a call. #[inline(always)] - fn push_candidate(&mut self, candidate: WorkerCandidate) { + fn push_candidate(&mut self, candidate: CandidateData) { // Build the picker's rows alongside the input snapshot. The scoring loop then only // writes costs, without growing vectors or copying optional columns across trait calls. push_picker_candidate( @@ -205,7 +210,11 @@ impl ComposedPolicyState { for (scorer_index, (inputs, scorer)) in scorers.iter_mut().enumerate() { score_contributions.fill(f64::NAN); let context = context.with_inputs(*inputs); - scorer.score(&context, unscored_candidates, score_contributions)?; + scorer.score( + &context, + WorkerCandidates::new(unscored_candidates, *inputs), + score_contributions, + )?; for (row, (contribution, scored)) in score_contributions .iter() .zip(candidates.iter_mut()) @@ -285,7 +294,10 @@ pub(super) fn collect_policy_candidates( ); for (inputs, filter) in &mut state.filters { filter_context.inputs = *inputs; - match filter.keep(&filter_context, &filter_candidate) { + match filter.keep( + &filter_context, + WorkerCandidate::new(&filter_candidate, *inputs), + ) { Ok(true) => {} Ok(false) => return false, Err(policy_error) => { @@ -512,7 +524,7 @@ mod tests { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { @@ -802,7 +814,7 @@ mod tests { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { assert!(candidate.cache().is_none()); assert!(candidate.load().is_none()); @@ -821,7 +833,7 @@ mod tests { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - _candidates: &[WorkerCandidate], + _candidates: WorkerCandidates<'_>, _costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { unreachable!("rejected worker must not reach scorers") @@ -860,7 +872,7 @@ mod tests { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, + _candidate: WorkerCandidate<'_>, ) -> Result { Ok(true) } @@ -875,7 +887,7 @@ mod tests { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (_candidate, cost) in candidates.iter().zip(costs) { diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index 16a98e354ec7..c551a00b135e 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -10,7 +10,7 @@ use parking_lot::Mutex; use super::policy::WorkerSelectionPolicyStateRef; use super::{ - MaterializedSelectionInput, WorkerCandidate, WorkerInputs, WorkerSelectionContext, + CandidateData, MaterializedSelectionInput, WorkerInputs, WorkerSelectionContext, WorkerSelectionInput, WorkerSelector, select_worker_with_policy, }; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank}; @@ -266,7 +266,7 @@ fn default_row( context: DefaultScoringContext, worker: WorkerWithDpRank, preferred_taint_multiplier: Option, -) -> WorkerCandidate { +) -> CandidateData { input.row_with_device_overlap( worker, preferred_taint_multiplier, @@ -282,7 +282,7 @@ impl> DefaultWorkerScorer { &self, context: &WorkerSelectionContext<'_>, default_context: DefaultScoringContext, - row: &WorkerCandidate, + row: &CandidateData, formula_name: &'static str, ) -> f64 { let kv_router_config = self.kv_router_config.borrow(); @@ -292,7 +292,13 @@ impl> DefaultWorkerScorer { let load = &row.load; let effective_overlap_blocks = cache.effective_overlap_blocks; let device_overlap_blocks = cache.device_overlap_blocks; - let shared_beyond_device_blocks = cache.shared_beyond_device_blocks; + let shared_beyond_device_blocks = context + .request + .shared_cache_hits + .as_ref() + .map_or(0, |hits| { + hits.hits_beyond(device_overlap_blocks.round().max(0.0) as u32) + }); let shared_overlap_blocks = weights.shared_cache_multiplier * shared_beyond_device_blocks as f64; // Normalize backlog above the least-loaded eligible worker by this request's @@ -414,7 +420,7 @@ impl> DefaultWorkerScorer { &self, context: &WorkerSelectionContext<'_>, default_context: DefaultScoringContext, - row: &WorkerCandidate, + row: &CandidateData, ) -> f64 { let base_score = self.worker_logit(context, default_context, row, "Formula"); match row.preferred_taint_multiplier { @@ -1911,9 +1917,23 @@ mod tests { let default_row = default_row(&input, default_context, worker, None); assert_eq!(custom_row.cache.device_overlap_blocks, 0.0); - assert_eq!(custom_row.cache.shared_beyond_device_blocks, 4); + assert_eq!( + request + .shared_cache_hits + .as_ref() + .unwrap() + .hits_beyond(custom_row.cache.device_overlap_blocks as u32), + 4 + ); assert_eq!(default_row.cache.device_overlap_blocks, 2.0); - assert_eq!(default_row.cache.shared_beyond_device_blocks, 2); + assert_eq!( + request + .shared_cache_hits + .as_ref() + .unwrap() + .hits_beyond(default_row.cache.device_overlap_blocks as u32), + 2 + ); } /// Without shared cache hits, the scoring should be unchanged. diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 06c694845c13..0735f11710d6 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -703,7 +703,14 @@ impl crate::scheduling::selector::WorkerPicker for CapturingPicker { .cache() .expect("CACHE inputs requested") .iter() - .map(|cache| cache.shared_beyond_device_blocks()) + .map(|cache| { + context + .cache() + .and_then(|request| request.shared_hits()) + .map_or(0, |hits| { + hits.hits_beyond(cache.device_overlap_blocks().round().max(0.0) as u32) + }) + }) .collect(), }); Ok(0) diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index d07eece55881..f2b6c73d5589 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -26,7 +26,7 @@ use crate::scheduling::WorkerSelectionPolicyError; use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::overlap::build_overlap_scores_response; use crate::scheduling::selector::{ - WorkerCandidate, WorkerFilter, WorkerInputView, WorkerPicker, WorkerScorer, + WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, }; use crate::{TrackingHashContext, TrackingHashScope}; @@ -51,7 +51,7 @@ impl WorkerScorer for WorkerIdScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { @@ -85,7 +85,7 @@ impl WorkerScorer for NonFiniteScorer { fn score( &mut self, _context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (_candidate, cost) in candidates.iter().zip(costs) { @@ -113,7 +113,7 @@ impl WorkerFilter for RejectAllFilter { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, + _candidate: WorkerCandidate<'_>, ) -> Result { Ok(false) } @@ -125,7 +125,7 @@ impl WorkerFilter for RejectWorker { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { Ok(candidate.worker().worker_id != self.0) } diff --git a/lib/kv-router/tests/plugin_api_compat.rs b/lib/kv-router/tests/plugin_api_compat.rs index fbc64640e451..c93ae21d0c99 100644 --- a/lib/kv-router/tests/plugin_api_compat.rs +++ b/lib/kv-router/tests/plugin_api_compat.rs @@ -10,9 +10,9 @@ use std::sync::Arc; use dynamo_kv_router::plugins::{RouterPluginRegistry, request_classifier, worker_selection}; use dynamo_kv_router::scheduling::{ClassifyFuture, ClassifyRequest, RequestClassifier}; use dynamo_kv_router::{ - KvRouterConfig, RoutingPartitionRef, WorkerCandidate, WorkerFilter, WorkerInputView, - WorkerInputs, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, - WorkerSelectionPolicyError, WorkerType, + KvRouterConfig, RoutingPartitionRef, WorkerCandidate, WorkerCandidates, WorkerFilter, + WorkerInputView, WorkerInputs, WorkerPicker, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerType, }; struct LegacyPolicy; @@ -25,7 +25,7 @@ impl WorkerFilter for LegacyPolicy { fn keep( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { let _: Option<&dynamo_kv_router::SessionContext> = context.session_context(); let _: Option = context.expected_output_tokens(); @@ -42,7 +42,7 @@ impl WorkerScorer for LegacyPolicy { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 0eba202a0505..70601699f757 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -2210,9 +2210,17 @@ mod tests { .iter() .position(|candidate| candidate.worker() == self.worker) .ok_or_else(|| WorkerSelectionPolicyError::failed("fixed worker not eligible"))?; - let shared = input + let shared = _context .cache() - .map(|cache| cache[row].shared_beyond_device_blocks()) + .and_then(|cache| cache.shared_hits()) + .map(|hits| { + hits.hits_beyond( + input.cache().unwrap()[row] + .device_overlap_blocks() + .round() + .max(0.0) as u32, + ) + }) .filter(|blocks| *blocks > 0); assert_eq!(shared, self.expected_shared_blocks); Ok(row) @@ -2225,7 +2233,7 @@ mod tests { fn keep( &mut self, _context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, + _candidate: WorkerCandidate<'_>, ) -> Result { Ok(false) } diff --git a/lib/router-plugins/builtin/benches/default_selection.rs b/lib/router-plugins/builtin/benches/default_selection.rs index 7c512a8a3251..cf9fbde1ba2b 100644 --- a/lib/router-plugins/builtin/benches/default_selection.rs +++ b/lib/router-plugins/builtin/benches/default_selection.rs @@ -132,7 +132,7 @@ impl dynamo_kv_router::WorkerScorer for LoadScorer { fn score( &mut self, _: &dynamo_kv_router::WorkerSelectionContext<'_>, - candidates: &[dynamo_kv_router::WorkerCandidate], + candidates: dynamo_kv_router::WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), dynamo_kv_router::WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { @@ -202,5 +202,61 @@ fn stacked(c: &mut Criterion) { group.finish(); } } -criterion_group!(benches, bench, stacked); + +// Shared ranges exercise both tier-aware scoring and the accounting-only fallback. +fn shared_cache(c: &mut Criterion) { + for tier_matches in [false, true] { + for range_count in [1, 5, 32] { + for credit in [0.0, 0.6] { + let mut group = c.benchmark_group(format!( + "shared_cache/tiers{tier_matches}/ranges{range_count}/credit{credit}" + )); + group + .warm_up_time(Duration::from_millis(200)) + .measurement_time(Duration::from_millis(500)) + .sample_size(30) + .nresamples(1_000); + for count in [8, 1024, 4096] { + let (workers, mut request) = support::fixture(count, 2048); + if !tier_matches { + request.overlap.tier_overlap_blocks = Default::default(); + } + request.shared_cache_hits = + Some(dynamo_kv_router::SharedCacheHits::from_ranges( + (0..range_count).map(|i| i * 4..i * 4 + 2).collect(), + )); + let policy = default_policy( + KvRouterConfig { + shared_cache_multiplier: credit, + overlap_score_credit_decay: 0.6, + ..Default::default() + }, + "prefill", + ); + let mut select = || { + black_box( + policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(), + ); + }; + let count_allocs = allocations(&mut select); + eprintln!( + "shared allocations per 100 warm selections: tiers={tier_matches} ranges={range_count} credit={credit} workers={count} allocations={count_allocs}" + ); + group + .bench_function(BenchmarkId::new("plugin", count), |b| b.iter(&mut select)); + } + group.finish(); + } + } + } +} + +criterion_group!(benches, bench, stacked, shared_cache); criterion_main!(benches); diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 4a802832d1f1..3268f3330932 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -5,7 +5,8 @@ use super::parameters::PolicyParameters; use dynamo_kv_router::plugins::worker_selection::{ - WorkerCandidate, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, + WorkerCandidate, WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicyError, }; /// Resolve optional terms once so disabled weights add neither branches nor conversions to @@ -56,7 +57,7 @@ struct DefaultScorer { } /// Values shared by every worker score in one selection. Preparation keeps their conversions -/// and role checks outside the per-worker trait call without changing floating-point arithmetic. +/// and role checks outside the per-worker calculation without changing floating-point arithmetic. #[derive(Default)] struct PreparedRequest { min_prefill: usize, @@ -129,7 +130,7 @@ impl fn prepare( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { // Plain disaggregated decode is load-only. Conditional decode retains cache credit. let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { @@ -140,8 +141,9 @@ impl self.prepared = PreparedRequest { min_prefill: 0, has_tier_matches: context - .has_tier_matches() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache context unavailable"))?, + .cache() + .ok_or_else(|| WorkerSelectionPolicyError::failed("cache context unavailable"))? + .has_tier_matches(), block_size: IntegerDivisor::new(u64::from(context.block_size())), request_blocks: IntegerDivisor::new(context.request_blocks()), overlap_credit, @@ -152,7 +154,7 @@ impl }; if self.prepared.use_decay { let mut minimum = usize::MAX; - for candidate in candidates { + for candidate in candidates.iter() { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; @@ -166,7 +168,7 @@ impl fn score_worker( &mut self, context: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { let cache = candidate .cache() @@ -181,13 +183,10 @@ impl estimated_overlap }; let shared_credit = if SHARED_CREDIT { - let shared = if self.prepared.has_tier_matches { - cache.shared_beyond_device_blocks() - } else { - context - .shared_cache_hits() - .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)) - }; + let shared = context + .cache() + .and_then(|cache| cache.shared_hits()) + .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)); self.shared_cache_multiplier * shared as f64 } else { // Preserve signed zero when the configured weight is -0.0. @@ -241,7 +240,7 @@ impl WorkerScorer fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.prepare(context, candidates)?; diff --git a/lib/router-plugins/builtin/tests/cache_context.rs b/lib/router-plugins/builtin/tests/cache_context.rs index a3ba47ef33af..03b071374c5a 100644 --- a/lib/router-plugins/builtin/tests/cache_context.rs +++ b/lib/router-plugins/builtin/tests/cache_context.rs @@ -4,16 +4,16 @@ mod support; use dynamo_kv_router::{ - KvRouterConfig, SharedCacheHits, WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, - WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, - WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, + KvRouterConfig, SharedCacheHits, WorkerCandidate, WorkerCandidates, WorkerFilter, + WorkerInputView, WorkerInputs, WorkerPicker, WorkerScorer, WorkerSelectionContext, + WorkerSelectionInput, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, }; use std::sync::{ Arc, atomic::{AtomicUsize, Ordering}, }; -struct CacheProbe { +struct InputProbe { inputs: WorkerInputs, tier_matches: bool, shared_hits: bool, @@ -21,30 +21,59 @@ struct CacheProbe { calls: Arc, } -impl CacheProbe { +impl InputProbe { fn declare(&self) -> WorkerInputs { self.declarations.fetch_add(1, Ordering::Relaxed); self.inputs } + fn check_candidate(&self, candidate: WorkerCandidate<'_>) { + assert_eq!( + candidate.cache().is_some(), + self.inputs.contains(WorkerInputs::CACHE) + ); + assert_eq!( + candidate.load().is_some(), + self.inputs.contains(WorkerInputs::LOAD) + ); + assert_eq!( + candidate.preferred_taint_multiplier().is_some(), + self.inputs.contains(WorkerInputs::PREFERRED_TAINT) + ); + if let Some(cache) = candidate.cache() { + assert_eq!( + cache.host_overlap_blocks(), + if self.tier_matches { 2.0 } else { 0.0 } + ); + } + if let Some(load) = candidate.load() { + assert!(load.is_available()); + } + } + fn check(&self, context: &WorkerSelectionContext<'_>) { self.calls.fetch_add(1, Ordering::Relaxed); // Ordinary request context stays available without CACHE. assert_eq!(context.prompt_tokens(), 17); if self.inputs.contains(WorkerInputs::CACHE) { - assert_eq!(context.has_tier_matches(), Some(self.tier_matches)); - assert_eq!(context.shared_cache_hits().is_some(), self.shared_hits); - if let Some(hits) = context.shared_cache_hits() { + assert_eq!( + context.cache().unwrap().has_tier_matches(), + self.tier_matches + ); + assert_eq!( + context.cache().unwrap().shared_hits().is_some(), + self.shared_hits + ); + if let Some(hits) = context.cache().unwrap().shared_hits() { assert_eq!(hits.hits_beyond(0), 4); } } else { - assert_eq!(context.has_tier_matches(), None); - assert!(context.shared_cache_hits().is_none()); + assert!(context.cache().is_none()); } } } -impl WorkerFilter for CacheProbe { +impl WorkerFilter for InputProbe { fn required_worker_inputs(&self) -> WorkerInputs { self.declare() } @@ -52,14 +81,15 @@ impl WorkerFilter for CacheProbe { fn keep( &mut self, context: &WorkerSelectionContext<'_>, - _candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { self.check(context); + self.check_candidate(candidate); Ok(true) } } -impl WorkerScorer for CacheProbe { +impl WorkerScorer for InputProbe { fn required_worker_inputs(&self) -> WorkerInputs { self.declare() } @@ -67,16 +97,28 @@ impl WorkerScorer for CacheProbe { fn score( &mut self, context: &WorkerSelectionContext<'_>, - _candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.check(context); + assert_eq!(candidates.len(), costs.len()); + assert!(!candidates.is_empty()); + assert!(candidates.get(candidates.len()).is_none()); + for (row, candidate) in candidates.iter().enumerate() { + self.check_candidate(candidate); + let indexed = candidates.get(row).unwrap(); + self.check_candidate(indexed); + assert_eq!(indexed.worker(), candidate.worker()); + if let Some(cache) = candidate.cache() { + assert!(std::ptr::eq(cache, indexed.cache().unwrap())); + } + } costs.fill(0.0); Ok(()) } } -impl WorkerPicker for CacheProbe { +impl WorkerPicker for InputProbe { fn required_worker_inputs(&self) -> WorkerInputs { self.declare() } @@ -84,64 +126,101 @@ impl WorkerPicker for CacheProbe { fn pick( &mut self, context: &WorkerSelectionContext<'_>, - _input: WorkerInputView<'_>, + input: WorkerInputView<'_>, ) -> Result { self.check(context); + assert_eq!( + input.cache().is_some(), + self.inputs.contains(WorkerInputs::CACHE) + ); + assert_eq!( + input.load().is_some(), + self.inputs.contains(WorkerInputs::LOAD) + ); + for candidate in input.candidates() { + assert_eq!( + candidate.preferred_taint_multiplier().is_some(), + self.inputs.contains(WorkerInputs::PREFERRED_TAINT) + ); + } Ok(0) } } +fn inputs(mask: usize) -> WorkerInputs { + [ + WorkerInputs::CACHE, + WorkerInputs::LOAD, + WorkerInputs::PREFERRED_TAINT, + ] + .into_iter() + .enumerate() + .fold(WorkerInputs::NONE, |inputs, (bit, group)| { + if mask & (1 << bit) == 0 { + inputs + } else { + inputs | group + } + }) +} + #[test] -fn cache_context_requires_each_components_own_declaration() { - for mask in 0..8 { - for unrequested in [WorkerInputs::NONE, WorkerInputs::LOAD] { - for tier_matches in [false, true] { - for shared_hits in [false, true] { - let (workers, mut request) = support::fixture(2, 17); - if !tier_matches { - request.overlap.tier_overlap_blocks = Default::default(); - } - if shared_hits { - request.shared_cache_hits = - Some(SharedCacheHits::from_ranges(vec![1..3, 5..7])); - } - let declarations = Arc::new(AtomicUsize::new(0)); - let calls = Arc::new(AtomicUsize::new(0)); - let probe = |inputs| CacheProbe { - inputs, - tier_matches, - shared_hits, - declarations: declarations.clone(), - calls: calls.clone(), - }; - let inputs = |bit| { - if mask & bit == 0 { - unrequested - } else { - WorkerInputs::CACHE +fn each_component_can_only_read_its_own_inputs() { + for filter_mask in 0..8 { + for scorer_mask in 0..8 { + for picker_mask in 0..8 { + for tier_matches in [false, true] { + for shared_hits in [false, true] { + let (workers, mut request) = support::fixture(2, 17); + if !tier_matches { + request.overlap.tier_overlap_blocks = Default::default(); + } + if shared_hits { + request.shared_cache_hits = + Some(SharedCacheHits::from_ranges(vec![1..3, 5..7])); + } + // Even unmatched preferences produce Some(1.0), so an undeclared + // component leaking another component's taint value fails this test. + request + .routing_constraints + .preferred_taints + .insert("preferred".into(), 1.0); + let declarations = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let probe = |mask| InputProbe { + inputs: inputs(mask), + tier_matches, + shared_hits, + declarations: declarations.clone(), + calls: calls.clone(), + }; + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + vec![ + Box::new(probe(filter_mask)), + Box::new(probe(7 ^ filter_mask)), + ], + vec![ + Box::new(probe(scorer_mask)), + Box::new(probe(7 ^ scorer_mask)), + ], + Box::new(probe(picker_mask)), + ); + assert_eq!(declarations.load(Ordering::Relaxed), 5); + for _ in 0..2 { + policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); } - }; - let policy = WorkerSelectionPolicy::new_with_filters( - KvRouterConfig::default(), - "test", - vec![Box::new(probe(inputs(1))), Box::new(probe(unrequested))], - vec![Box::new(probe(inputs(2))), Box::new(probe(unrequested))], - Box::new(probe(inputs(4))), - ); - assert_eq!(declarations.load(Ordering::Relaxed), 5); - for _ in 0..2 { - policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) - .unwrap(); + assert_eq!(calls.load(Ordering::Relaxed), 22); + assert_eq!(declarations.load(Ordering::Relaxed), 5); } - // Four worker/ranks pass two filters, then two scorers and one picker run. - assert_eq!(calls.load(Ordering::Relaxed), 22); - assert_eq!(declarations.load(Ordering::Relaxed), 5); } } } diff --git a/lib/router-plugins/builtin/tests/default_lifecycle.rs b/lib/router-plugins/builtin/tests/default_lifecycle.rs index 2ac05945629f..ac6b11b09007 100644 --- a/lib/router-plugins/builtin/tests/default_lifecycle.rs +++ b/lib/router-plugins/builtin/tests/default_lifecycle.rs @@ -90,7 +90,7 @@ async fn shared_core_books_and_releases_with_the_builtin_default() { #[tokio::test] async fn preparation_error_does_not_book_and_allows_retry() { use dynamo_kv_router::{ - WorkerCandidate, WorkerInputView, WorkerPicker, WorkerScorer, WorkerSelectionContext, + WorkerCandidates, WorkerInputView, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, }; use std::sync::Arc; @@ -100,7 +100,7 @@ async fn preparation_error_does_not_book_and_allows_retry() { fn prepare( &mut self, _: &WorkerSelectionContext<'_>, - _: &[WorkerCandidate], + _: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { if std::mem::take(&mut self.0) { return Err(WorkerSelectionPolicyError::failed("preparation failed")); @@ -113,7 +113,7 @@ async fn preparation_error_does_not_book_and_allows_retry() { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.prepare(context, candidates)?; diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index d76b9ab0fbb9..24ff97ca1ca1 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -172,7 +172,7 @@ fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { ) -> Result { assert_eq!(context.prompt_tokens(), 17); assert_eq!(context.request_blocks(), 2); - assert_eq!(context.has_tier_matches(), Some(true)); + assert!(context.cache().unwrap().has_tier_matches()); assert!(input.load().unwrap().iter().all(|load| load.is_available())); for cache in input.cache().unwrap() { let (blocks, tokens) = cache.accounting_cache_estimate(); diff --git a/lib/router-plugins/builtin/tests/scorer_prepare.rs b/lib/router-plugins/builtin/tests/scorer_prepare.rs index fe91168a1c61..5c55faa9adb1 100644 --- a/lib/router-plugins/builtin/tests/scorer_prepare.rs +++ b/lib/router-plugins/builtin/tests/scorer_prepare.rs @@ -7,9 +7,9 @@ mod support; use std::sync::{Arc, Mutex}; use dynamo_kv_router::{ - KvRouterConfig, WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, WorkerPicker, - WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, - WorkerSelectionPolicyError, WorkerSelector, + KvRouterConfig, WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerInputs, + WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, + WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, }; use support::fixture; @@ -29,7 +29,7 @@ impl RelativeLoadScorer { fn prepare( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { assert_eq!(context.prompt_tokens(), 17); assert!(candidates.iter().all(|c| c.cache().is_none())); @@ -56,7 +56,7 @@ impl WorkerScorer for RelativeLoadScorer { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.prepare(context, candidates)?; @@ -73,7 +73,7 @@ impl WorkerFilter for ExcludeWorkerOne { fn keep( &mut self, _: &WorkerSelectionContext<'_>, - candidate: &WorkerCandidate, + candidate: WorkerCandidate<'_>, ) -> Result { Ok(candidate.worker().worker_id != 1) } @@ -171,7 +171,7 @@ impl FailPreparation { fn prepare( &mut self, _: &WorkerSelectionContext<'_>, - _: &[WorkerCandidate], + _: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { *self.0.lock().unwrap() += 1; Err(WorkerSelectionPolicyError::failed("prepare failed")) @@ -182,7 +182,7 @@ impl WorkerScorer for FailPreparation { fn score( &mut self, context: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, _costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { self.prepare(context, candidates)?; @@ -245,7 +245,7 @@ fn rejects_nonfinite_contributions_and_overflow_before_picking() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (_candidate, cost) in candidates.iter().zip(costs) { @@ -291,7 +291,7 @@ fn picker_columns_and_costs_stay_aligned_after_a_scoring_error() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - candidates: &[WorkerCandidate], + candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { for (candidate, cost) in candidates.iter().zip(costs) { @@ -363,7 +363,7 @@ fn unwritten_costs_cannot_reuse_a_previous_selection_or_scorer() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - _: &[WorkerCandidate], + _: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { costs.fill(7.0); @@ -375,7 +375,7 @@ fn unwritten_costs_cannot_reuse_a_previous_selection_or_scorer() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - _: &[WorkerCandidate], + _: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { if std::mem::replace(&mut self.0, false) { From fdd12d1b38a5f43b26beac77a345d9a5b3de714e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 23:16:37 +0000 Subject: [PATCH 17/31] fix(router): gate cache setup on worker inputs Signed-off-by: Ishan Dhanani --- .../configuration/groups/kv_router_args.py | 5 +- .../router/configuration-and-tuning.md | 6 +- .../router/custom-worker-selection.mdx | 4 +- .../router/router-operations.md | 2 +- .../components/frontend-configuration.mdx | 4 +- .../observability/metrics-catalog.mdx | 2 +- .../router/custom-policy-example/README.md | 2 + lib/kv-router/src/scheduling/config.rs | 44 +-- lib/llm/src/discovery/model_manager.rs | 19 +- lib/llm/src/kv_router.rs | 328 ++++++++++-------- lib/llm/src/kv_router/indexer/ingress.rs | 14 +- lib/llm/src/kv_router/indexer/mod.rs | 8 - 12 files changed, 223 insertions(+), 215 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 15f844989af7..372f63d8f6a0 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -295,11 +295,12 @@ def add_arguments(self, parser) -> None: default=False, dest="load_aware", help=( - "KV Router: Enable load-aware routing without cache-reuse signals. " + "KV Router: Enable the load-aware routing preset. " "On the frontend, this implies --router-mode kv. " "This preset sets overlap_score_credit=0, disables KV events and " "KV-reuse assumptions, enables active-block " - "and prefill-token load tracking, and disables remote/shared cache indexers." + "and prefill-token load tracking, and disables remote/shared cache indexers. " + "The policy's CACHE input declaration still controls local approximate indexing." ), ) add_argument( diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md index 08026aa76981..466274e299a5 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md @@ -70,13 +70,13 @@ its configuration even if the endpoint is unchanged. ## Routing Behavior -- `--router-kv-overlap-score-credit`: Device-local prefix-overlap credit multiplier in the prefill cost calculation. It must be finite and nonnegative. Values greater than `1.0` give overlap extra credit, but the adjusted prefill contribution is clamped at zero. When set to `0`, the router ignores prefix caches and skips creating a local indexer. Defaults to `1.0`. +- `--router-kv-overlap-score-credit`: Device-local prefix-overlap credit multiplier in the prefill cost calculation. It must be finite and nonnegative. Values greater than `1.0` give overlap extra credit, but the adjusted prefill contribution is clamped at zero. When set to `0`, device-local overlap receives no score credit. Cache setup still follows the policy's `WorkerInputs::CACHE` declaration. Defaults to `1.0`. - `--router-kv-overlap-score-credit-decay`: Decays device-local overlap credit for workers whose active prefill load exceeds the least-loaded eligible worker. `0` disables decay. Defaults to 0. - `--router-prefill-load-scale`: Scale applied to adjusted prompt-side prefill load after device, lower-tier, and shared-cache credits are subtracted. Defaults to 1. - `--router-decode-active-request-weight`: Experimental finite, nonnegative block-equivalent decode cost added for each active request on a candidate worker. Defaults to 0. - `--router-host-cache-hit-weight`: Credit multiplier for host-pinned (CPU offload) prefix overlap, from 0.0 to 1.0. Symmetric to `--router-kv-overlap-score-credit` but applied to the host-pinned tier when a backend exposes CPU offload via a KV connector. Defaults to 0.75. - `--router-disk-cache-hit-weight`: Credit multiplier for disk/lower-tier (e.g. NVMe-backed) prefix overlap, from 0.0 to 1.0. Defaults to 0.25. -- `--load-aware`: Preset for load-aware KV routing without cache-reuse signals. On the frontend, it implies `--router-mode kv`. It sets `overlap_score_credit=0`, disables KV events and KV reuse assumptions, enables active-block and prefill-token load tracking, disables remote/shared cache indexers, and preserves `--router-prefill-load-scale`, `--router-host-cache-hit-weight`, and `--router-disk-cache-hit-weight`. +- `--load-aware`: Preset for load-aware KV routing. On the frontend, it implies `--router-mode kv`. It sets `overlap_score_credit=0`, disables KV events and KV reuse assumptions, enables active-block and prefill-token load tracking, disables remote/shared cache indexers, and preserves `--router-prefill-load-scale`, `--router-host-cache-hit-weight`, and `--router-disk-cache-hit-weight`. The policy's `WorkerInputs::CACHE` declaration still controls local approximate indexing; the builtin default declares `CACHE`. - `--router-temperature`: Controls worker selection randomness through softmax sampling of normalized router cost logits. A value of 0 (default) ensures deterministic selection of the lowest-cost worker, while higher values introduce more randomness. - `--router-conditional-disagg`: **Experimental.** Enables conditional disaggregation in frontend-embedded disaggregated serving. Requires `--router-mode kv`, `--router-kv-events`, separate prefill/decode worker pools, and decode-worker KV event publishing. Use `--router-conditional-disagg-config` for policy settings. See [Conditional Disaggregation](../../../advanced-customizations/conditional-disaggregation.md) for backend requirements and policy tuning. - `--router-track-prefill-tokens`: Enables prompt-side load accounting in the worker cost model. This should stay enabled if you want queue thresholds, `active_prefill_tokens`, and AIC prefill load decay to reflect prompt work. @@ -445,7 +445,7 @@ Use `--router-kv-overlap-score-credit-decay` to reduce that device-local credit Use `--load-aware` when you want the KV scheduler's active load model without prefix/cache reuse. This is equivalent to using KV mode with overlap credit set to 0, KV events disabled, KV reuse assumptions disabled, active load tracking enabled, and shared-cache routing disabled. `--router-prefill-load-scale` remains available to tune prompt-side load relative to decode blocks. -Deprecated: `--router-kv-overlap-score-weight`, `--kv-overlap-score-weight`, `DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT`, and `DYN_OVERLAP_SCORE_WEIGHT` are still accepted, but emit deprecation warnings. Nonzero legacy values map to `prefill_load_scale` to preserve existing behavior without changing overlap credit. A legacy value of 0 maps to both `prefill_load_scale=0` and `overlap_score_credit=0`, which preserves the old no-overlap/no-indexer behavior. If a deprecated overlap score weight is still present, it takes precedence over the newer prefill load scale field; a legacy value of 0 also takes precedence over the newer overlap credit field. When migrating to `--router-prefill-load-scale` or `DYN_ROUTER_PREFILL_LOAD_SCALE`, remove the deprecated flag, env var, or JSON field from the deployment config. Use `--router-kv-overlap-score-credit` or `DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT` only when you mean to tune the cache-overlap credit itself. +Deprecated: `--router-kv-overlap-score-weight`, `--kv-overlap-score-weight`, `DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT`, and `DYN_OVERLAP_SCORE_WEIGHT` are still accepted, but emit deprecation warnings. Nonzero legacy values map to `prefill_load_scale` to preserve existing behavior without changing overlap credit. A legacy value of 0 maps to both `prefill_load_scale=0` and `overlap_score_credit=0`, which zeros those scoring weights. Indexer setup follows the policy's `WorkerInputs::CACHE` declaration independently of these values. If a deprecated overlap score weight is still present, it takes precedence over the newer prefill load scale field; a legacy value of 0 also takes precedence over the newer overlap credit field. When migrating to `--router-prefill-load-scale` or `DYN_ROUTER_PREFILL_LOAD_SCALE`, remove the deprecated flag, env var, or JSON field from the deployment config. Use `--router-kv-overlap-score-credit` or `DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT` only when you mean to tune the cache-overlap credit itself. When migrating the deprecated overlap score weight, use `--router-prefill-load-scale` to preserve its scaling role. Tune `--router-kv-overlap-score-credit` separately only when you intend to change device-local cache credit; values above `1.0` are supported, with adjusted prefill cost clamped at zero. diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index b711ba2f3528..ef37fdacf97f 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -433,6 +433,8 @@ The disaggregated prefill hop and query-only probes do not run classifiers. Stat Worker identity and ordinary request metadata are always available. Request-level cache context requires `WorkerInputs::CACHE`. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. +The combined `WorkerInputs` declarations also control cache setup. Without `CACHE`, the router creates no routing indexer, KV event subscription, or shared-cache client. With `CACHE`, cache settings choose event-backed, approximate, or remote indexing. Worker role, conditional disaggregation, and scoring weights do not enable or disable this setup. `serve_indexer` and `enable_session_prefix_index` require a policy that declares `CACHE`; incompatible configurations fail before cache resources start. + ### Request Context | Accessor | Meaning | @@ -662,7 +664,7 @@ worker_selection: router_temperature: 0.0 ``` -The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. The default plugin uses configured parameters and does not apply per-request score or temperature overrides. Plain disaggregated decode remains load-only when prefill tracking is off and conditional disaggregation is disabled. +The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. The default scorer declares `CACHE`, including when its device credit is zero; select a policy without `CACHE` to omit routing cache setup. The default plugin uses configured parameters and does not apply per-request score or temperature overrides. Plain disaggregated decode remains load-only when prefill tracking is off and conditional disaggregation is disabled. `WorkerCacheInput::accounting_cache_estimate()` returns the host's weighted block estimate and rounded token estimate. Missing estimates are zero. These estimates are not clamped to prompt length and do not measure physical GPU residency. The plugin can use them without changing the load that the host records. `WorkerLoadInput::is_available()` distinguishes a missing load observation from an observed idle worker. Both values describe the current selection snapshot. diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/router-operations.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/router-operations.md index 2f925bd67d0f..bbbf00ed09bf 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/router-operations.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/router-operations.md @@ -130,7 +130,7 @@ graph TD Request-plane transport is independent of KV event transport. The request plane (`DYN_REQUEST_PLANE` or `--request-plane`) controls how requests reach workers. ZMQ is the default event plane for every discovery backend, including etcd. Set `--event-plane nats` or `DYN_EVENT_PLANE=nats` to opt into NATS Core. With the default ZMQ event plane, the router does not require NATS; the selected discovery backend still determines whether etcd is required. When using the NATS event plane, NATS is initialized automatically; set `NATS_SERVER=nats://...` to override the default `localhost:4222`. -When `--router-kv-overlap-score-credit` is set to 0, no KV indexer is created and prefix matching is disabled. When `--no-router-kv-events` is set, a KV indexer is still created but no event subscriber is launched; the router predicts cache state from its own routing decisions using the selected approximate cache policy. +The worker-selection policy's `WorkerInputs::CACHE` declaration controls whether the router creates a routing indexer. Setting `--router-kv-overlap-score-credit` to 0 changes device-cache scoring but does not disable indexing. When the policy requests `CACHE` and `--no-router-kv-events` is set, the router creates an approximate indexer without an event subscriber and predicts cache state from its own routing decisions. A policy without `CACHE` starts neither an indexer nor an event subscriber. Backend KV event publishing is independent of the frontend's `--no-router-kv-events` flag. The frontend flag controls whether the router consumes events; backend flags control whether workers publish them. If the router is not consuming events, workers that still publish will waste resources but cause no harm. diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 76179995cdbc..2ccb48174d03 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -127,13 +127,13 @@ router. See [Router Guide](../../developer-guide/knowledge-base/modular-componen ### KV scoring and cache locality - Preset for KV load-aware routing without cache-reuse signals; implies `--router-mode kv`. + Preset for KV load-aware routing; implies `--router-mode kv`. Disables device-cache credit, KV event consumption, and remote/shared cache indexers. The worker-selection policy's `CACHE` declaration still controls local approximate indexing. Environment variable: `DYN_ROUTER_LOAD_AWARE` - Credit multiplier for device-local prefix overlap. The value must be finite and nonnegative. + Credit multiplier for device-local prefix overlap. The value must be finite and nonnegative. Setting it to zero does not disable indexing; cache setup follows the worker-selection policy's `WorkerInputs::CACHE` declaration. Values greater than `1.0` give device overlap extra credit, but adjusted prefill cost is clamped at zero. See [Configuration and Tuning](../../developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md#tuning-guidelines). diff --git a/docs/fern/pages/reference/observability/metrics-catalog.mdx b/docs/fern/pages/reference/observability/metrics-catalog.mdx index a5962adcd4df..4a6c54d47466 100644 --- a/docs/fern/pages/reference/observability/metrics-catalog.mdx +++ b/docs/fern/pages/reference/observability/metrics-catalog.mdx @@ -348,7 +348,7 @@ The frontend registers these metrics when queueing is enabled by `--router-queue ### KV indexer metrics - KV cache events applied to the router's radix-tree index. Includes events applied to the device tier and lower tiers such as host-pinned memory and disk. Appears only when `--router-kv-overlap-score-credit` is greater than 0 and workers publish KV events. Labeled by `status` and `event_type` — see [Metric-specific labels](metric-labels.mdx#metric-specific-labels). + KV cache events applied to the router's radix-tree index. Includes events applied to the device tier and lower tiers such as host-pinned memory and disk. Appears when the worker-selection policy requests `WorkerInputs::CACHE`, the router consumes KV events, and workers publish them. Labeled by `status` and `event_type` — see [Metric-specific labels](metric-labels.mdx#metric-specific-labels). The standalone indexer exposes the device-tier-only counter `dynamo_kvrouter_kv_cache_events_applied`; see [Standalone KV Indexer](../../developer-guide/knowledge-base/modular-components/router/standalone-indexer.md). diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 55de546142c0..45e940bf87bf 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -377,6 +377,8 @@ Send the same `curl` request from a fourth terminal. The frontend log records se Filters, scorers, and pickers can read `context.prompt_tokens()` for the exact input token count. This returns the existing request value without a prompt copy or allocation. Use `context.request_blocks()` when the policy needs rounded KV blocks instead. For example, at 16 tokens per block, a 17-token prompt has 2 blocks. +The union of component declarations controls routing cache setup. Without `WorkerInputs::CACHE`, no routing indexer, KV event subscription, or shared-cache client is started, regardless of score weights or worker role. With `CACHE`, router settings select event-backed, approximate, or remote indexing. Serving an indexer or enabling the session prefix index requires a policy that declares `CACHE`. + Request-level cache information also requires `WorkerInputs::CACHE`: `context.cache()` returns `None` without that declaration. With it, the borrowed view exposes `has_tier_matches()` and `shared_hits()`. The former is false when the snapshot has no tier matches; the latter is `None` when no shared-cache result was supplied. A policy can call `shared_hits().map(|hits| hits.hits_beyond(prefix))` to calculate shared hits after its chosen prefix, only when needed. Scorers receive a `WorkerCandidates` view with `iter()`, `get(row)`, `len()`, and `is_empty()`. Filters receive one borrowed `WorkerCandidate`. Both expose `cache()`, `load()`, and `preferred_taint_multiplier()` only when that component declares the corresponding `CACHE`, `LOAD`, or `PREFERRED_TAINT` group. The same rule applies to picker inputs. Another component requesting a group does not grant access. Views share the host's stored rows without allocating or copying a candidate table. diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 4bc3e3c75043..d0802d296a68 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -1091,16 +1091,8 @@ fn validate_kv_router_config(config: &KvRouterConfig) -> Result<(), String> { "enable_session_prefix_index is not supported with use_remote_indexer=true".to_string(), ); } - if config.enable_session_prefix_index - && (!config.use_kv_events || config.overlap_score_credit <= 0.0) - { - return Err( - "enable_session_prefix_index requires use_kv_events=true and overlap_score_credit > 0" - .to_string(), - ); - } - if config.serve_indexer && config.overlap_score_credit == 0.0 { - return Err("serve_indexer requires overlap_score_credit > 0".to_string()); + if config.enable_session_prefix_index && !config.use_kv_events { + return Err("enable_session_prefix_index requires use_kv_events=true".to_string()); } if config.router_predicted_ttl_secs.is_some() && !config.use_kv_events { return Err("router_predicted_ttl_secs requires use_kv_events=true".to_string()); @@ -1583,18 +1575,6 @@ impl KvRouterConfig { precomputed_block_hashes, )) } - - /// Check if KV event subscription should be started. - /// - /// Returns false if: - /// - KV events are disabled (`use_kv_events=false`) - /// - Overlap scoring is disabled (`overlap_score_credit=0`) - /// - /// When false, the router skips starting the KV event subscription entirely, - /// avoiding the need to query workers for their local indexer state. - pub fn should_subscribe_to_kv_events(&self) -> bool { - self.use_kv_events && self.overlap_score_credit > 0.0 - } } fn random_sequence_hashes(num_blocks: usize) -> Vec { @@ -2484,7 +2464,6 @@ models: assert_eq!(config.overlap_score_credit, 0.0); assert_eq!(config.prefill_load_scale, 0.0); - assert!(!config.should_subscribe_to_kv_events()); } #[test] @@ -2585,14 +2564,17 @@ models: } #[test] - fn test_overlap_credit_zero_skips_kv_event_subscription() { - let config = KvRouterConfig { - overlap_score_credit: 0.0, - use_kv_events: true, - ..Default::default() - }; - - assert!(!config.should_subscribe_to_kv_events()); + fn indexer_features_do_not_depend_on_scoring_credit() { + for (serve_indexer, enable_session_prefix_index) in [(true, false), (false, true)] { + let config = KvRouterConfig { + overlap_score_credit: 0.0, + use_kv_events: true, + serve_indexer, + enable_session_prefix_index, + ..Default::default() + }; + assert!(config.validate().is_ok()); + } } #[test] diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 754c4631f7d0..641d89fb8c71 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -2323,19 +2323,12 @@ impl ModelManager { None }; - let kv_event_source_requirement = - KvEventSourceRequirement::derive(worker_role, &effective_kv_router_config); - let cache_required = wants_cache - || effective_kv_router_config.serve_indexer - || effective_kv_router_config.enable_session_prefix_index - || matches!( - kv_event_source_requirement, - KvEventSourceRequirement::ConditionalDisaggDecodeCache - | KvEventSourceRequirement::Unknown - ); - let kv_source_membership = if cache_required - && kv_event_source_requirement.should_subscribe(&effective_kv_router_config) - { + let kv_event_source_requirement = KvEventSourceRequirement::derive( + worker_role, + &effective_kv_router_config, + policy.inputs(), + ); + let kv_source_membership = if kv_event_source_requirement.should_subscribe() { Some( self.get_or_create_kv_source_membership_watch(&endpoint) .await?, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 70601699f757..60b7e5fea42e 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -139,16 +139,26 @@ impl SelectionPolicySource { label: &'static str, model_name: Option<&str>, ) -> Result { - if let Self::Prepared(prepared) = self { - return Ok(prepared); + let prepared = match self { + Self::Prepared(prepared) => prepared, + source => PreparedSelectionPolicy::prepare( + source.resolve(config, worker_type, label)?, + config, + worker_type, + model_name, + ), + }; + if !prepared.inputs().contains(WorkerInputs::CACHE) { + anyhow::ensure!( + !config.serve_indexer, + "serve_indexer requires the worker-selection policy to declare WorkerInputs::CACHE" + ); + anyhow::ensure!( + !config.enable_session_prefix_index, + "enable_session_prefix_index requires the worker-selection policy to declare WorkerInputs::CACHE" + ); } - let factory = self.resolve(config, worker_type, label)?; - Ok(PreparedSelectionPolicy::prepare( - factory, - config, - worker_type, - model_name, - )) + Ok(prepared) } } @@ -396,20 +406,24 @@ pub(crate) enum KvEventSourceRequirement { } impl KvEventSourceRequirement { - pub(crate) fn derive(worker_role: Option, config: &KvRouterConfig) -> Self { - let Some(worker_role) = worker_role else { - return Self::Unknown; - }; - if config.use_remote_indexer || !config.should_subscribe_to_kv_events() { + pub(crate) fn derive( + worker_role: Option, + config: &KvRouterConfig, + inputs: WorkerInputs, + ) -> Self { + if !inputs.contains(WorkerInputs::CACHE) + || config.use_remote_indexer + || !config.use_kv_events + { return Self::NotRequired; } match worker_role { - WorkerType::Prefill | WorkerType::Aggregated => Self::CacheAwareRouting, - WorkerType::Decode if config.conditional_disagg_enabled => { + Some(WorkerType::Decode) if config.conditional_disagg_enabled => { Self::ConditionalDisaggDecodeCache } - WorkerType::Decode | WorkerType::Encode => Self::NotRequired, + Some(_) => Self::CacheAwareRouting, + None => Self::Unknown, } } @@ -429,11 +443,8 @@ impl KvEventSourceRequirement { ) } - pub(crate) fn should_subscribe(self, config: &KvRouterConfig) -> bool { - match self { - Self::Unknown => !config.use_remote_indexer && config.should_subscribe_to_kv_events(), - requirement => requirement.requires_source(), - } + pub(crate) fn should_subscribe(self) -> bool { + self != Self::NotRequired } } @@ -735,16 +746,11 @@ impl KvRouter { let tracking_hash = TrackingHashContext::from_config(&kv_router_config)?; let tracking_model_name = resolve_tracking_model_name(tracking_hash.algorithm(), model_name.as_deref())?; - let kv_event_source_requirement = - KvEventSourceRequirement::derive(worker_role, &kv_router_config); - let cache_required = required_worker_inputs.contains(WorkerInputs::CACHE) - || kv_router_config.serve_indexer - || kv_router_config.enable_session_prefix_index - || matches!( - kv_event_source_requirement, - KvEventSourceRequirement::ConditionalDisaggDecodeCache - | KvEventSourceRequirement::Unknown - ); + let kv_event_source_requirement = KvEventSourceRequirement::derive( + worker_role, + &kv_router_config, + required_worker_inputs, + ); let component = endpoint.component(); // All chooser tasks are children of the routing load context owner. let cancellation_token = parent_token.child_token(); @@ -761,7 +767,7 @@ impl KvRouter { model_name: model_name.as_deref(), worker_role, metric_worker_type, - cache_required, + worker_inputs: required_worker_inputs, kv_event_source_requirement, kv_source_membership, cancellation_token: cancellation_token.clone(), @@ -771,7 +777,7 @@ impl KvRouter { let indexer = ingress.indexer().clone(); let approximate_lru_metrics = metrics::ApproximateLruMetrics::from_component(component); let configured_policy = kv_router_config.router_approximate_cache_policy.to_string(); - let effective_policy = if kv_router_config.overlap_score_credit <= 0.0 { + let effective_policy = if matches!(indexer, Indexer::None) { "disabled" } else if indexer.uses_approximate_lru() { "lru" @@ -2075,95 +2081,53 @@ mod tests { #[test] fn kv_event_source_requirement_matrix() { - let default = KvRouterConfig::default(); - let mut cases = vec![ - ( - Some(WorkerType::Prefill), - default.clone(), - KvEventSourceRequirement::CacheAwareRouting, - true, - ), - ( - Some(WorkerType::Aggregated), - default.clone(), - KvEventSourceRequirement::CacheAwareRouting, - true, - ), - ( - Some(WorkerType::Decode), - default.clone(), - KvEventSourceRequirement::NotRequired, - false, - ), - ( - Some(WorkerType::Encode), - default.clone(), - KvEventSourceRequirement::NotRequired, - false, - ), - ( - None, - default.clone(), - KvEventSourceRequirement::Unknown, - true, - ), - ]; - for policy in [ - dynamo_kv_router::ConditionalDisaggPolicyKind::IslBounding, - dynamo_kv_router::ConditionalDisaggPolicyKind::PrefillLoad, - dynamo_kv_router::ConditionalDisaggPolicyKind::IslOrLoad, - ] { - cases.push(( - Some(WorkerType::Decode), - KvRouterConfig { - conditional_disagg_enabled: true, - conditional_disagg_policy: policy, - ..default.clone() - }, - KvEventSourceRequirement::ConditionalDisaggDecodeCache, - true, - )); - } - for config in [ - KvRouterConfig { - use_remote_indexer: true, - ..Default::default() - }, - KvRouterConfig { - use_kv_events: false, - ..Default::default() - }, - KvRouterConfig { - overlap_score_credit: 0.0, - ..Default::default() - }, + for role in [ + None, + Some(WorkerType::Aggregated), + Some(WorkerType::Prefill), + Some(WorkerType::Decode), + Some(WorkerType::Encode), ] { - cases.extend([ - ( - None, - config.clone(), - KvEventSourceRequirement::Unknown, - false, - ), - ( - Some(WorkerType::Aggregated), - config.clone(), - KvEventSourceRequirement::NotRequired, - false, - ), - ( - Some(WorkerType::Decode), - config, - KvEventSourceRequirement::NotRequired, - false, - ), - ]); - } - - for (role, config, expected, should_subscribe) in cases { - let requirement = KvEventSourceRequirement::derive(role, &config); - assert_eq!(requirement, expected); - assert_eq!(requirement.should_subscribe(&config), should_subscribe); + for inputs in [WorkerInputs::LOAD, WorkerInputs::CACHE | WorkerInputs::LOAD] { + for overlap_score_credit in [0.0, 1.0, 2.0] { + for use_kv_events in [false, true] { + for use_remote_indexer in [false, true] { + for conditional_disagg_enabled in [false, true] { + let config = KvRouterConfig { + overlap_score_credit, + use_kv_events, + use_remote_indexer, + conditional_disagg_enabled, + ..Default::default() + }; + let expected = if !inputs.contains(WorkerInputs::CACHE) + || !use_kv_events + || use_remote_indexer + { + KvEventSourceRequirement::NotRequired + } else if role.is_none() { + KvEventSourceRequirement::Unknown + } else if role == Some(WorkerType::Decode) + && conditional_disagg_enabled + { + KvEventSourceRequirement::ConditionalDisaggDecodeCache + } else { + KvEventSourceRequirement::CacheAwareRouting + }; + let actual = + KvEventSourceRequirement::derive(role, &config, inputs); + assert_eq!(actual, expected); + assert_eq!( + actual.should_subscribe(), + inputs.contains(WorkerInputs::CACHE) + && use_kv_events + && !use_remote_indexer, + ); + } + } + } + } + } } } @@ -2296,7 +2260,7 @@ mod tests { "role-aware-subscription", HashMap::from([(7, ModelRuntimeConfig::default())]), 16, - SelectionPolicySource::Registry, + picker_policy(|| Box::new(LoadOnlyPicker)), None, worker_role, "decode", @@ -2359,21 +2323,77 @@ mod tests { } #[tokio::test] - async fn constructor_skips_sources_for_decode_but_preserves_unknown_behavior() { - let router = make_router_without_membership(Some(WorkerType::Decode)) - .await - .expect("ordinary decode must not require KV source membership"); - assert!(!router.ingress.has_subscription()); + async fn cache_inputs_control_indexer_and_subscription() { + for role in [None, Some(WorkerType::Aggregated), Some(WorkerType::Decode)] { + for use_kv_events in [false, true] { + for overlap_score_credit in [0.0, 1.0] { + for wants_cache in [false, true] { + let policy = if wants_cache { + fixed_policy(None, WorkerWithDpRank::from_worker_id(7)) + } else { + picker_policy(|| Box::new(LoadOnlyPicker)) + }; + let result = make_router( + "cache-input-gate", + HashMap::from([(7, ModelRuntimeConfig::default())]), + 16, + policy, + None, + role, + "decode", + KvRouterConfig { + use_kv_events, + overlap_score_credit, + conditional_disagg_enabled: true, + skip_initial_worker_wait: true, + router_event_threads: 1, + ..Default::default() + }, + ) + .await; + if wants_cache && use_kv_events { + // No membership watch was provided: reaching subscription setup is + // required even for zero scoring credit and decode/unknown roles. + let error = + result.err().expect("CACHE needs a source membership watch"); + assert!( + error + .to_string() + .contains("KV source membership watch is required") + ); + } else { + let router = result.unwrap(); + assert_eq!(matches!(router.indexer, Indexer::None), !wants_cache); + assert!(!router.ingress.has_subscription()); + router.cancellation_token.cancel(); + } + } + } + } + } + } - let error = make_router_without_membership(None) - .await - .err() - .expect("unknown role must preserve config-driven subscription"); - assert!( - error - .to_string() - .contains("KV source membership watch is required") - ); + #[tokio::test] + async fn indexer_features_require_declared_cache_inputs() { + for feature in ["serve_indexer", "enable_session_prefix_index"] { + let config = KvRouterConfig { + serve_indexer: feature == "serve_indexer", + enable_session_prefix_index: feature == "enable_session_prefix_index", + overlap_score_credit: 0.0, + ..Default::default() + }; + config.validate().unwrap(); + let source = picker_policy(|| Box::new(LoadOnlyPicker)); + let error = source + .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) + .err() + .expect("host features must not enable undeclared cache input"); + assert!(error.to_string().contains(feature)); + assert!(error.to_string().contains("WorkerInputs::CACHE")); + fixed_policy(None, WorkerWithDpRank::from_worker_id(7)) + .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) + .expect("CACHE allows indexer features even with zero scoring credit"); + } } #[tokio::test] @@ -2557,24 +2577,42 @@ mod tests { async fn session_prefix_tracking_survives_shared_core_selection(#[case] threads: u32) { use dynamo_kv_router::protocols::{ExternalSequenceBlockHash, StorageTier}; - let router = make_router( - "session-prefix-core", - HashMap::from([(7, ModelRuntimeConfig::default())]), + let component = make_test_component("session-prefix-core").await; + let endpoint = component.endpoint("backend"); + let client = endpoint.client().await.unwrap(); + let (_tx, workers) = watch::channel(HashMap::from([(7, ModelRuntimeConfig::default())])); + let membership = crate::discovery::KvSourceMembershipCoordinator::start( + endpoint.id(), + workers.clone(), + component.drt().discovery(), + ) + .subscribe(); + let router = KvRouter::new_with_worker_role( + endpoint, + client, + workers, + Some(membership), 2, SelectionPolicySource::Registry, - None, - Some(WorkerType::Decode), - "decode", - KvRouterConfig { + Some(KvRouterConfig { enable_session_prefix_index: true, router_event_threads: threads, router_temperature: 0.0, + overlap_score_credit: 0.0, skip_initial_worker_wait: true, ..Default::default() - }, + }), + None, + Some(WorkerType::Decode), + "decode", + None, + false, + None, + None, ) .await .unwrap(); + assert!(router.ingress.has_subscription()); let session_index = router.session_prefix_index.as_ref().unwrap(); let worker = WorkerWithDpRank::new(7, 0); let tokens = [11, 12, 13, 14]; diff --git a/lib/llm/src/kv_router/indexer/ingress.rs b/lib/llm/src/kv_router/indexer/ingress.rs index 987144a41665..e3f2089cc849 100644 --- a/lib/llm/src/kv_router/indexer/ingress.rs +++ b/lib/llm/src/kv_router/indexer/ingress.rs @@ -13,7 +13,7 @@ use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::identity::RoutingPartitionId; use dynamo_kv_router::services::indexer::registry::WorkerRegistry; use dynamo_kv_router::services::selection::KvEventIngress; -use dynamo_kv_router::{SessionPrefixIndexer, WorkerType}; +use dynamo_kv_router::{SessionPrefixIndexer, WorkerInputs, WorkerType}; use dynamo_runtime::component::Endpoint; use tokio_util::sync::CancellationToken; @@ -31,8 +31,8 @@ pub(crate) struct RuntimeIngressArgs<'a> { pub model_name: Option<&'a str>, pub worker_role: Option, pub metric_worker_type: &'static str, - /// Whether any consumer needs KV knowledge; without it the index is `None`. - pub cache_required: bool, + /// Combined policy inputs. Without CACHE, no routing indexer is created. + pub worker_inputs: WorkerInputs, pub kv_event_source_requirement: KvEventSourceRequirement, pub kv_source_membership: Option, pub cancellation_token: CancellationToken, @@ -54,13 +54,14 @@ impl RuntimeIngress { model_name, worker_role, metric_worker_type, - cache_required, + worker_inputs, kv_event_source_requirement, kv_source_membership, cancellation_token, session_prefix_index, } = args; let component = endpoint.component(); + let cache_required = worker_inputs.contains(WorkerInputs::CACHE); let indexer = if cache_required { super::build( component, @@ -75,9 +76,7 @@ impl RuntimeIngress { Indexer::None }; - let subscription = if cache_required - && kv_event_source_requirement.should_subscribe(kv_router_config) - { + let subscription = if cache_required && kv_event_source_requirement.should_subscribe() { let membership_watch = kv_source_membership.ok_or_else(|| { anyhow::anyhow!( "KV source membership watch is required when local KV event subscription is enabled" @@ -102,7 +101,6 @@ impl RuntimeIngress { requirement = %kv_event_source_requirement, cache_required = cache_required, use_kv_events = kv_router_config.use_kv_events, - overlap_score_credit = kv_router_config.overlap_score_credit, use_remote_indexer = kv_router_config.use_remote_indexer, "Skipping KV event subscription" ); diff --git a/lib/llm/src/kv_router/indexer/mod.rs b/lib/llm/src/kv_router/indexer/mod.rs index 98d4800ee492..f74ab7d9b720 100644 --- a/lib/llm/src/kv_router/indexer/mod.rs +++ b/lib/llm/src/kv_router/indexer/mod.rs @@ -54,7 +54,6 @@ pub(crate) use recovery::{ #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ResolvedApproximatePrimaryPolicy { - Disabled, Ttl, Lru, TtlRemoteFallback, @@ -70,9 +69,6 @@ fn resolve_approximate_primary_policy( "router_approximate_cache_policy=lru requires use_kv_events=false; the local side indexer is TTL-only" ); } - if config.overlap_score_credit <= 0.0 { - return Ok(ResolvedApproximatePrimaryPolicy::Disabled); - } if config.use_kv_events || config.router_approximate_cache_policy == ApproximateCachePolicyKind::Ttl { @@ -94,10 +90,6 @@ pub(crate) async fn build( session_prefix_index: Option>, ) -> Result { let approximate_policy = resolve_approximate_primary_policy(kv_router_config)?; - if approximate_policy == ResolvedApproximatePrimaryPolicy::Disabled { - return Ok(Indexer::None); - } - if approximate_policy == ResolvedApproximatePrimaryPolicy::TtlRemoteFallback { tracing::warn!( use_remote_indexer = kv_router_config.use_remote_indexer, From be4c9d12fc789f52cbbf811a084c94ed111d82d8 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Mon, 21 Sep 2026 23:38:55 +0000 Subject: [PATCH 18/31] test(router): deduplicate builtin selection setup Signed-off-by: Ishan Dhanani --- .../builtin/benches/default_selection.rs | 44 +++-------------- .../builtin/tests/cache_context.rs | 9 +--- .../builtin/tests/default_policy.rs | 48 ++++--------------- .../builtin/tests/scorer_prepare.rs | 48 ++++--------------- .../builtin/tests/support/mod.rs | 10 +++- 5 files changed, 35 insertions(+), 124 deletions(-) diff --git a/lib/router-plugins/builtin/benches/default_selection.rs b/lib/router-plugins/builtin/benches/default_selection.rs index cf9fbde1ba2b..b942b5605943 100644 --- a/lib/router-plugins/builtin/benches/default_selection.rs +++ b/lib/router-plugins/builtin/benches/default_selection.rs @@ -4,7 +4,7 @@ mod support; use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; use dynamo_custom_policy_builtin::default_policy; -use dynamo_kv_router::{KvRouterConfig, WorkerSelectionInput, WorkerSelector}; +use dynamo_kv_router::{KvRouterConfig, WorkerSelector}; use std::alloc::{GlobalAlloc, Layout, System}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; @@ -67,24 +67,14 @@ fn bench(c: &mut Criterion) { let reference_allocs = allocations(|| { black_box( reference - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(), ); }); let plugin_allocs = allocations(|| { black_box( plugin - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(), ); }); @@ -95,12 +85,7 @@ fn bench(c: &mut Criterion) { b.iter(|| { black_box( reference - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(), ) }) @@ -109,12 +94,7 @@ fn bench(c: &mut Criterion) { b.iter(|| { black_box( plugin - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(), ) }) @@ -184,12 +164,7 @@ fn stacked(c: &mut Criterion) { let mut select = || { black_box( policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(), ); }; @@ -236,12 +211,7 @@ fn shared_cache(c: &mut Criterion) { let mut select = || { black_box( policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(), ); }; diff --git a/lib/router-plugins/builtin/tests/cache_context.rs b/lib/router-plugins/builtin/tests/cache_context.rs index 03b071374c5a..3be3299f8ed4 100644 --- a/lib/router-plugins/builtin/tests/cache_context.rs +++ b/lib/router-plugins/builtin/tests/cache_context.rs @@ -6,7 +6,7 @@ mod support; use dynamo_kv_router::{ KvRouterConfig, SharedCacheHits, WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerInputs, WorkerPicker, WorkerScorer, WorkerSelectionContext, - WorkerSelectionInput, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, + WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, }; use std::sync::{ Arc, @@ -210,12 +210,7 @@ fn each_component_can_only_read_its_own_inputs() { assert_eq!(declarations.load(Ordering::Relaxed), 5); for _ in 0..2 { policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); } assert_eq!(calls.load(Ordering::Relaxed), 22); diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 24ff97ca1ca1..7a2dc0b0dd76 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -5,8 +5,8 @@ use dynamo_custom_policy_builtin::{DefaultWorkerSelector, default_policy, defaul use dynamo_kv_router::protocols::WorkerWithDpRank; use dynamo_kv_router::{ KvRouterConfig, RoutingPartitionRef, WorkerInputView, WorkerInputs, WorkerPicker, - WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, - WorkerSelectionPolicyError, WorkerSelector, WorkerType, + WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, + WorkerType, }; use support::*; @@ -45,12 +45,7 @@ fn seeded_selection_matches_reference_across_cache_and_load_shapes() { ); let plugin = DefaultWorkerSelector::new_seeded(Some(config), "test", 42); for _ in 0..64 { - let input = WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - ); + let input = support::selection_input(&workers, &request, 16); let expected = reference.select_worker(input).unwrap(); let actual = plugin.select_worker(input).unwrap(); assert_eq!( @@ -92,12 +87,7 @@ fn unseeded_minimum_picker_only_selects_workers_tied_for_lowest_cost() { } for _ in 0..32 { let selected = policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); assert!(best.contains(&selected.worker.worker_id)); } @@ -136,12 +126,7 @@ fn prepared_values_follow_each_request_when_reusing_a_policy() { load.active_prefill_tokens += round * 19; } } - let input = WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - block_size, - ); + let input = support::selection_input(&workers, &request, block_size); let expected = reference.select_worker(input).unwrap(); let actual = plugin.select_worker(input).unwrap(); assert_eq!( @@ -183,12 +168,7 @@ fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { } let (workers, request) = fixture(2, 17); WorkerSelectionPolicy::new(KvRouterConfig::default(), "test", vec![], Box::new(Inspect)) - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); } @@ -220,12 +200,7 @@ fn mandatory_pin_is_preserved() { let (workers, mut request) = fixture(4, 17); request.pinned_worker = Some(WorkerWithDpRank::new(3, 1)); let selected = default_policy(KvRouterConfig::default(), "test") - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); assert_eq!(Some(selected.worker), request.pinned_worker); } @@ -260,12 +235,7 @@ fn configured_parameters_replace_request_score_overrides() { let policy = default_policy(config, "prefill"); for _ in 0..32 { let result = policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); assert_eq!(result.worker.worker_id, 0); } @@ -303,7 +273,7 @@ fn pin_does_not_advance_seeded_random_stream() { let plugin = DefaultWorkerSelector::new_seeded(Some(config), "prefill", 42); for pinned in [true, false, true, false, false] { request.pinned_worker = pinned.then_some(WorkerWithDpRank::new(3, 1)); - let input = WorkerSelectionInput::configured(&workers, &request, request.eligibility(), 16); + let input = support::selection_input(&workers, &request, 16); assert_eq!( reference.select_worker(input).unwrap().worker, plugin.select_worker(input).unwrap().worker diff --git a/lib/router-plugins/builtin/tests/scorer_prepare.rs b/lib/router-plugins/builtin/tests/scorer_prepare.rs index 5c55faa9adb1..3e83592997c2 100644 --- a/lib/router-plugins/builtin/tests/scorer_prepare.rs +++ b/lib/router-plugins/builtin/tests/scorer_prepare.rs @@ -8,8 +8,8 @@ use std::sync::{Arc, Mutex}; use dynamo_kv_router::{ KvRouterConfig, WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerInputs, - WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionInput, - WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, + WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, + WorkerSelectionPolicyError, WorkerSelector, }; use support::fixture; @@ -139,12 +139,7 @@ fn scores_batches_from_surviving_workers_and_resets_each_selection() { request.allowed_worker_ids = Some([3].into_iter().collect()); } let result = policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); assert_eq!(result.worker.worker_id, 3); let mut calls = calls.lock().unwrap(); @@ -215,24 +210,14 @@ fn skips_preparation_for_empty_sets_and_aborts_on_preparation_error() { request.allowed_worker_ids = Some(allowed.into_iter().collect()); assert!( policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .is_err() ); assert_eq!(*calls.lock().unwrap(), 0); } request.allowed_worker_ids = None; let error = policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap_err(); assert!(error.to_string().contains("prepare failed")); assert_eq!(*calls.lock().unwrap(), 1); @@ -270,12 +255,7 @@ fn rejects_nonfinite_contributions_and_overflow_before_picking() { Box::new(NeverPick), ); let error = policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) + .select_worker(support::selection_input(&workers, &request, 16)) .unwrap_err(); assert!(error.to_string().contains("non-finite")); } @@ -337,12 +317,7 @@ fn picker_columns_and_costs_stay_aligned_after_a_scoring_error() { Box::new(CheckColumns), ); let select = |request: &dynamo_kv_router::scheduling::SchedulingRequest| { - policy.select_worker(WorkerSelectionInput::configured( - &workers, - request, - request.eligibility(), - 16, - )) + policy.select_worker(support::selection_input(&workers, request, 16)) }; assert!( select(&request) @@ -402,14 +377,7 @@ fn unwritten_costs_cannot_reuse_a_previous_selection_or_scorer() { vec![Box::new(Fill), Box::new(OmitOnSecondCall(true))], Box::new(First), ); - let select = || { - policy.select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) - }; + let select = || policy.select_worker(support::selection_input(&workers, &request, 16)); select().unwrap(); let error = select().unwrap_err(); assert!(matches!( diff --git a/lib/router-plugins/builtin/tests/support/mod.rs b/lib/router-plugins/builtin/tests/support/mod.rs index 626dafd225ad..86d0fc7e918f 100644 --- a/lib/router-plugins/builtin/tests/support/mod.rs +++ b/lib/router-plugins/builtin/tests/support/mod.rs @@ -1,8 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use dynamo_kv_router::WorkerLoadProjection; use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank}; use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode, SchedulingRequest}; +use dynamo_kv_router::{WorkerLoadProjection, WorkerSelectionInput}; use std::collections::HashMap; #[derive(Clone, Copy)] @@ -82,3 +82,11 @@ pub fn fixture( } (workers, request) } + +pub fn selection_input<'a>( + workers: &'a HashMap, + request: &'a SchedulingRequest, + block_size: u32, +) -> WorkerSelectionInput<'a, TestWorker> { + WorkerSelectionInput::configured(workers, request, request.eligibility(), block_size) +} From 7445237c4c91721aac70f7d7aa23621f76ad853c Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 00:43:31 +0000 Subject: [PATCH 19/31] test(router): simplify builtin batch scoring fixtures Signed-off-by: Ishan Dhanani --- .../{scorer_prepare.rs => batch_scoring.rs} | 91 +++++++------------ .../builtin/tests/default_lifecycle.rs | 45 ++------- 2 files changed, 39 insertions(+), 97 deletions(-) rename lib/router-plugins/builtin/tests/{scorer_prepare.rs => batch_scoring.rs} (84%) diff --git a/lib/router-plugins/builtin/tests/scorer_prepare.rs b/lib/router-plugins/builtin/tests/batch_scoring.rs similarity index 84% rename from lib/router-plugins/builtin/tests/scorer_prepare.rs rename to lib/router-plugins/builtin/tests/batch_scoring.rs index 3e83592997c2..c083d1e99431 100644 --- a/lib/router-plugins/builtin/tests/scorer_prepare.rs +++ b/lib/router-plugins/builtin/tests/batch_scoring.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Exercise batch scoring and private preparation using only public plugin inputs. +//! Exercise the batch scoring contract using only public plugin inputs. mod support; use std::sync::{Arc, Mutex}; @@ -15,21 +15,24 @@ use support::fixture; #[derive(Debug, PartialEq)] enum Call { - Prepare(usize, Vec), - Score(usize), + Score(usize, Vec), Pick, } struct RelativeLoadScorer { index: usize, - minimum: usize, calls: Arc>>, } -impl RelativeLoadScorer { - fn prepare( +impl WorkerScorer for RelativeLoadScorer { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } + + fn score( &mut self, context: &WorkerSelectionContext<'_>, candidates: WorkerCandidates<'_>, + costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { assert_eq!(context.prompt_tokens(), 17); assert!(candidates.iter().all(|c| c.cache().is_none())); @@ -38,31 +41,14 @@ impl RelativeLoadScorer { self.calls .lock() .unwrap() - .push(Call::Prepare(self.index, ids)); - self.minimum = candidates + .push(Call::Score(self.index, ids)); + let minimum = candidates .iter() .map(|c| c.load().unwrap().active_requests()) .min() .unwrap(); - Ok(()) - } -} - -impl WorkerScorer for RelativeLoadScorer { - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::LOAD - } - - fn score( - &mut self, - context: &WorkerSelectionContext<'_>, - candidates: WorkerCandidates<'_>, - costs: &mut [f64], - ) -> Result<(), WorkerSelectionPolicyError> { - self.prepare(context, candidates)?; for (candidate, cost) in candidates.iter().zip(costs) { - self.calls.lock().unwrap().push(Call::Score(self.index)); - *cost = (candidate.load().unwrap().active_requests() - self.minimum) as f64; + *cost = (candidate.load().unwrap().active_requests() - minimum) as f64; } Ok(()) } @@ -119,14 +105,13 @@ fn scores_batches_from_surviving_workers_and_resets_each_selection() { .map(|index| { Box::new(RelativeLoadScorer { index, - minimum: usize::MAX, calls: calls.clone(), }) as Box }) .collect(), Box::new(InspectCosts(calls.clone())), ); - // Excluded workers have zero load. Neither may lower the preparation minimum. + // Excluded workers have zero load. Neither may lower the batch minimum. for round in 0..3 { for (worker, load) in &mut request.worker_loads { load.active_requests = match worker.worker_id { @@ -148,40 +133,28 @@ fn scores_batches_from_surviving_workers_and_resets_each_selection() { } else { vec![2, 2, 3, 3] }; - assert_eq!(calls.len(), 2 + expected.len() * 2 + 1); - for index in 0..2 { - let start = index * (expected.len() + 1); - assert_eq!(calls[start], Call::Prepare(index, expected.clone())); - for call in &calls[start + 1..start + 1 + expected.len()] { - assert_eq!(*call, Call::Score(index)); - } - } - assert_eq!(calls.last(), Some(&Call::Pick)); + assert_eq!( + *calls, + vec![ + Call::Score(0, expected.clone()), + Call::Score(1, expected), + Call::Pick + ] + ); calls.clear(); } } -struct FailPreparation(Arc>); -impl FailPreparation { - fn prepare( +struct FailScoring(Arc>); +impl WorkerScorer for FailScoring { + fn score( &mut self, _: &WorkerSelectionContext<'_>, _: WorkerCandidates<'_>, + _: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { *self.0.lock().unwrap() += 1; - Err(WorkerSelectionPolicyError::failed("prepare failed")) - } -} - -impl WorkerScorer for FailPreparation { - fn score( - &mut self, - context: &WorkerSelectionContext<'_>, - candidates: WorkerCandidates<'_>, - _costs: &mut [f64], - ) -> Result<(), WorkerSelectionPolicyError> { - self.prepare(context, candidates)?; - panic!("preparation failure must stop scoring") + Err(WorkerSelectionPolicyError::failed("score failed")) } } struct NeverPick; @@ -196,14 +169,14 @@ impl WorkerPicker for NeverPick { } #[test] -fn skips_preparation_for_empty_sets_and_aborts_on_preparation_error() { +fn skips_scoring_for_empty_sets_and_aborts_on_error() { let (workers, mut request) = fixture(2, 17); let calls = Arc::new(Mutex::new(0)); let policy = WorkerSelectionPolicy::new_with_filters( KvRouterConfig::default(), "test", vec![Box::new(ExcludeWorkerOne)], - vec![Box::new(FailPreparation(calls.clone()))], + vec![Box::new(FailScoring(calls.clone()))], Box::new(NeverPick), ); for allowed in [vec![], vec![1]] { @@ -219,7 +192,7 @@ fn skips_preparation_for_empty_sets_and_aborts_on_preparation_error() { let error = policy .select_worker(support::selection_input(&workers, &request, 16)) .unwrap_err(); - assert!(error.to_string().contains("prepare failed")); + assert!(error.to_string().contains("score failed")); assert_eq!(*calls.lock().unwrap(), 1); } @@ -230,12 +203,10 @@ fn rejects_nonfinite_contributions_and_overflow_before_picking() { fn score( &mut self, _: &WorkerSelectionContext<'_>, - candidates: WorkerCandidates<'_>, + _: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { - for (_candidate, cost) in candidates.iter().zip(costs) { - *cost = self.0; - } + costs.fill(self.0); Ok(()) } } diff --git a/lib/router-plugins/builtin/tests/default_lifecycle.rs b/lib/router-plugins/builtin/tests/default_lifecycle.rs index ac6b11b09007..857bbf38332b 100644 --- a/lib/router-plugins/builtin/tests/default_lifecycle.rs +++ b/lib/router-plugins/builtin/tests/default_lifecycle.rs @@ -6,7 +6,6 @@ use dynamo_kv_router::services::selection::{ PromptRequest, SelectAndReserveRequest, SelectionCacheConfig, SelectionCore, WorkerRequest, }; use dynamo_kv_router::{KvRouterConfig, RoutingPartitionId}; -use std::collections::{HashMap, HashSet}; fn worker(worker_id: u64) -> WorkerRequest { WorkerRequest { @@ -14,24 +13,9 @@ fn worker(worker_id: u64) -> WorkerRequest { model_name: "model".to_string(), routing_group: "default".to_string(), endpoint: Some(format!("http://worker-{worker_id}:8000")), - kv_events_endpoint: None, - kv_events_endpoints: HashMap::new(), - replay_endpoint: None, block_size: Some(4), - data_parallel_start_rank: None, - data_parallel_size: None, max_num_batched_tokens: Some(1024), - total_kv_blocks: None, - stable_routing_id: None, - is_eagle: None, - taints: HashSet::new(), - topology_domains: HashMap::new(), - kv_transfer_domain: None, - kv_transfer_enforcement: None, - kv_transfer_preferred_weight: None, - router_hint_worker_type: None, - router_hint_source_control_endpoints: HashMap::new(), - kv_event_source_mode: None, + ..Default::default() } } @@ -88,7 +72,7 @@ async fn shared_core_books_and_releases_with_the_builtin_default() { } #[tokio::test] -async fn preparation_error_does_not_book_and_allows_retry() { +async fn scoring_error_does_not_book_and_allows_retry() { use dynamo_kv_router::{ WorkerCandidates, WorkerInputView, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, @@ -96,30 +80,17 @@ async fn preparation_error_does_not_book_and_allows_retry() { use std::sync::Arc; struct FailOnce(bool); - impl FailOnce { - fn prepare( - &mut self, - _: &WorkerSelectionContext<'_>, - _: WorkerCandidates<'_>, - ) -> Result<(), WorkerSelectionPolicyError> { - if std::mem::take(&mut self.0) { - return Err(WorkerSelectionPolicyError::failed("preparation failed")); - } - Ok(()) - } - } - impl WorkerScorer for FailOnce { fn score( &mut self, - context: &WorkerSelectionContext<'_>, - candidates: WorkerCandidates<'_>, + _: &WorkerSelectionContext<'_>, + _: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { - self.prepare(context, candidates)?; - for (_candidate, cost) in candidates.iter().zip(costs) { - *cost = 0.0; + if std::mem::take(&mut self.0) { + return Err(WorkerSelectionPolicyError::failed("score failed")); } + costs.fill(0.0); Ok(()) } } @@ -160,7 +131,7 @@ async fn preparation_error_does_not_book_and_allows_retry() { .select_and_reserve(reserve_request("request")) .await .unwrap_err(); - assert!(error.to_string().contains("preparation failed")); + assert!(error.to_string().contains("score failed")); assert!(!partition.scheduler().has_request("request")); core.select_and_reserve(reserve_request("request")) .await From 71ae532424f7d7c4e109521ef43c3e7e39c1a870 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 00:58:21 +0000 Subject: [PATCH 20/31] docs(router): fix plugin scoring and registry examples Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index ef37fdacf97f..7ac5651b1fcc 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -163,7 +163,7 @@ impl ActiveRequestsScorer { candidates: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { self.minimum = usize::MAX; - for candidate in candidates { + for candidate in candidates.iter() { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; @@ -395,13 +395,19 @@ The discovery-backed HTTP frontend and Python `KvRouter.generate` path install c The catalog registers plugin providers once. `RouterPluginRegistry::resolve_plugins` validates the selected configuration and returns a `RouterPlugins` bundle containing the configured factories. Both the HTTP frontend and Python router pass this bundle to the shared router construction path. Each router receives fresh plugin instances; factories run during construction, not for each request. -Rust embedders can use the same setup: +Rust embedders can use the same setup. From the embedding application's directory, add the builtin policy crate from the same Dynamo checkout: + +```bash +cargo add --path "$DYNAMO_DIR/lib/router-plugins/builtin" dynamo-custom-policy-builtin +``` + +Start with `default_registry()` so omitted worker roles use the builtin policy, then register the custom catalog: ```rust -use dynamo_kv_router::plugins::RouterPluginRegistry; +use dynamo_custom_policy_builtin::default_registry; use dynamo_llm::entrypoint::input::http::HttpFrontend; -let mut registry = RouterPluginRegistry::default(); +let mut registry = default_registry(); acme_routing_catalog::register(&mut registry)?; let plugins = registry.resolve_plugins(&kv_router_config)?; HttpFrontend::default() From c3c9fe9cfc98d483634dae42bb6c7cc24b584f3c Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 01:08:18 +0000 Subject: [PATCH 21/31] docs(router): hide builtin policy setup behind host API Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 22 ++++++------------- .../router/custom-policy-example/README.md | 2 +- 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 7ac5651b1fcc..d6d217a05b9e 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -41,7 +41,7 @@ Write your own policy when no shipped policy expresses the rule you need. The re | `WorkerPicker` | All eligible workers and their total costs | One row index | The rule samples, breaks ties, tracks policy state, or ignores total cost | | Policy factory | Router configuration, worker type, and routing partition | One `WorkerSelectionPolicy` | Prefill, decode, models, or routing groups need different components | -An external policy owns its filters, scorers, and picker. The default scorer and picker are private components of `dynamo-custom-policy-builtin`. They implement the same public traits as external policies. +An external policy owns its filters, scorers, and picker. Dynamo's default scorer and picker implement the same public traits as external policies. ## Public Plugin API and Compatibility @@ -395,19 +395,13 @@ The discovery-backed HTTP frontend and Python `KvRouter.generate` path install c The catalog registers plugin providers once. `RouterPluginRegistry::resolve_plugins` validates the selected configuration and returns a `RouterPlugins` bundle containing the configured factories. Both the HTTP frontend and Python router pass this bundle to the shared router construction path. Each router receives fresh plugin instances; factories run during construction, not for each request. -Rust embedders can use the same setup. From the embedding application's directory, add the builtin policy crate from the same Dynamo checkout: - -```bash -cargo add --path "$DYNAMO_DIR/lib/router-plugins/builtin" dynamo-custom-policy-builtin -``` - -Start with `default_registry()` so omitted worker roles use the builtin policy, then register the custom catalog: +Rust embedders using `dynamo-llm` start with the host's `router_plugin_registry()`, which includes Dynamo's default policy. Register the custom catalog before resolving the configuration. Plugin authors and frontend embedders do not need a direct dependency on the default policy's implementation crate: ```rust -use dynamo_custom_policy_builtin::default_registry; use dynamo_llm::entrypoint::input::http::HttpFrontend; +use dynamo_llm::kv_router::plugins::router_plugin_registry; -let mut registry = default_registry(); +let mut registry = router_plugin_registry(); acme_routing_catalog::register(&mut registry)?; let plugins = registry.resolve_plugins(&kv_router_config)?; HttpFrontend::default() @@ -650,9 +644,9 @@ The [example README](https://github.com/ai-dynamo/dynamo/blob/main/examples/rout ## Default Policy Plugin -The default worker selector lives in `dynamo-custom-policy-builtin`, beside the two-tier policy. Routing hosts install its default factory once, even without a custom catalog. An omitted policy or the reserved `default` selection uses that factory. Optional catalogs register named providers without replacing the default or existing classifiers. +Routing hosts install Dynamo's default policy automatically, even without a custom catalog. An omitted policy or the reserved `default` selection uses that policy. Optional catalogs register named providers without replacing the default or existing classifiers. -A Rust embedding that constructs `SelectionServiceBuilder` must supply a registry with a default factory, such as `default_registry()`, or an explicit policy factory. A missing factory fails construction before workers or synchronization resources start. Each routing partition receives its own policy instance. +Applications that construct the low-level `SelectionServiceBuilder` directly must supply a registry with a default factory or an explicit policy factory. A missing factory fails construction before workers or synchronization resources start. Each routing partition receives its own policy instance. Resolving `RouterPlugins` from the host registry includes the default factory even when no named policy is selected. `RouterPlugins::default()` remains an empty builder input. Use `has_custom_worker_selection()` to detect an explicit named or programmatic policy, and `has_custom_plugins()` to detect that policy or a classifier. Factory presence alone does not imply a custom selection. These checks preserve stock frontend modes and avoid waiting for a typed model card solely because the builtin factory is present. A classifier-only configuration still uses the builtin default for worker selection. @@ -676,6 +670,4 @@ The provider also accepts `overlap_score_credit_decay`, `decode_active_request_w The default scorer owns the cache and load formula, including preferred-taint costs. Its private `prepare` helper finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. -Direct Rust callers and replay import `DefaultWorkerSelector` from `dynamo_custom_policy_builtin`. Scheduler actors use `default_policy()` directly. The core's old selector is available only as a test and benchmark reference, not in production builds. - -The builtin crate keeps composition in `default/mod.rs`, the cost formula in `default/scorer.rs`, selection in `default/picker.rs`, startup parameters in `default/parameters.rs`, and the synchronous adapter in `default/selector.rs`. The two-tier policy remains in its own module. Preparation reads the existing snapshot; it adds no tracker synchronization. The routing core still projects load, selects and books serially, and processes request lifecycle updates. +Preparation reads the existing snapshot; it adds no tracker synchronization. The routing core still projects load, selects and books serially, and processes request lifecycle updates. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 45e940bf87bf..81f3b69dc677 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -383,4 +383,4 @@ Request-level cache information also requires `WorkerInputs::CACHE`: `context.ca Scorers receive a `WorkerCandidates` view with `iter()`, `get(row)`, `len()`, and `is_empty()`. Filters receive one borrowed `WorkerCandidate`. Both expose `cache()`, `load()`, and `preferred_taint_multiplier()` only when that component declares the corresponding `CACHE`, `LOAD`, or `PREFERRED_TAINT` group. The same rule applies to picker inputs. Another component requesting a group does not grant access. Views share the host's stored rows without allocating or copying a candidate table. -The builtin default now uses the public policy API too. Rust hosts install `dynamo_custom_policy_builtin::default_registry()` before adding custom providers. This supplies default selection for roles that the custom configuration does not replace. +Dynamo's default policy uses the public policy API too. The frontend and EPP install it automatically. Rust embedders using `dynamo-llm` start with `dynamo_llm::kv_router::plugins::router_plugin_registry()` before adding custom providers. This supplies default selection for roles that the custom configuration does not replace, without a direct dependency on the default policy's implementation crate. From 16d2c67a407097b911c799d83ec85db063ed9f90 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 01:26:38 +0000 Subject: [PATCH 22/31] refactor(router): simplify batch scoring example Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 43 +++++++------------ .../router/custom-policy-example/README.md | 4 +- .../simple-filter-score-pick/src/lib.rs | 2 +- .../simple-filter-score-pick/src/scorer.rs | 35 +++++---------- 4 files changed, 29 insertions(+), 55 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index d6d217a05b9e..0380c19a85a1 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -152,26 +152,7 @@ use dynamo_kv_router::plugins::worker_selection::{ }; /// Scores active requests above the least-loaded surviving candidate. -#[derive(Default)] -struct ActiveRequestsScorer { - minimum: usize, -} -impl ActiveRequestsScorer { - fn prepare( - &mut self, - _context: &WorkerSelectionContext<'_>, - candidates: WorkerCandidates<'_>, - ) -> Result<(), WorkerSelectionPolicyError> { - self.minimum = usize::MAX; - for candidate in candidates.iter() { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - self.minimum = self.minimum.min(load.active_requests()); - } - Ok(()) - } -} +struct ActiveRequestsScorer; impl WorkerScorer for ActiveRequestsScorer { /// Requests load inputs for the active-request count. @@ -179,26 +160,32 @@ impl WorkerScorer for ActiveRequestsScorer { WorkerInputs::LOAD } - /// Prepare the batch minimum and write one relative load cost per worker. + /// Find the batch minimum and write one relative load cost per worker. fn score( &mut self, - context: &WorkerSelectionContext<'_>, + _context: &WorkerSelectionContext<'_>, candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { - self.prepare(context, candidates)?; + let mut minimum = usize::MAX; + for candidate in candidates.iter() { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + minimum = minimum.min(load.active_requests()); + } for (candidate, cost) in candidates.iter().zip(costs) { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - *cost = (load.active_requests() - self.minimum) as f64; + *cost = (load.active_requests() - minimum) as f64; } Ok(()) } } ``` -Dynamo calls `score` once per scorer after filtering. The `WorkerCandidates` view borrows a fixed request snapshot; its order is unspecified and it cannot be retained. Use `iter()` to visit workers, `get(row)` to read one worker, and `len()` or `is_empty()` to inspect the batch. Each returned `WorkerCandidate` borrows the stored row and exposes only that scorer's declared inputs. Creating the view does not copy the worker table. The output slice has the same length and starts with NaN in every slot. Write every cost; an unwritten slot or a non-finite contribution or total stops selection before picking or booking. Empty sets skip scoring. Any preparation helper belongs to the scorer implementation and runs inside its `score` call. +Dynamo calls `score` once per scorer after filtering. The `WorkerCandidates` view borrows a fixed request snapshot; its order is unspecified and it cannot be retained. Use `iter()` to visit workers, `get(row)` to read one worker, and `len()` or `is_empty()` to inspect the batch. Each returned `WorkerCandidate` borrows the stored row and exposes only that scorer's declared inputs. Creating the view does not copy the worker table. The output slice has the same length and starts with NaN in every slot. Write every cost; an unwritten slot or a non-finite contribution or total stops selection before picking or booking. Empty sets skip scoring. Compute request-wide values, such as the minimum load, inside `score` before filling the output slice. A policy can stack multiple scorers. Dynamo calls them in declaration order and adds their costs. Dynamo rejects a non-finite contribution or total. @@ -270,7 +257,7 @@ fn provider( let filters: Vec> = vec![ Box::new(MinimumDeviceOverlapFilter { minimum_blocks }), ]; - let scorers: Vec> = vec![Box::new(ActiveRequestsScorer::default())]; + let scorers: Vec> = vec![Box::new(ActiveRequestsScorer)]; WorkerSelectionPolicy::new_with_filters( config.clone(), worker_type.as_str(), @@ -668,6 +655,6 @@ The provider also accepts `overlap_score_credit_decay`, `decode_active_request_w `WorkerCacheInput::accounting_cache_estimate()` returns the host's weighted block estimate and rounded token estimate. Missing estimates are zero. These estimates are not clamped to prompt length and do not measure physical GPU residency. The plugin can use them without changing the load that the host records. `WorkerLoadInput::is_available()` distinguishes a missing load observation from an observed idle worker. Both values describe the current selection snapshot. -The default scorer owns the cache and load formula, including preferred-taint costs. Its private `prepare` helper finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. +The default scorer owns the cache and load formula, including preferred-taint costs. Inside `score`, it finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. -Preparation reads the existing snapshot; it adds no tracker synchronization. The routing core still projects load, selects and books serially, and processes request lifecycle updates. +Scoring reads the existing snapshot; it adds no tracker synchronization. The routing core still projects load, selects and books serially, and processes request lifecycle updates. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 81f3b69dc677..5d764749bbb5 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -31,7 +31,7 @@ Preferred routing taints are optional candidate metadata. A filter, scorer, or p | `disagg-filter-score-pick` | Prefill and decode workers each need the complete policy flow | | `simple-stacked-score-pick` | Multiple scorer costs compose before one picker runs | -The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap, receives all surviving workers in one `score` call, and scores active requests above the least-loaded worker. A private preparation helper computes the minimum from that batch. The scorer writes one finite contribution per row into the host-owned output slice. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. +The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap, receives all surviving workers in one `score` call, and scores active requests above the least-loaded worker. The scorer computes the minimum inside `score` and keeps it local to that call. The scorer writes one finite contribution per row into the host-owned output slice. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. The [`soft-pin-repin` policy](soft-pin-repin/README.md) documents its load threshold, soft-binding behavior, and two-Mocker `A -> B -> B` walkthrough. @@ -110,7 +110,7 @@ fn provider( config.clone(), worker_type.as_str(), filters, - vec![Box::new(ActiveRequestsScorer::default())], + vec![Box::new(ActiveRequestsScorer)], Box::new(RequestAwarePicker), ) })) diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs index f9439a8de1ff..a63e0e6504d0 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/lib.rs @@ -56,7 +56,7 @@ fn provider( config.clone(), worker_type.as_str(), filters, - vec![Box::new(ActiveRequestsScorer::default())], + vec![Box::new(ActiveRequestsScorer)], Box::new(RequestAwarePicker), ) }, diff --git a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs index a7f4a551cc6b..13c3b3cfbd12 100644 --- a/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs +++ b/examples/router/custom-policy-example/simple-filter-score-pick/src/scorer.rs @@ -9,26 +9,7 @@ use dynamo_kv_router::plugins::worker_selection::{ }; /// Scores active requests above the least-loaded surviving candidate. -#[derive(Default)] -pub(crate) struct ActiveRequestsScorer { - minimum: usize, -} -impl ActiveRequestsScorer { - fn prepare( - &mut self, - _context: &WorkerSelectionContext<'_>, - candidates: WorkerCandidates<'_>, - ) -> Result<(), WorkerSelectionPolicyError> { - self.minimum = usize::MAX; - for candidate in candidates.iter() { - let load = candidate - .load() - .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - self.minimum = self.minimum.min(load.active_requests()); - } - Ok(()) - } -} +pub(crate) struct ActiveRequestsScorer; impl WorkerScorer for ActiveRequestsScorer { /// Requests load inputs for the active-request count. @@ -36,19 +17,25 @@ impl WorkerScorer for ActiveRequestsScorer { WorkerInputs::LOAD } - /// Prepare the batch minimum and write one relative load cost per worker. + /// Find the batch minimum and write one relative load cost per worker. fn score( &mut self, - context: &WorkerSelectionContext<'_>, + _context: &WorkerSelectionContext<'_>, candidates: WorkerCandidates<'_>, costs: &mut [f64], ) -> Result<(), WorkerSelectionPolicyError> { - self.prepare(context, candidates)?; + let mut minimum = usize::MAX; + for candidate in candidates.iter() { + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + minimum = minimum.min(load.active_requests()); + } for (candidate, cost) in candidates.iter().zip(costs) { let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - *cost = (load.active_requests() - self.minimum) as f64; + *cost = (load.active_requests() - minimum) as f64; } Ok(()) } From f137b6838eca3c82b1f0c87abb5f2962ed1901c9 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 01:33:31 +0000 Subject: [PATCH 23/31] docs(router): trim plugin guide to external contracts Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 52 ++++++------------- .../router/custom-policy-example/README.md | 14 +---- 2 files changed, 17 insertions(+), 49 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 0380c19a85a1..68bf99c9ef16 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -17,7 +17,7 @@ This feature replaces the worker-ranking part of Dynamo's routing pipeline. A `W flowchart LR Request["Incoming request"] --> Eligibility["Dynamo eligibility filters"] Eligibility --> Filters["WorkerFilter(s)
keep or reject each worker"] - Filters --> Scorers["WorkerScorer.score
all surviving workers, one cost per worker"] + Filters --> Scorers["WorkerScorer(s)
one cost per worker"] Scorers --> Sum["Dynamo sums scorer costs"] Sum --> Picker["WorkerPicker
one candidate row"] Picker --> Validate["Dynamo validates the row"] @@ -41,8 +41,7 @@ Write your own policy when no shipped policy expresses the rule you need. The re | `WorkerPicker` | All eligible workers and their total costs | One row index | The rule samples, breaks ties, tracks policy state, or ignores total cost | | Policy factory | Router configuration, worker type, and routing partition | One `WorkerSelectionPolicy` | Prefill, decode, models, or routing groups need different components | -An external policy owns its filters, scorers, and picker. Dynamo's default scorer and picker implement the same public traits as external policies. - +An external policy owns its filters, scorers, and picker. Dynamo's default scorer and picker are internal and can change with the built-in routing algorithm. ## Public Plugin API and Compatibility @@ -54,9 +53,9 @@ Use `dynamo_kv_router::plugins` for plugin registration and these modules for pl | `plugins::request_classifier` | `RequestClassifier`, `ClassifyRequest`, `ClassifyEvent`, configuration, and factories | | `plugins` | `RouterPluginRegistry` and the resolved `RouterPlugins` bundle | -The module organization changes no signal accessor, callback signature, signal meaning, or configuration format. Cache overlap, worker load, session metadata, and classification inputs retain their existing APIs. Dynamo owns scheduler queues, eligibility checks, reservations, lifecycle synchronization, and the default selector implementation. Low-level host interfaces such as `WorkerSelector` are used to connect Dynamo components; external plugins use `WorkerSelectionPolicy`. +Dynamo owns scheduler queues, eligibility checks, reservations, lifecycle synchronization, and the default selector implementation. Low-level host interfaces such as `WorkerSelector` are used to connect Dynamo components; external plugins use `WorkerSelectionPolicy`. -Existing plugin imports and the `WorkerSelectionPolicyRegistry` name remain available through compatibility re-exports until version 1.7. They refer to the same types and traits, so existing policies can register into `RouterPluginRegistry` without adapters. Migrate imports to `plugins::worker_selection` or `plugins::request_classifier`, rename the registry to `RouterPluginRegistry`, and use `register_worker_selection` for worker policies. The compatibility exports, registry alias, worker-only `register` method, and frontend `worker_selection_policy_factory` setter are scheduled for removal in 1.7. The signal methods themselves are unchanged. +Existing plugin imports and the `WorkerSelectionPolicyRegistry` name remain available through compatibility re-exports until version 1.7. They refer to the same types and traits, so existing policies can register into `RouterPluginRegistry` without adapters. Migrate imports to `plugins::worker_selection` or `plugins::request_classifier`, rename the registry to `RouterPluginRegistry`, and use `register_worker_selection` for worker policies. The compatibility exports, registry alias, worker-only `register` method, and frontend `worker_selection_policy_factory` setter are scheduled for removal in 1.7. ## Build the Policy @@ -144,23 +143,20 @@ Pass filters to `WorkerSelectionPolicy::new_with_filters`. If the policy has no A scorer expresses a preference without excluding a worker. It writes one finite cost for each worker into a host-owned output slice. Lower total cost is better by convention. -This scorer first finds the lowest active-request count among surviving workers, then scores each worker by the number of requests above that minimum. Subtracting the same minimum from each cost preserves the ranking: +This scorer ranks workers by active requests above the batch minimum: ```rust use dynamo_kv_router::plugins::worker_selection::{ WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicyError, }; -/// Scores active requests above the least-loaded surviving candidate. struct ActiveRequestsScorer; impl WorkerScorer for ActiveRequestsScorer { - /// Requests load inputs for the active-request count. fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::LOAD } - /// Find the batch minimum and write one relative load cost per worker. fn score( &mut self, _context: &WorkerSelectionContext<'_>, @@ -185,7 +181,7 @@ impl WorkerScorer for ActiveRequestsScorer { } ``` -Dynamo calls `score` once per scorer after filtering. The `WorkerCandidates` view borrows a fixed request snapshot; its order is unspecified and it cannot be retained. Use `iter()` to visit workers, `get(row)` to read one worker, and `len()` or `is_empty()` to inspect the batch. Each returned `WorkerCandidate` borrows the stored row and exposes only that scorer's declared inputs. Creating the view does not copy the worker table. The output slice has the same length and starts with NaN in every slot. Write every cost; an unwritten slot or a non-finite contribution or total stops selection before picking or booking. Empty sets skip scoring. Compute request-wide values, such as the minimum load, inside `score` before filling the output slice. +Dynamo calls `score` once per scorer after filtering. The output slice matches the candidate order and length. Write every entry with a finite cost; an unwritten or invalid cost stops selection before picking. A policy can stack multiple scorers. Dynamo calls them in declaration order and adds their costs. Dynamo rejects a non-finite contribution or total. @@ -382,7 +378,7 @@ The discovery-backed HTTP frontend and Python `KvRouter.generate` path install c The catalog registers plugin providers once. `RouterPluginRegistry::resolve_plugins` validates the selected configuration and returns a `RouterPlugins` bundle containing the configured factories. Both the HTTP frontend and Python router pass this bundle to the shared router construction path. Each router receives fresh plugin instances; factories run during construction, not for each request. -Rust embedders using `dynamo-llm` start with the host's `router_plugin_registry()`, which includes Dynamo's default policy. Register the custom catalog before resolving the configuration. Plugin authors and frontend embedders do not need a direct dependency on the default policy's implementation crate: +Rust embedders start with the host registry, which includes the default policy, then register their catalog: ```rust use dynamo_llm::entrypoint::input::http::HttpFrontend; @@ -420,15 +416,15 @@ The disaggregated prefill hop and query-only probes do not run classifiers. Stat Worker identity and ordinary request metadata are always available. Request-level cache context requires `WorkerInputs::CACHE`. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. -The combined `WorkerInputs` declarations also control cache setup. Without `CACHE`, the router creates no routing indexer, KV event subscription, or shared-cache client. With `CACHE`, cache settings choose event-backed, approximate, or remote indexing. Worker role, conditional disaggregation, and scoring weights do not enable or disable this setup. `serve_indexer` and `enable_session_prefix_index` require a policy that declares `CACHE`; incompatible configurations fail before cache resources start. +Routing cache resources start only if a component requests `CACHE`. `serve_indexer` and `enable_session_prefix_index` also require `CACHE`. Scoring weights do not control cache setup. ### Request Context | Accessor | Meaning | |---|---| -| `prompt_tokens()` | Exact incoming prompt token count, borrowed from the request without allocation | +| `prompt_tokens()` | Exact incoming prompt token count | | `request_blocks()` | Incoming prompt size in KV blocks, rounded up | -| `cache()` | With `CACHE`, returns a borrowed `RequestCacheInput`; otherwise returns `None` | +| `cache()` | Request cache data when the component declares `CACHE`; otherwise `None` | | `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 | @@ -438,16 +434,12 @@ The combined `WorkerInputs` declarations also control cache setup. Without `CACH | `strict_priority()` | Strict integer priority. The queue orders larger values first | | `router_temperature_override()` | Optional per-request router temperature override | -Each filter, scorer, and picker must declare `WorkerInputs::CACHE` in its own `required_worker_inputs()` to access `context.cache()`. Another component requesting cache does not enable it. The returned view exposes: +`context.cache()` exposes: | Accessor | Meaning | |---|---| | `has_tier_matches()` | Whether the lookup contains any GPU, CPU, or disk matches before filtering; `false` means only accounting estimates or no cache data were supplied | -| `shared_hits()` | Optional unweighted shared-cache ranges from the current lookup, with `total_hits` and `hits_beyond(prefix)` | - -A declared cache view exists even when no shared-cache result was supplied; in that case, `shared_hits()` returns `None`. The host borrows the existing lookup snapshot. Reading the view adds no cache lookup, range copy, or allocation. - -Calculate shared-cache credit inside the policy from `shared_hits()` and the worker's cache data. For example, call `hits_beyond(device_prefix)` to count shared blocks beyond the GPU prefix. The default policy uses its accounting estimate as the prefix when tier matches are absent. It performs this calculation only when shared-cache credit is enabled. +| `shared_hits()` | `None` if no shared result was supplied; otherwise unweighted ranges with `total_hits` and `hits_beyond(prefix)` | ### Session Context @@ -497,7 +489,7 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `CACHE` | `device_overlap_blocks()` | Device-resident prefix overlap in blocks | | `CACHE` | `host_overlap_blocks()` | Host-pinned prefix overlap in blocks | | `CACHE` | `disk_overlap_blocks()` | Disk prefix overlap in blocks | -| `CACHE` | `accounting_cache_estimate()` | Host accounting estimate as weighted KV blocks and rounded tokens; not clamped to prompt length | +| `CACHE` | `accounting_cache_estimate()` | Accounting estimate as weighted KV blocks and rounded tokens; zero if missing, not clamped to prompt length | | `LOAD` | `is_available()` | Whether the host supplied a load observation; distinguishes missing data from an observed idle worker | | `LOAD` | `active_prefill_tokens()` | Tokens currently active in the worker's prefill stage | | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks | @@ -506,7 +498,7 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi 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 requested groups through `WorkerCandidate::cache()` or `load()`. A picker reads index-aligned arrays through `WorkerInputView`. Each component must request every optional group that it reads. Candidate and context views enforce that declaration independently for every filter, scorer, and picker, including preferred-taint access. The host materializes the combined requirements once and shares borrowed views; it does not copy or rewrite every row for each scorer. +A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. A picker reads index-aligned arrays through `WorkerInputView`. 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. @@ -631,13 +623,7 @@ The [example README](https://github.com/ai-dynamo/dynamo/blob/main/examples/rout ## Default Policy Plugin -Routing hosts install Dynamo's default policy automatically, even without a custom catalog. An omitted policy or the reserved `default` selection uses that policy. Optional catalogs register named providers without replacing the default or existing classifiers. - -Applications that construct the low-level `SelectionServiceBuilder` directly must supply a registry with a default factory or an explicit policy factory. A missing factory fails construction before workers or synchronization resources start. Each routing partition receives its own policy instance. - -Resolving `RouterPlugins` from the host registry includes the default factory even when no named policy is selected. `RouterPlugins::default()` remains an empty builder input. Use `has_custom_worker_selection()` to detect an explicit named or programmatic policy, and `has_custom_plugins()` to detect that policy or a classifier. Factory presence alone does not imply a custom selection. These checks preserve stock frontend modes and avoid waiting for a typed model card solely because the builtin factory is present. A classifier-only configuration still uses the builtin default for worker selection. - -To configure a separate default-policy instance, use the named provider: +To give a worker role its own default-policy settings, select `dynamo-default-cost-fn`: ```yaml worker_selection: @@ -651,10 +637,4 @@ worker_selection: router_temperature: 0.0 ``` -The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. The default scorer declares `CACHE`, including when its device credit is zero; select a policy without `CACHE` to omit routing cache setup. The default plugin uses configured parameters and does not apply per-request score or temperature overrides. Plain disaggregated decode remains load-only when prefill tracking is off and conditional disaggregation is disabled. - -`WorkerCacheInput::accounting_cache_estimate()` returns the host's weighted block estimate and rounded token estimate. Missing estimates are zero. These estimates are not clamped to prompt length and do not measure physical GPU residency. The plugin can use them without changing the load that the host records. `WorkerLoadInput::is_available()` distinguishes a missing load observation from an observed idle worker. Both values describe the current selection snapshot. - -The default scorer owns the cache and load formula, including preferred-taint costs. Inside `score`, it finds the minimum prefill load used to reduce cache credit on busier workers. The picker reads only the final costs and owns temperature sampling and tie-breaking. The host checks scorer costs and the selected row before booking; diagnostics use the selected row's summed scorer cost. `with_exclusive_affinity(true)` asks the host to constrain candidates to an eligible session-affinity target; explicit request pins remain mandatory for every policy. - -Scoring reads the existing snapshot; it adds no tracker synchronization. The routing core still projects load, selects and books serially, and processes request lifecycle updates. +The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. Per-request score and temperature overrides do not apply to this policy. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 5d764749bbb5..0a6d7cfa41f4 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -31,7 +31,7 @@ Preferred routing taints are optional candidate metadata. A filter, scorer, or p | `disagg-filter-score-pick` | Prefill and decode workers each need the complete policy flow | | `simple-stacked-score-pick` | Multiple scorer costs compose before one picker runs | -The `simple-filter-score-pick` policy shows the complete pipeline. It filters on minimum device overlap, receives all surviving workers in one `score` call, and scores active requests above the least-loaded worker. The scorer computes the minimum inside `score` and keeps it local to that call. The scorer writes one finite contribution per row into the host-owned output slice. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. +The `simple-filter-score-pick` policy filters on minimum device overlap and scores active requests above the least-loaded worker. Its picker normally selects the lowest cost. Tool-result turns select the worker with the most device overlap through `session_context().input_trigger()`. The [`soft-pin-repin` policy](soft-pin-repin/README.md) documents its load threshold, soft-binding behavior, and two-Mocker `A -> B -> B` walkthrough. @@ -372,15 +372,3 @@ python3 -m dynamo.mocker \ ``` Send the same `curl` request from a fourth terminal. The frontend log records separate prefill and decode selections. Each worker set runs its own filter, scorer, and picker. - -## Exact Prompt Length - -Filters, scorers, and pickers can read `context.prompt_tokens()` for the exact input token count. This returns the existing request value without a prompt copy or allocation. Use `context.request_blocks()` when the policy needs rounded KV blocks instead. For example, at 16 tokens per block, a 17-token prompt has 2 blocks. - -The union of component declarations controls routing cache setup. Without `WorkerInputs::CACHE`, no routing indexer, KV event subscription, or shared-cache client is started, regardless of score weights or worker role. With `CACHE`, router settings select event-backed, approximate, or remote indexing. Serving an indexer or enabling the session prefix index requires a policy that declares `CACHE`. - -Request-level cache information also requires `WorkerInputs::CACHE`: `context.cache()` returns `None` without that declaration. With it, the borrowed view exposes `has_tier_matches()` and `shared_hits()`. The former is false when the snapshot has no tier matches; the latter is `None` when no shared-cache result was supplied. A policy can call `shared_hits().map(|hits| hits.hits_beyond(prefix))` to calculate shared hits after its chosen prefix, only when needed. - -Scorers receive a `WorkerCandidates` view with `iter()`, `get(row)`, `len()`, and `is_empty()`. Filters receive one borrowed `WorkerCandidate`. Both expose `cache()`, `load()`, and `preferred_taint_multiplier()` only when that component declares the corresponding `CACHE`, `LOAD`, or `PREFERRED_TAINT` group. The same rule applies to picker inputs. Another component requesting a group does not grant access. Views share the host's stored rows without allocating or copying a candidate table. - -Dynamo's default policy uses the public policy API too. The frontend and EPP install it automatically. Rust embedders using `dynamo-llm` start with `dynamo_llm::kv_router::plugins::router_plugin_registry()` before adding custom providers. This supplies default selection for roles that the custom configuration does not replace, without a direct dependency on the default policy's implementation crate. From 3584802a8a40f2396f9535c673f88135a67f55ff Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 02:23:37 +0000 Subject: [PATCH 24/31] refactor(router): expose cache facts through worker inputs Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 14 +- .../router/custom-policy-example/README.md | 2 + lib/kv-router/src/lib.rs | 2 +- lib/kv-router/src/plugins/worker_selection.rs | 8 +- .../src/plugins/worker_selection/context.rs | 57 +------- .../src/plugins/worker_selection/inputs.rs | 131 +++++++++++++++--- lib/kv-router/src/scheduling/selector/mod.rs | 31 +++-- .../src/scheduling/selector/policy.rs | 32 +++-- .../src/services/selection/core/tests.rs | 9 +- lib/llm/src/kv_router.rs | 15 +- .../builtin/src/default/scorer.rs | 12 +- .../builtin/src/two_tier_cost_fn.rs | 4 +- .../builtin/tests/batch_scoring.rs | 5 +- .../{cache_context.rs => cache_inputs.rs} | 71 ++++++---- .../builtin/tests/default_policy.rs | 13 +- 15 files changed, 231 insertions(+), 175 deletions(-) rename lib/router-plugins/builtin/tests/{cache_context.rs => cache_inputs.rs} (74%) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 68bf99c9ef16..f7e57f1ada39 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -414,7 +414,7 @@ The disaggregated prefill hop and query-only probes do not run classifiers. Stat ## Available Signals -Worker identity and ordinary request metadata are always available. Request-level cache context requires `WorkerInputs::CACHE`. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. +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`. Scoring weights do not control cache setup. @@ -424,7 +424,6 @@ Routing cache resources start only if a component requests `CACHE`. `serve_index |---|---| | `prompt_tokens()` | Exact incoming prompt token count | | `request_blocks()` | Incoming prompt size in KV blocks, rounded up | -| `cache()` | Request cache data when the component declares `CACHE`; otherwise `None` | | `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 | @@ -434,13 +433,6 @@ Routing cache resources start only if a component requests `CACHE`. `serve_index | `strict_priority()` | Strict integer priority. The queue orders larger values first | | `router_temperature_override()` | Optional per-request router temperature override | -`context.cache()` exposes: - -| Accessor | Meaning | -|---|---| -| `has_tier_matches()` | Whether the lookup contains any GPU, CPU, or disk matches before filtering; `false` means only accounting estimates or no cache data were supplied | -| `shared_hits()` | `None` if no shared result was supplied; otherwise unweighted ranges with `total_hits` and `hits_beyond(prefix)` | - ### Session Context `session_context()` returns `None` when the request has no session metadata. This policy-facing view contains selected session metadata; it is not Dynamo's internal request envelope. When present, it provides: @@ -490,6 +482,8 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `CACHE` | `host_overlap_blocks()` | Host-pinned prefix overlap in blocks | | `CACHE` | `disk_overlap_blocks()` | Disk prefix overlap in blocks | | `CACHE` | `accounting_cache_estimate()` | Accounting estimate as weighted KV blocks and rounded tokens; zero if missing, not clamped to prompt length | +| `CACHE` | `has_tier_matches()` | Whether the lookup contains any GPU, CPU, or disk matches before filtering; `false` means only accounting estimates or no cache data were supplied | +| `CACHE` | `shared_hits()` | `None` if no shared result was supplied; otherwise unweighted ranges with `total_hits` and `hits_beyond(prefix)` | | `LOAD` | `is_available()` | Whether the host supplied a load observation; distinguishes missing data from an observed idle worker | | `LOAD` | `active_prefill_tokens()` | Tokens currently active in the worker's prefill stage | | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks | @@ -498,7 +492,7 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi 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 arrays through `WorkerInputView`. Each component must declare every group it reads; another component requesting a group does not grant access. +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. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 0a6d7cfa41f4..7b26c598bc62 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -20,6 +20,8 @@ policy crate -> catalog crate -> router-policy YAML -> frontend or EPP binary Dynamo owns discovery, eligibility, queueing, validation, reservations, accounting, and metrics. A policy sees only eligible workers and returns one candidate row. +Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (filter/scorer) or `input.cache()` (picker); otherwise these return `None`. Each cache row includes worker overlaps, `has_tier_matches()`, and `shared_hits()`. + 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. ## Pick a Starting Point diff --git a/lib/kv-router/src/lib.rs b/lib/kv-router/src/lib.rs index 921049a584a2..0e2580e861ac 100644 --- a/lib/kv-router/src/lib.rs +++ b/lib/kv-router/src/lib.rs @@ -81,7 +81,7 @@ pub use selector::DefaultWorkerSelector; pub use selector::{WorkerSelectionInput, WorkerSelector}; // TODO(v1.7): Remove these compatibility re-exports; use crate::plugins instead. pub use plugins::worker_selection::{ - RequestCacheInput, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, + ScoredWorkerCandidate, WorkerCacheInput, WorkerCacheInputs, WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, }; diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index 22b9a0a23a7f..6d50a3f8f477 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -7,11 +7,11 @@ mod config; mod context; mod inputs; -pub use context::{RequestCacheInput, WorkerSelectionContext}; -pub(crate) use inputs::CandidateData; +pub use context::WorkerSelectionContext; +pub(crate) use inputs::{CacheSnapshot, CandidateData, WorkerCacheData}; pub use inputs::{ - ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, WorkerInputView, - WorkerInputs, WorkerLoadInput, + ScoredWorkerCandidate, WorkerCacheInput, WorkerCacheInputs, WorkerCandidate, WorkerCandidates, + WorkerInputView, WorkerInputs, WorkerLoadInput, }; pub use super::registry::{ diff --git a/lib/kv-router/src/plugins/worker_selection/context.rs b/lib/kv-router/src/plugins/worker_selection/context.rs index 6c50dbae98d0..0b4c0cafc09c 100644 --- a/lib/kv-router/src/plugins/worker_selection/context.rs +++ b/lib/kv-router/src/plugins/worker_selection/context.rs @@ -1,10 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Request metadata and the borrowed, CACHE-gated lookup snapshot. +//! Request metadata available to worker-selection components. -use super::{SessionContext, WorkerInputs}; -use crate::protocols::{SharedCacheHits, WorkerAffinityTarget, WorkerWithDpRank}; +use super::SessionContext; +use crate::protocols::{WorkerAffinityTarget, WorkerWithDpRank}; use crate::scheduling::SchedulingRequest; /// Request-level values available to custom filters, scorers, and pickers. @@ -13,36 +13,10 @@ pub struct WorkerSelectionContext<'a> { pub(crate) request_blocks: u64, pub(crate) block_size: u32, pub(crate) track_prefill_tokens: bool, - pub(crate) has_tier_matches: bool, - pub(crate) inputs: WorkerInputs, pub(crate) pinned_worker: Option, pub(crate) router_temperature_override: Option, } -/// Request-wide cache facts from the host's current lookup snapshot. -/// Borrowed for the selection callback; no worker data is copied or looked up here. -#[derive(Clone, Copy)] -pub struct RequestCacheInput<'a> { - shared_hits: Option<&'a SharedCacheHits>, - has_tier_matches: bool, -} - -impl<'a> RequestCacheInput<'a> { - /// Unweighted shared-cache ranges in KV block positions, or None if no result was supplied. - /// The host owns this snapshot. It does not change during the callback, and may lag engine state. - /// Use `hits_beyond(prefix)` to exclude hits already covered by the policy's chosen prefix. - pub fn shared_hits(self) -> Option<&'a SharedCacheHits> { - self.shared_hits - } - - /// Whether the request snapshot contains any tier-specific matches before worker filtering. - /// False means only accounting estimates (or no cache data) were supplied. This describes - /// observed matches, not worker cache capacity or whether a particular worker has a match. - pub fn has_tier_matches(self) -> bool { - self.has_tier_matches - } -} - impl WorkerSelectionContext<'_> { /// The exact worker/rank imposed by the host for this selection, if any. /// Includes explicit pins and eligible exclusive-affinity targets. This is @@ -57,31 +31,6 @@ impl WorkerSelectionContext<'_> { self.request.isl_tokens } - /// Borrow request-wide cache facts when this component declared [`WorkerInputs::CACHE`]. - /// This is the existing lookup snapshot; accessing it adds no lookup or allocation. - pub fn cache(&self) -> Option> { - self.inputs - .contains(WorkerInputs::CACHE) - .then_some(RequestCacheInput { - shared_hits: self.request.shared_cache_hits.as_ref(), - has_tier_matches: self.has_tier_matches, - }) - } - - /// Restrict request-level signals to this component's startup input declaration. - pub(crate) fn with_inputs(&self, inputs: WorkerInputs) -> Self { - Self { - request: self.request, - request_blocks: self.request_blocks, - block_size: self.block_size, - track_prefill_tokens: self.track_prefill_tokens, - has_tier_matches: self.has_tier_matches, - inputs, - pinned_worker: self.pinned_worker, - router_temperature_override: self.router_temperature_override, - } - } - /// Return the incoming prompt size in KV blocks. pub fn request_blocks(&self) -> u64 { self.request_blocks diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index 6466997449ea..4ce3dd6aaf4c 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -3,14 +3,14 @@ //! Worker input groups, stored snapshots, and component-scoped borrowed views. -use crate::protocols::WorkerWithDpRank; +use crate::protocols::{SharedCacheHits, WorkerWithDpRank}; use std::ops::BitOr; /// Host-owned row materialized once for the union of component input requirements. pub(crate) struct CandidateData { pub(crate) worker: WorkerWithDpRank, pub(crate) inputs: WorkerInputs, - pub(crate) cache: WorkerCacheInput, + pub(crate) cache: WorkerCacheData, pub(crate) load: WorkerLoadInput, pub(crate) preferred_taint_multiplier: Option, } @@ -20,11 +20,20 @@ pub(crate) struct CandidateData { pub struct WorkerCandidate<'a> { data: &'a CandidateData, inputs: WorkerInputs, + cache_snapshot: &'a CacheSnapshot<'a>, } impl<'a> WorkerCandidate<'a> { - pub(crate) fn new(data: &'a CandidateData, inputs: WorkerInputs) -> Self { - Self { data, inputs } + pub(crate) fn new( + data: &'a CandidateData, + inputs: WorkerInputs, + cache_snapshot: &'a CacheSnapshot<'a>, + ) -> Self { + Self { + data, + inputs, + cache_snapshot, + } } /// Worker identity is always available. @@ -33,10 +42,13 @@ impl<'a> WorkerCandidate<'a> { } /// Read this worker's cache snapshot only when this component declared CACHE. - pub fn cache(self) -> Option<&'a WorkerCacheInput> { + pub fn cache(self) -> Option> { self.inputs .contains(WorkerInputs::CACHE) - .then_some(&self.data.cache) + .then_some(WorkerCacheInput { + data: &self.data.cache, + snapshot: self.cache_snapshot, + }) } /// Read this worker's load snapshot only when this component declared LOAD. @@ -64,11 +76,20 @@ impl<'a> WorkerCandidate<'a> { pub struct WorkerCandidates<'a> { rows: &'a [CandidateData], inputs: WorkerInputs, + cache_snapshot: &'a CacheSnapshot<'a>, } impl<'a> WorkerCandidates<'a> { - pub(crate) fn new(rows: &'a [CandidateData], inputs: WorkerInputs) -> Self { - Self { rows, inputs } + pub(crate) fn new( + rows: &'a [CandidateData], + inputs: WorkerInputs, + cache_snapshot: &'a CacheSnapshot<'a>, + ) -> Self { + Self { + rows, + inputs, + cache_snapshot, + } } /// Number of surviving workers, equal to the scorer's cost-buffer length. @@ -85,7 +106,7 @@ impl<'a> WorkerCandidates<'a> { pub fn get(self, row: usize) -> Option> { self.rows .get(row) - .map(|data| WorkerCandidate::new(data, self.inputs)) + .map(|data| WorkerCandidate::new(data, self.inputs, self.cache_snapshot)) } /// Iterate over the surviving workers without copying their data. @@ -94,7 +115,7 @@ impl<'a> WorkerCandidates<'a> { ) -> impl ExactSizeIterator> + DoubleEndedIterator + Clone { self.rows .iter() - .map(move |data| WorkerCandidate::new(data, self.inputs)) + .map(move |data| WorkerCandidate::new(data, self.inputs, self.cache_snapshot)) } } @@ -113,7 +134,7 @@ pub struct WorkerInputs(u8); impl WorkerInputs { /// Request no optional worker inputs. pub const NONE: Self = Self(0); - /// Request worker KV-cache overlap inputs and request-level cache context. + /// Request worker KV-cache overlaps and shared-cache lookup results. pub const CACHE: Self = Self(1 << 0); /// Request active-load inputs. pub const LOAD: Self = Self(1 << 1); @@ -141,9 +162,15 @@ impl BitOr for WorkerInputs { } } -/// KV-cache overlap values for one worker. +/// Request-wide facts borrowed once for all cache views in a selection. +pub(crate) struct CacheSnapshot<'a> { + pub(crate) shared_hits: Option<&'a SharedCacheHits>, + pub(crate) has_tier_matches: bool, +} + +/// Numeric cache row retained in host buffers between selections. #[derive(Clone, Copy, Default)] -pub struct WorkerCacheInput { +pub(crate) struct WorkerCacheData { pub(crate) effective_overlap_blocks: f64, pub(crate) estimated_cached_tokens: usize, pub(crate) device_overlap_blocks: f64, @@ -151,6 +178,51 @@ pub struct WorkerCacheInput { pub(crate) disk_overlap_blocks: f64, } +/// Borrowed worker overlaps and shared-cache facts from one lookup snapshot. +/// Available only to components that declare [`WorkerInputs::CACHE`]. +#[derive(Clone, Copy)] +pub struct WorkerCacheInput<'a> { + data: &'a WorkerCacheData, + snapshot: &'a CacheSnapshot<'a>, +} + +/// Borrowed cache rows in the same order as a picker's candidates. +#[derive(Clone, Copy)] +pub struct WorkerCacheInputs<'a> { + pub(crate) rows: &'a [WorkerCacheData], + pub(crate) snapshot: &'a CacheSnapshot<'a>, +} + +impl<'a> WorkerCacheInputs<'a> { + /// Number of cache rows, equal to the picker candidate count. + pub fn len(self) -> usize { + self.rows.len() + } + + /// Whether there are no cache rows. + pub fn is_empty(self) -> bool { + self.rows.is_empty() + } + + /// Borrow one cache row, or None for an out-of-range index. + pub fn get(self, row: usize) -> Option> { + self.rows.get(row).map(|data| WorkerCacheInput { + data, + snapshot: self.snapshot, + }) + } + + /// Iterate over cache rows without copying their data. + pub fn iter( + self, + ) -> impl ExactSizeIterator> + DoubleEndedIterator + Clone { + self.rows.iter().map(move |data| WorkerCacheInput { + data, + snapshot: self.snapshot, + }) + } +} + /// Active-load values for one worker. #[derive(Clone, Copy, Default)] pub struct WorkerLoadInput { @@ -164,7 +236,7 @@ pub struct WorkerLoadInput { #[derive(Clone, Copy)] pub struct WorkerInputView<'a> { pub(crate) candidates: &'a [ScoredWorkerCandidate], - pub(crate) cache: Option<&'a [WorkerCacheInput]>, + pub(crate) cache: Option>, pub(crate) load: Option<&'a [WorkerLoadInput]>, } @@ -181,7 +253,7 @@ impl CandidateData { additional.cache } } else { - WorkerCacheInput::default() + WorkerCacheData::default() }, load: if inputs.contains(WorkerInputs::LOAD) { if self.inputs.contains(WorkerInputs::LOAD) { @@ -217,28 +289,45 @@ impl ScoredWorkerCandidate { } } -impl WorkerCacheInput { +impl<'a> WorkerCacheInput<'a> { + /// Unweighted shared-cache ranges in KV block positions, or None if no result was supplied. + /// The host owns this snapshot. It does not change during the callback, and may lag engine state. + /// Use `hits_beyond(prefix)` to exclude hits already covered by the policy's chosen prefix. + pub fn shared_hits(self) -> Option<&'a SharedCacheHits> { + self.snapshot.shared_hits + } + + /// Whether the request snapshot contains any tier-specific matches before worker filtering. + /// False means only accounting estimates (or no cache data) were supplied. This describes + /// observed matches, not worker cache capacity or whether a particular worker has a match. + pub fn has_tier_matches(self) -> bool { + self.snapshot.has_tier_matches + } + /// Host accounting estimate for this worker, in weighted KV blocks and rounded tokens. /// Lower-tier matches use the host's cache weights. Missing estimates are zero; /// neither value is clamped to prompt length. This is the current lookup snapshot, /// not a count of physically resident GPU tokens. Policy scores do not alter it. pub fn accounting_cache_estimate(&self) -> (f64, usize) { - (self.effective_overlap_blocks, self.estimated_cached_tokens) + ( + self.data.effective_overlap_blocks, + self.data.estimated_cached_tokens, + ) } /// Return device-resident prefix overlap in KV blocks. pub fn device_overlap_blocks(&self) -> f64 { - self.device_overlap_blocks + self.data.device_overlap_blocks } /// Return host-pinned prefix overlap in KV blocks. pub fn host_overlap_blocks(&self) -> f64 { - self.host_overlap_blocks + self.data.host_overlap_blocks } /// Return disk prefix overlap in KV blocks. pub fn disk_overlap_blocks(&self) -> f64 { - self.disk_overlap_blocks + self.data.disk_overlap_blocks } } @@ -272,7 +361,7 @@ impl<'a> WorkerInputView<'a> { } /// Return index-aligned KV-cache inputs when the picker requested them. - pub fn cache(self) -> Option<&'a [WorkerCacheInput]> { + pub fn cache(self) -> Option> { self.cache } diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 403bc92e629f..27918d598ecd 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -12,14 +12,14 @@ pub use reference::DefaultWorkerSelector; // TODO(v1.7): Remove these compatibility re-exports; use crate::plugins instead. pub use crate::plugins::worker_selection::{ - RequestCacheInput, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, + ScoredWorkerCandidate, WorkerCacheInput, WorkerCacheInputs, WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, }; #[cfg(any(test, feature = "bench"))] use reference::{DefaultWorkerPicker, DefaultWorkerScorer}; -use crate::plugins::worker_selection::CandidateData; +use crate::plugins::worker_selection::{CacheSnapshot, CandidateData, WorkerCacheData}; pub use policy::WorkerSelectionPolicy; use policy::{ComposedPolicyState, WorkerSelectionPolicyStateRef, collect_policy_candidates}; #[cfg(any(test, feature = "bench"))] @@ -131,22 +131,25 @@ impl<'a, C: WorkerConfigLike> WorkerSelectionInput<'a, C> { struct MaterializedSelectionInput<'a> { request: &'a SchedulingRequest, context: WorkerSelectionContext<'a>, + cache_snapshot: CacheSnapshot<'a>, } impl<'a> MaterializedSelectionInput<'a> { fn new(request: &'a SchedulingRequest, block_size: u32) -> Self { Self { request, + cache_snapshot: CacheSnapshot { + shared_hits: request.shared_cache_hits.as_ref(), + has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() + || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() + || !request.overlap.tier_overlap_blocks.disk.is_empty(), + }, context: WorkerSelectionContext { request, request_blocks: request.request_blocks(block_size), block_size, track_prefill_tokens: request.track_prefill_tokens, - inputs: WorkerInputs::NONE, pinned_worker: request.pinned_worker, - has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() - || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() - || !request.overlap.tier_overlap_blocks.disk.is_empty(), router_temperature_override: request .router_config_override .as_ref() @@ -201,7 +204,7 @@ impl<'a> MaterializedSelectionInput<'a> { .unwrap_or(0.0); let device_overlap_blocks = select_device_overlap(effective_overlap_blocks, reported_device_overlap_blocks); - WorkerCacheInput { + WorkerCacheData { effective_overlap_blocks, estimated_cached_tokens: cached_tokens, device_overlap_blocks, @@ -223,7 +226,7 @@ impl<'a> MaterializedSelectionInput<'a> { .unwrap_or(0) as f64, } } else { - WorkerCacheInput::default() + WorkerCacheData::default() }; let load = if inputs.contains(WorkerInputs::LOAD) { let available = worker_load.is_some(); @@ -394,15 +397,17 @@ fn select_worker_with_policy( ); let picker_input = WorkerInputView { candidates, - cache: picker_inputs - .contains(WorkerInputs::CACHE) - .then_some(cache_inputs.as_slice()), + cache: picker_inputs.contains(WorkerInputs::CACHE).then_some( + WorkerCacheInputs { + rows: cache_inputs, + snapshot: &input.cache_snapshot, + }, + ), load: picker_inputs .contains(WorkerInputs::LOAD) .then_some(load_inputs.as_slice()), }; - let context = input.context.with_inputs(*picker_inputs); - let row = picker.pick(&context, picker_input)?; + let row = picker.pick(&input.context, picker_input)?; let Some(candidate) = candidates.get(row) else { return Err(WorkerSelectionPolicyError::InvalidPickerRow { row, diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index e0e7a06ddb8a..8d6feeae81fe 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -16,8 +16,8 @@ use crate::scheduling::filter::RoutingEligibility; use crate::scheduling::types::{KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError}; use crate::plugins::worker_selection::{ - CandidateData, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerCandidates, - WorkerFilter, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, + CacheSnapshot, CandidateData, ScoredWorkerCandidate, WorkerCacheData, WorkerCandidate, + WorkerCandidates, WorkerFilter, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, }; @@ -45,7 +45,7 @@ pub(super) struct ComposedPolicyState { unscored_candidates: Vec, score_contributions: Vec, pub(super) candidates: Vec, - pub(super) cache_inputs: Vec, + pub(super) cache_inputs: Vec, pub(super) load_inputs: Vec, } @@ -151,7 +151,7 @@ fn push_picker_candidate( cost: f64, picker_inputs: WorkerInputs, candidates: &mut Vec, - cache_inputs: &mut Vec, + cache_inputs: &mut Vec, load_inputs: &mut Vec, ) { candidates.push(ScoredWorkerCandidate { @@ -194,6 +194,7 @@ impl ComposedPolicyState { fn score_candidates( &mut self, context: &WorkerSelectionContext<'_>, + cache_snapshot: &CacheSnapshot<'_>, ) -> Result<(), KvSchedulerError> { let Self { scorers, @@ -209,10 +210,9 @@ impl ComposedPolicyState { score_contributions.resize(candidates.len(), f64::NAN); for (scorer_index, (inputs, scorer)) in scorers.iter_mut().enumerate() { score_contributions.fill(f64::NAN); - let context = context.with_inputs(*inputs); scorer.score( - &context, - WorkerCandidates::new(unscored_candidates, *inputs), + context, + WorkerCandidates::new(unscored_candidates, *inputs, cache_snapshot), score_contributions, )?; for (row, (contribution, scored)) in score_contributions @@ -265,7 +265,7 @@ pub(super) fn collect_policy_candidates( state.push_candidate(candidate); false }); - state.score_candidates(&input.context)?; + state.score_candidates(&input.context, &input.cache_snapshot)?; return Ok(!state.candidates.is_empty()); } @@ -274,8 +274,6 @@ pub(super) fn collect_policy_candidates( && state.filter_inputs.contains(WorkerInputs::PREFERRED_TAINT); let materialize_additional_preferred_taint = eligibility.pinned_worker().is_none() && additional_inputs.contains(WorkerInputs::PREFERRED_TAINT); - // Reuse the borrowed context across rows; only the current filter's access mask changes. - let mut filter_context = input.context.with_inputs(WorkerInputs::NONE); let mut has_eligible_worker = false; let mut error = None; eligibility.any_eligible_worker_rank(workers, |worker, config| { @@ -293,10 +291,9 @@ pub(super) fn collect_policy_candidates( state.filter_inputs, ); for (inputs, filter) in &mut state.filters { - filter_context.inputs = *inputs; match filter.keep( - &filter_context, - WorkerCandidate::new(&filter_candidate, *inputs), + &input.context, + WorkerCandidate::new(&filter_candidate, *inputs, &input.cache_snapshot), ) { Ok(true) => {} Ok(false) => return false, @@ -326,7 +323,7 @@ pub(super) fn collect_policy_candidates( if let Some(error) = error { return Err(error); } - state.score_candidates(&input.context)?; + state.score_candidates(&input.context, &input.cache_snapshot)?; Ok(has_eligible_worker) } @@ -672,7 +669,12 @@ mod tests { input: WorkerInputView<'_>, ) -> Result { assert_eq!( - input.cache().expect("cache input")[0].device_overlap_blocks(), + input + .cache() + .expect("cache input") + .get(0) + .unwrap() + .device_overlap_blocks(), 0.0 ); Ok(0) diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 0735f11710d6..ee7319628181 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -704,12 +704,9 @@ impl crate::scheduling::selector::WorkerPicker for CapturingPicker { .expect("CACHE inputs requested") .iter() .map(|cache| { - context - .cache() - .and_then(|request| request.shared_hits()) - .map_or(0, |hits| { - hits.hits_beyond(cache.device_overlap_blocks().round().max(0.0) as u32) - }) + cache.shared_hits().map_or(0, |hits| { + hits.hits_beyond(cache.device_overlap_blocks().round().max(0.0) as u32) + }) }) .collect(), }); diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 60b7e5fea42e..eb0b61dbbcca 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -2174,17 +2174,10 @@ mod tests { .iter() .position(|candidate| candidate.worker() == self.worker) .ok_or_else(|| WorkerSelectionPolicyError::failed("fixed worker not eligible"))?; - let shared = _context - .cache() - .and_then(|cache| cache.shared_hits()) - .map(|hits| { - hits.hits_beyond( - input.cache().unwrap()[row] - .device_overlap_blocks() - .round() - .max(0.0) as u32, - ) - }) + let cache = input.cache().unwrap().get(row).unwrap(); + let shared = cache + .shared_hits() + .map(|hits| hits.hits_beyond(cache.device_overlap_blocks().round().max(0.0) as u32)) .filter(|blocks| *blocks > 0); assert_eq!(shared, self.expected_shared_blocks); Ok(row) diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 3268f3330932..ab0a5c801d32 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -61,7 +61,6 @@ struct DefaultScorer { #[derive(Default)] struct PreparedRequest { min_prefill: usize, - has_tier_matches: bool, block_size: IntegerDivisor, request_blocks: IntegerDivisor, overlap_credit: f64, @@ -140,10 +139,6 @@ impl }; self.prepared = PreparedRequest { min_prefill: 0, - has_tier_matches: context - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache context unavailable"))? - .has_tier_matches(), block_size: IntegerDivisor::new(u64::from(context.block_size())), request_blocks: IntegerDivisor::new(context.request_blocks()), overlap_credit, @@ -177,15 +172,14 @@ impl .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); - let device = if self.prepared.has_tier_matches { + let device = if cache.has_tier_matches() { cache.device_overlap_blocks() } else { estimated_overlap }; let shared_credit = if SHARED_CREDIT { - let shared = context - .cache() - .and_then(|cache| cache.shared_hits()) + let shared = cache + .shared_hits() .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)); self.shared_cache_multiplier * shared as f64 } else { diff --git a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs index f4a2a3f3bf37..e77eda8b120c 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs @@ -31,7 +31,7 @@ use std::sync::Arc; use dynamo_kv_router::KvRouterConfig; use dynamo_kv_router::plugins::worker_selection::{ - WorkerCacheInput, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, + WorkerCacheInputs, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelectionPolicyFactory, }; @@ -97,7 +97,7 @@ fn least_loaded(load: &[WorkerLoadInput], rows: impl Iterator) -> fn select_row( parameters: &Parameters, - cache: &[WorkerCacheInput], + cache: WorkerCacheInputs<'_>, load: &[WorkerLoadInput], request_blocks: u64, ) -> Option { diff --git a/lib/router-plugins/builtin/tests/batch_scoring.rs b/lib/router-plugins/builtin/tests/batch_scoring.rs index c083d1e99431..16c0131a14e6 100644 --- a/lib/router-plugins/builtin/tests/batch_scoring.rs +++ b/lib/router-plugins/builtin/tests/batch_scoring.rs @@ -272,7 +272,10 @@ fn picker_columns_and_costs_stay_aligned_after_a_scoring_error() { for (row, candidate) in input.candidates().iter().enumerate() { let worker = candidate.worker(); let expected_overlap = (worker.worker_id * 7 + u64::from(worker.dp_rank)) % 9; - assert_eq!(cache[row].device_overlap_blocks(), expected_overlap as f64); + assert_eq!( + cache.get(row).unwrap().device_overlap_blocks(), + expected_overlap as f64 + ); assert_eq!(load[row].active_requests(), worker.worker_id as usize % 5); assert_eq!(candidate.cost(), load[row].active_requests() as f64); } diff --git a/lib/router-plugins/builtin/tests/cache_context.rs b/lib/router-plugins/builtin/tests/cache_inputs.rs similarity index 74% rename from lib/router-plugins/builtin/tests/cache_context.rs rename to lib/router-plugins/builtin/tests/cache_inputs.rs index 3be3299f8ed4..f23ad8fcfdeb 100644 --- a/lib/router-plugins/builtin/tests/cache_context.rs +++ b/lib/router-plugins/builtin/tests/cache_inputs.rs @@ -4,9 +4,9 @@ mod support; use dynamo_kv_router::{ - KvRouterConfig, SharedCacheHits, WorkerCandidate, WorkerCandidates, WorkerFilter, - WorkerInputView, WorkerInputs, WorkerPicker, WorkerScorer, WorkerSelectionContext, - WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, + KvRouterConfig, SharedCacheHits, WorkerCacheInput, WorkerCandidate, WorkerCandidates, + WorkerFilter, WorkerInputView, WorkerInputs, WorkerPicker, WorkerScorer, + WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelector, }; use std::sync::{ Arc, @@ -16,7 +16,8 @@ use std::sync::{ struct InputProbe { inputs: WorkerInputs, tier_matches: bool, - shared_hits: bool, + shared_hits: Option, + shared_identity: Arc, declarations: Arc, calls: Arc, } @@ -41,6 +42,7 @@ impl InputProbe { self.inputs.contains(WorkerInputs::PREFERRED_TAINT) ); if let Some(cache) = candidate.cache() { + self.check_cache(cache); assert_eq!( cache.host_overlap_blocks(), if self.tier_matches { 2.0 } else { 0.0 } @@ -55,20 +57,22 @@ impl InputProbe { self.calls.fetch_add(1, Ordering::Relaxed); // Ordinary request context stays available without CACHE. assert_eq!(context.prompt_tokens(), 17); - if self.inputs.contains(WorkerInputs::CACHE) { - assert_eq!( - context.cache().unwrap().has_tier_matches(), - self.tier_matches - ); - assert_eq!( - context.cache().unwrap().shared_hits().is_some(), - self.shared_hits - ); - if let Some(hits) = context.cache().unwrap().shared_hits() { - assert_eq!(hits.hits_beyond(0), 4); + } + + fn check_cache(&self, cache: WorkerCacheInput<'_>) { + assert_eq!(cache.has_tier_matches(), self.tier_matches); + assert_eq!(cache.shared_hits().is_some(), self.shared_hits.is_some()); + if let Some(hits) = cache.shared_hits() { + assert_eq!(hits.hits_beyond(0), self.shared_hits.unwrap()); + let address = std::ptr::from_ref(hits) as usize; + if let Err(previous) = self.shared_identity.compare_exchange( + 0, + address, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + assert_eq!(address, previous); } - } else { - assert!(context.cache().is_none()); } } } @@ -109,9 +113,6 @@ impl WorkerScorer for InputProbe { let indexed = candidates.get(row).unwrap(); self.check_candidate(indexed); assert_eq!(indexed.worker(), candidate.worker()); - if let Some(cache) = candidate.cache() { - assert!(std::ptr::eq(cache, indexed.cache().unwrap())); - } } costs.fill(0.0); Ok(()) @@ -137,6 +138,15 @@ impl WorkerPicker for InputProbe { input.load().is_some(), self.inputs.contains(WorkerInputs::LOAD) ); + if let Some(cache) = input.cache() { + assert_eq!(cache.len(), input.candidates().len()); + assert!(!cache.is_empty()); + assert!(cache.get(cache.len()).is_none()); + for (row, entry) in cache.iter().enumerate() { + self.check_cache(entry); + self.check_cache(cache.get(row).unwrap()); + } + } for candidate in input.candidates() { assert_eq!( candidate.preferred_taint_multiplier().is_some(), @@ -170,15 +180,18 @@ fn each_component_can_only_read_its_own_inputs() { for scorer_mask in 0..8 { for picker_mask in 0..8 { for tier_matches in [false, true] { - for shared_hits in [false, true] { + for shared_hits in [None, Some(0), Some(4)] { let (workers, mut request) = support::fixture(2, 17); if !tier_matches { request.overlap.tier_overlap_blocks = Default::default(); } - if shared_hits { - request.shared_cache_hits = - Some(SharedCacheHits::from_ranges(vec![1..3, 5..7])); - } + request.shared_cache_hits = shared_hits.map(|count| { + SharedCacheHits::from_ranges(if count == 0 { + vec![] + } else { + vec![1..3, 5..7] + }) + }); // Even unmatched preferences produce Some(1.0), so an undeclared // component leaking another component's taint value fails this test. request @@ -187,10 +200,12 @@ fn each_component_can_only_read_its_own_inputs() { .insert("preferred".into(), 1.0); let declarations = Arc::new(AtomicUsize::new(0)); let calls = Arc::new(AtomicUsize::new(0)); + let shared_identity = Arc::new(AtomicUsize::new(0)); let probe = |mask| InputProbe { inputs: inputs(mask), tier_matches, shared_hits, + shared_identity: shared_identity.clone(), declarations: declarations.clone(), calls: calls.clone(), }; @@ -213,6 +228,12 @@ fn each_component_can_only_read_its_own_inputs() { .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); } + if let Some(hits) = request.shared_cache_hits.as_ref() { + assert_eq!( + shared_identity.load(Ordering::Relaxed), + std::ptr::from_ref(hits) as usize + ); + } assert_eq!(calls.load(Ordering::Relaxed), 22); assert_eq!(declarations.load(Ordering::Relaxed), 5); } diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 7a2dc0b0dd76..61c4d67bacc7 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -157,16 +157,23 @@ fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { ) -> Result { assert_eq!(context.prompt_tokens(), 17); assert_eq!(context.request_blocks(), 2); - assert!(context.cache().unwrap().has_tier_matches()); assert!(input.load().unwrap().iter().all(|load| load.is_available())); - for cache in input.cache().unwrap() { + for cache in input.cache().unwrap().iter() { + assert!(cache.has_tier_matches()); let (blocks, tokens) = cache.accounting_cache_estimate(); assert_eq!((blocks * 16.0) as usize, tokens); } Ok(0) } } - let (workers, request) = fixture(2, 17); + let (workers, mut request) = fixture(2, 17); + // Tier availability belongs to the lookup, even when its only entry is ineligible. + request.overlap.tier_overlap_blocks = Default::default(); + request + .overlap + .tier_overlap_blocks + .disk + .insert(WorkerWithDpRank::from_worker_id(999), 1); WorkerSelectionPolicy::new(KvRouterConfig::default(), "test", vec![], Box::new(Inspect)) .select_worker(support::selection_input(&workers, &request, 16)) .unwrap(); From 7101ccae979cf3f05fbe8b01a7fcf32ce4717845 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 02:46:24 +0000 Subject: [PATCH 25/31] fix(python): keep router input validation independent of Python Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/entrypoint.rs | 23 ++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 6dab6ba64939..7d1347bae907 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -1035,7 +1035,8 @@ pub fn run_input<'p>( .router_config() .router_mode .is_kv_routing(), - )?; + ) + .map_err(PyValueError::new_err)?; crate::future_into_py(py, async move { if matches!(&input_enum, Input::Http) { HttpFrontend::default() @@ -1062,17 +1063,13 @@ fn validate_router_plugin_input( plugins: &dynamo_kv_router::plugins::RouterPlugins, input: &Input, is_kv_routing: bool, -) -> PyResult<()> { +) -> Result<(), &'static str> { if plugins.has_custom_plugins() { if !is_kv_routing { - return Err(PyValueError::new_err( - "linked router plugins require --router-mode kv", - )); + return Err("linked router plugins require --router-mode kv"); } if !matches!(input, Input::Http) { - return Err(PyValueError::new_err( - "linked router plugins require HTTP frontend input", - )); + return Err("linked router plugins require HTTP frontend input"); } } Ok(()) @@ -1106,7 +1103,13 @@ mod plugin_input_tests { let plugins = RouterPlugins::default() .with_worker_selection(dynamo_custom_policy_builtin::default_factory()); assert!(validate_router_plugin_input(&plugins, &Input::Http, true).is_ok()); - assert!(validate_router_plugin_input(&plugins, &Input::Http, false).is_err()); - assert!(validate_router_plugin_input(&plugins, &Input::Grpc, true).is_err()); + assert_eq!( + validate_router_plugin_input(&plugins, &Input::Http, false), + Err("linked router plugins require --router-mode kv") + ); + assert_eq!( + validate_router_plugin_input(&plugins, &Input::Grpc, true), + Err("linked router plugins require HTTP frontend input") + ); } } From 98d3c70fcd41a923e2ec388da8830725d95cf922 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 02:46:38 +0000 Subject: [PATCH 26/31] refactor(router): apply builtin policy review cleanups Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/local.rs | 4 +-- .../src/{default/mod.rs => default.rs} | 8 ++--- .../builtin/src/default/scorer.rs | 32 +++++++++---------- .../builtin/src/default/selector.rs | 2 +- 4 files changed, 23 insertions(+), 23 deletions(-) rename lib/router-plugins/builtin/src/{default/mod.rs => default.rs} (92%) diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index fdf8ef64e3bb..332960b7a3bd 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -1300,8 +1300,8 @@ mod tests { 64, Some(vec![1, 2, 3, 4]), TierOverlapBlocks::default(), - rustc_hash::FxHashMap::from_iter([(worker, 0.75)]), - rustc_hash::FxHashMap::from_iter([(worker, 48)]), + FxHashMap::from_iter([(worker, 0.75)]), + FxHashMap::from_iter([(worker, 48)]), None, true, None, diff --git a/lib/router-plugins/builtin/src/default/mod.rs b/lib/router-plugins/builtin/src/default.rs similarity index 92% rename from lib/router-plugins/builtin/src/default/mod.rs rename to lib/router-plugins/builtin/src/default.rs index 0657694637e5..e6ab04b0b4d2 100644 --- a/lib/router-plugins/builtin/src/default/mod.rs +++ b/lib/router-plugins/builtin/src/default.rs @@ -31,9 +31,9 @@ fn policy_with_rng( parameters: PolicyParameters, worker_label: &'static str, rng: Option>>, - plain_decode: bool, + is_plain_decode: bool, ) -> WorkerSelectionPolicy { - let scorer = scorer::build(¶meters, worker_label, plain_decode); + let scorer = scorer::build(¶meters, worker_label, is_plain_decode); let picker = picker::DefaultPicker::new(parameters.router_temperature, rng); WorkerSelectionPolicy::new(config, worker_label, vec![scorer], Box::new(picker)) .with_exclusive_affinity(true) @@ -51,13 +51,13 @@ fn policy_for_role( role: dynamo_kv_router::WorkerType, parameters: PolicyParameters, ) -> WorkerSelectionPolicy { - let plain_decode = + let is_plain_decode = role == dynamo_kv_router::WorkerType::Decode && !config.conditional_disagg_enabled; policy_with_rng( config, parameters, role.default_selector_label(), None, - plain_decode, + is_plain_decode, ) } diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index ab0a5c801d32..610f3f11c612 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -14,7 +14,7 @@ use dynamo_kv_router::plugins::worker_selection::{ pub(super) fn build( config: &PolicyParameters, worker_label: &'static str, - plain_decode: bool, + is_plain_decode: bool, ) -> Box { match ( config.decode_active_request_weight != 0.0, @@ -23,22 +23,22 @@ pub(super) fn build( (false, false) => Box::new(DefaultScorer::::new( config, worker_label, - plain_decode, + is_plain_decode, )), (false, true) => Box::new(DefaultScorer::::new( config, worker_label, - plain_decode, + is_plain_decode, )), (true, false) => Box::new(DefaultScorer::::new( config, worker_label, - plain_decode, + is_plain_decode, )), (true, true) => Box::new(DefaultScorer::::new( config, worker_label, - plain_decode, + is_plain_decode, )), } } @@ -52,7 +52,7 @@ struct DefaultScorer { decode_active_request_weight: f64, prefill_load_scale: f64, is_decode: bool, - plain_decode: bool, + is_plain_decode: bool, prepared: PreparedRequest, } @@ -64,8 +64,8 @@ struct PreparedRequest { block_size: IntegerDivisor, request_blocks: IntegerDivisor, overlap_credit: f64, - use_decay: bool, - subtract_from_decode: bool, + needs_decay: bool, + needs_decode_subtraction: bool, } /// Division by a power-of-two integer has an exact reciprocal. Both inputs to this helper @@ -106,7 +106,7 @@ impl pub(super) fn new( config: &PolicyParameters, worker_label: &'static str, - plain_decode: bool, + is_plain_decode: bool, ) -> Self { Self { overlap_score_credit: config.overlap_score_credit, @@ -117,7 +117,7 @@ impl decode_active_request_weight: config.decode_active_request_weight, prefill_load_scale: config.prefill_load_scale, is_decode: worker_label == "decode", - plain_decode, + is_plain_decode, prepared: PreparedRequest::default(), } } @@ -132,7 +132,7 @@ impl candidates: WorkerCandidates<'_>, ) -> Result<(), WorkerSelectionPolicyError> { // Plain disaggregated decode is load-only. Conditional decode retains cache credit. - let overlap_credit = if self.plain_decode && !context.tracks_prefill_tokens() { + let overlap_credit = if self.is_plain_decode && !context.tracks_prefill_tokens() { 0.0 } else { self.overlap_score_credit @@ -142,12 +142,12 @@ impl block_size: IntegerDivisor::new(u64::from(context.block_size())), request_blocks: IntegerDivisor::new(context.request_blocks()), overlap_credit, - use_decay: context.tracks_prefill_tokens() && self.overlap_score_credit_decay > 0.0, - subtract_from_decode: self.is_decode + needs_decay: context.tracks_prefill_tokens() && self.overlap_score_credit_decay > 0.0, + needs_decode_subtraction: self.is_decode && !context.tracks_prefill_tokens() && overlap_credit > 0.0, }; - if self.prepared.use_decay { + if self.prepared.needs_decay { let mut minimum = usize::MAX; for candidate in candidates.iter() { let load = candidate @@ -186,7 +186,7 @@ impl // Preserve signed zero when the configured weight is -0.0. self.shared_cache_multiplier }; - let decay = if self.prepared.use_decay { + let decay = if self.prepared.needs_decay { let excess = self.prepared.block_size.divide( load.active_prefill_tokens() .saturating_sub(self.prepared.min_prefill) as f64, @@ -205,7 +205,7 @@ impl } else { self.decode_active_request_weight }; - let logit = if self.prepared.subtract_from_decode { + let logit = if self.prepared.needs_decode_subtraction { (load.decode_cost_blocks() - credit).max(0.0) + request_cost } else { let raw_tokens = if !context.tracks_prefill_tokens() { diff --git a/lib/router-plugins/builtin/src/default/selector.rs b/lib/router-plugins/builtin/src/default/selector.rs index 96782ae00dae..e07dbb1d4b08 100644 --- a/lib/router-plugins/builtin/src/default/selector.rs +++ b/lib/router-plugins/builtin/src/default/selector.rs @@ -23,7 +23,7 @@ impl DefaultWorkerSelector { pub fn new(config: Option, worker_type: &'static str) -> Self { Self::with_rng(config.unwrap_or_default(), worker_type, None) } - /// Construct a reproducible selector. Clones share its random stream, as before. + /// Construct a reproducible selector. Clones share its random stream. pub fn new_seeded( config: Option, worker_type: &'static str, From feb8b346c75630e49558cc02d14f9c21dad25ca3 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 02:53:23 +0000 Subject: [PATCH 27/31] perf(router): avoid copying costs for unseeded sampling Signed-off-by: Ishan Dhanani --- .../builtin/src/default/picker.rs | 9 +++++ .../builtin/tests/default_policy.rs | 33 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/lib/router-plugins/builtin/src/default/picker.rs b/lib/router-plugins/builtin/src/default/picker.rs index 3b38f77e0413..e703e2f70c4c 100644 --- a/lib/router-plugins/builtin/src/default/picker.rs +++ b/lib/router-plugins/builtin/src/default/picker.rs @@ -110,6 +110,15 @@ impl WorkerPicker for DefaultPicker { } return Ok(best_row); } + if self.rng.is_none() { + return Ok(softmax_sample_index( + candidates, + |candidate| candidate.cost(), + self.temperature, + fastrand::f64(), + &mut self.probabilities, + )); + } self.entries.clear(); self.entries.extend( candidates diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 61c4d67bacc7..77eca6b96c7d 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -63,6 +63,39 @@ fn seeded_selection_matches_reference_across_cache_and_load_shapes() { } } +#[test] +fn unseeded_sampling_matches_reference_with_the_same_random_draw() { + for count in [1, 8, 64] { + for temperature in [0.1, 0.7, 1.0, 2.0] { + for equal_costs in [false, true] { + let (workers, mut request) = fixture(count, 2048); + if equal_costs { + request.worker_loads.clear(); + request.overlap = Default::default(); + } + let config = KvRouterConfig { + router_temperature: temperature, + ..Default::default() + }; + let reference = + dynamo_kv_router::DefaultWorkerSelector::new(Some(config.clone()), "prefill"); + let plugin = default_policy(config, "prefill"); + for seed in 0..64 { + let input = selection_input(&workers, &request, 16); + fastrand::seed(seed); + let expected = reference.select_worker(input).unwrap(); + let next_random = fastrand::u64(..); + fastrand::seed(seed); + let actual = plugin.select_worker(input).unwrap(); + assert_eq!(actual.worker, expected.worker); + assert_eq!(actual.cached_tokens, expected.cached_tokens); + assert_eq!(fastrand::u64(..), next_random); + } + } + } + } +} + #[test] fn unseeded_minimum_picker_only_selects_workers_tied_for_lowest_cost() { let (workers, mut request) = fixture(8, 17); From 3bf711fe4d89a83ce68c0fd8bcb081f525e9fa76 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 05:05:04 +0000 Subject: [PATCH 28/31] refactor(router): simplify seeded builtin picker Signed-off-by: Ishan Dhanani --- .../builtin/src/default/picker.rs | 28 ++++++------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/lib/router-plugins/builtin/src/default/picker.rs b/lib/router-plugins/builtin/src/default/picker.rs index e703e2f70c4c..8f42239e4100 100644 --- a/lib/router-plugins/builtin/src/default/picker.rs +++ b/lib/router-plugins/builtin/src/default/picker.rs @@ -110,7 +110,7 @@ impl WorkerPicker for DefaultPicker { } return Ok(best_row); } - if self.rng.is_none() { + let Some(rng) = &self.rng else { return Ok(softmax_sample_index( candidates, |candidate| candidate.cost(), @@ -118,7 +118,7 @@ impl WorkerPicker for DefaultPicker { fastrand::f64(), &mut self.probabilities, )); - } + }; self.entries.clear(); self.entries.extend( candidates @@ -127,13 +127,11 @@ impl WorkerPicker for DefaultPicker { .map(|(row, candidate)| (row, candidate.cost())), ); // Canonical order is required only for deterministic replay, never for production ties. - if self.rng.is_some() { - self.entries.sort_unstable_by_key(|(row, _)| { - let worker = candidates[*row].worker(); - (worker.worker_id, worker.dp_rank) - }); - } - let mut rng = self.rng.as_ref().map(|rng| rng.lock()); + self.entries.sort_unstable_by_key(|(row, _)| { + let worker = candidates[*row].worker(); + (worker.worker_id, worker.dp_rank) + }); + let mut rng = rng.lock(); let selected = if self.temperature == 0.0 { let mut best = 0; let mut best_cost = f64::INFINITY; @@ -145,26 +143,18 @@ impl WorkerPicker for DefaultPicker { ties = 1; } else if *cost == best_cost { ties += 1; - let sample = match rng.as_mut() { - Some(rng) => rng.usize(0..ties), - None => fastrand::usize(0..ties), - }; - if sample == 0 { + if rng.usize(0..ties) == 0 { best = index; } } } best } else { - let sample = match rng.as_mut() { - Some(rng) => rng.f64(), - None => fastrand::f64(), - }; softmax_sample_index( &self.entries, |(_, cost)| *cost, self.temperature, - sample, + rng.f64(), &mut self.probabilities, ) }; From d83839ce7b9a930362ed7499848b89baab7fb882 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 05:05:04 +0000 Subject: [PATCH 29/31] fix(router): preserve builtin load-only cache requirements Signed-off-by: Ishan Dhanani --- .../configuration/groups/kv_router_args.py | 2 +- .../router/custom-worker-selection.mdx | 4 +- lib/llm/src/kv_router.rs | 57 ++++++++++++++- .../builtin/src/default/scorer.rs | 71 +++++++++++-------- .../builtin/src/default/selector.rs | 2 +- .../builtin/tests/default_policy.rs | 47 +++++++----- 6 files changed, 130 insertions(+), 53 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 372f63d8f6a0..0fc85150cfec 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -300,7 +300,7 @@ def add_arguments(self, parser) -> None: "This preset sets overlap_score_credit=0, disables KV events and " "KV-reuse assumptions, enables active-block " "and prefill-token load tracking, and disables remote/shared cache indexers. " - "The policy's CACHE input declaration still controls local approximate indexing." + "The builtin policy does not request cache inputs in this mode." ), ) add_argument( diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index f7e57f1ada39..b4669b7a6fd2 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -416,7 +416,7 @@ The disaggregated prefill hop and query-only probes do not run classifiers. Stat 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`. Scoring weights do not control cache setup. +Routing cache resources start only if a component requests `CACHE`. `serve_indexer` and `enable_session_prefix_index` also require `CACHE`. ### Request Context @@ -631,4 +631,4 @@ worker_selection: router_temperature: 0.0 ``` -The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. Per-request score and temperature overrides do not apply to this policy. +The provider also accepts `overlap_score_credit_decay`, `decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, and `shared_cache_multiplier`. Omitted parameters inherit router configuration. Values must be finite and non-negative. Per-request score and temperature overrides do not apply to this policy. The builtin skips cache inputs at zero overlap credit and for plain disaggregated decode. diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index eb0b61dbbcca..b6c5cb2e8d94 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -2366,6 +2366,61 @@ mod tests { } } + #[tokio::test] + async fn builtin_load_only_modes_skip_indexing_and_subscriptions() { + for role in [WorkerType::Aggregated, WorkerType::Decode] { + let router = make_router( + "builtin-load-only", + HashMap::from([(7, ModelRuntimeConfig::default())]), + 2, + SelectionPolicySource::Registry, + None, + Some(role), + role.default_selector_label(), + KvRouterConfig { + overlap_score_credit: if role == WorkerType::Decode { 1.0 } else { 0.0 }, + use_kv_events: role == WorkerType::Decode, + router_assume_kv_reuse: false, + skip_initial_worker_wait: true, + ..Default::default() + }, + ) + .await + .unwrap(); + assert!( + !router + .required_worker_inputs() + .contains(WorkerInputs::CACHE) + ); + assert!(matches!(router.indexer, Indexer::None)); + assert!(!router.ingress.has_subscription()); + let tokens = vec![11, 12, 13, 14]; + for _ in 0..2 { + router + .record_routing_decision( + TokensWithHashes::new(tokens.clone(), 2), + WorkerWithDpRank::from_worker_id(7), + ) + .await + .unwrap(); + let FindBestMatchOutcome::Routed { cached_tokens, .. } = + find_best_match(&router, &tokens, false).await.unwrap() + else { + panic!("load-only request must route"); + }; + assert_eq!(cached_tokens, 0); + assert_eq!( + router + .prefill_load_hint_for(tokens.len(), cached_tokens, true) + .unwrap() + .initial_effective_prefill_tokens, + tokens.len() + ); + } + router.cancellation_token.cancel(); + } + } + #[tokio::test] async fn indexer_features_require_declared_cache_inputs() { for feature in ["serve_indexer", "enable_session_prefix_index"] { @@ -2586,7 +2641,7 @@ mod tests { workers, Some(membership), 2, - SelectionPolicySource::Registry, + fixed_policy(None, WorkerWithDpRank::from_worker_id(7)), Some(KvRouterConfig { enable_session_prefix_index: true, router_event_threads: threads, diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 610f3f11c612..57d3afdeab2a 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -165,41 +165,43 @@ impl context: &WorkerSelectionContext<'_>, candidate: WorkerCandidate<'_>, ) -> Result { - let cache = candidate - .cache() - .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); - let device = if cache.has_tier_matches() { - cache.device_overlap_blocks() - } else { - estimated_overlap - }; - let shared_credit = if SHARED_CREDIT { - let shared = cache - .shared_hits() - .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)); - self.shared_cache_multiplier * shared as f64 - } else { - // Preserve signed zero when the configured weight is -0.0. - self.shared_cache_multiplier - }; - let decay = if self.prepared.needs_decay { - let excess = self.prepared.block_size.divide( - load.active_prefill_tokens() - .saturating_sub(self.prepared.min_prefill) as f64, - ); - 1.0 / (1.0 - + self.overlap_score_credit_decay * self.prepared.request_blocks.divide(excess)) + let (cached_tokens, credit) = if let Some(cache) = candidate.cache() { + let (estimated_overlap, cached_tokens) = cache.accounting_cache_estimate(); + let device = if cache.has_tier_matches() { + cache.device_overlap_blocks() + } else { + estimated_overlap + }; + let shared_credit = if SHARED_CREDIT { + let shared = cache + .shared_hits() + .map_or(0, |hits| hits.hits_beyond(device.round().max(0.0) as u32)); + self.shared_cache_multiplier * shared as f64 + } else { + // Preserve signed zero when the configured weight is -0.0. + self.shared_cache_multiplier + }; + let decay = if self.prepared.needs_decay { + let excess = self.prepared.block_size.divide( + load.active_prefill_tokens() + .saturating_sub(self.prepared.min_prefill) as f64, + ); + 1.0 / (1.0 + + self.overlap_score_credit_decay * self.prepared.request_blocks.divide(excess)) + } else { + 1.0 + }; + let credit = self.prepared.overlap_credit * decay * device + + self.host_cache_hit_weight * cache.host_overlap_blocks() + + self.disk_cache_hit_weight * cache.disk_overlap_blocks() + + shared_credit; + (cached_tokens, credit) } else { - 1.0 + (0, 0.0) }; - let credit = self.prepared.overlap_credit * decay * device - + self.host_cache_hit_weight * cache.host_overlap_blocks() - + self.disk_cache_hit_weight * cache.disk_overlap_blocks() - + shared_credit; let request_cost = if REQUEST_COST { self.decode_active_request_weight * load.active_requests() as f64 } else { @@ -228,7 +230,14 @@ impl WorkerScorer for DefaultScorer { fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT + // The builtin's load-only modes must book the full prompt without cache credit. + // Hosts still start indexing solely from the policy's declared inputs. + let inputs = WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT; + if !self.is_plain_decode && self.overlap_score_credit != 0.0 { + inputs | WorkerInputs::CACHE + } else { + inputs + } } fn score( diff --git a/lib/router-plugins/builtin/src/default/selector.rs b/lib/router-plugins/builtin/src/default/selector.rs index e07dbb1d4b08..1fb9caeeb9e2 100644 --- a/lib/router-plugins/builtin/src/default/selector.rs +++ b/lib/router-plugins/builtin/src/default/selector.rs @@ -73,7 +73,7 @@ impl std::fmt::Debug for DefaultWorkerSelector { } impl WorkerSelector for DefaultWorkerSelector { fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::CACHE | WorkerInputs::LOAD + >::required_worker_inputs(&self.policy.lock()) } fn uses_exclusive_affinity_target(&self) -> bool { true diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 77eca6b96c7d..0694990a0bca 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -215,23 +215,36 @@ fn exact_prompt_and_accounting_inputs_are_available_to_external_pickers() { #[test] fn configured_default_resolves_for_every_role_and_uses_exclusive_affinity() { let registry = default_registry(); - for role in [ - WorkerType::Aggregated, - WorkerType::Prefill, - WorkerType::Decode, - WorkerType::Encode, - ] { - let config = KvRouterConfig::default(); - let factory = registry - .resolve_for_worker_type(&config, role) - .unwrap() - .unwrap(); - let policy = factory(&config, role, RoutingPartitionRef::new("model", "default")); - assert!( - >::uses_exclusive_affinity_target( - &policy - ) - ); + for conditional_disagg_enabled in [false, true] { + for role in [ + WorkerType::Aggregated, + WorkerType::Prefill, + WorkerType::Decode, + WorkerType::Encode, + ] { + let config = KvRouterConfig { + conditional_disagg_enabled, + ..Default::default() + }; + let factory = registry + .resolve_for_worker_type(&config, role) + .unwrap() + .unwrap(); + let policy = factory(&config, role, RoutingPartitionRef::new("model", "default")); + let inputs = + >::required_worker_inputs( + &policy, + ); + assert_eq!( + inputs.contains(WorkerInputs::CACHE), + role != WorkerType::Decode || conditional_disagg_enabled + ); + assert!( + >::uses_exclusive_affinity_target( + &policy + ) + ); + } } } From e1360a17b2f9251e309c91ac01c140b4f11633ce Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 17:05:20 +0000 Subject: [PATCH 30/31] docs(router): defer request classifier documentation Signed-off-by: Ishan Dhanani --- .../router/custom-worker-selection.mdx | 64 +------------------ .../use-cases/agents/priority-scheduling.md | 2 +- 2 files changed, 4 insertions(+), 62 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index b4669b7a6fd2..a7030c8511a9 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -50,12 +50,11 @@ Use `dynamo_kv_router::plugins` for plugin registration and these modules for pl | Module | Public API | |---|---| | `plugins::worker_selection` | `WorkerFilter`, `WorkerScorer`, `WorkerPicker`, `WorkerSelectionPolicy`, request context, worker signal accessors, configuration, and factories | -| `plugins::request_classifier` | `RequestClassifier`, `ClassifyRequest`, `ClassifyEvent`, configuration, and factories | | `plugins` | `RouterPluginRegistry` and the resolved `RouterPlugins` bundle | Dynamo owns scheduler queues, eligibility checks, reservations, lifecycle synchronization, and the default selector implementation. Low-level host interfaces such as `WorkerSelector` are used to connect Dynamo components; external plugins use `WorkerSelectionPolicy`. -Existing plugin imports and the `WorkerSelectionPolicyRegistry` name remain available through compatibility re-exports until version 1.7. They refer to the same types and traits, so existing policies can register into `RouterPluginRegistry` without adapters. Migrate imports to `plugins::worker_selection` or `plugins::request_classifier`, rename the registry to `RouterPluginRegistry`, and use `register_worker_selection` for worker policies. The compatibility exports, registry alias, worker-only `register` method, and frontend `worker_selection_policy_factory` setter are scheduled for removal in 1.7. +Existing plugin imports and the `WorkerSelectionPolicyRegistry` name remain available through compatibility re-exports until version 1.7. They refer to the same types and traits, so existing policies can register into `RouterPluginRegistry` without adapters. Migrate imports to `plugins::worker_selection`, rename the registry to `RouterPluginRegistry`, and use `register_worker_selection` for worker policies. The compatibility exports, registry alias, worker-only `register` method, and frontend `worker_selection_policy_factory` setter are scheduled for removal in 1.7. ## Build the Policy @@ -332,49 +331,7 @@ cargo check --manifest-path "$POLICY_DIR/catalog/Cargo.toml" Add one focused test for each policy decision and one registration test for every type name. -## Register a Request Classifier - -**Experimental.** Request classifiers use the same statically linked catalog as worker filters, scorers, and pickers. Keep the catalog's single `register` function and call `register_request_classifier` for each classifier type. `RouterPluginRegistry` also accepts the existing worker-selection providers; `WorkerSelectionPolicyRegistry` remains an alias, so existing catalogs compile unchanged. - -Implement `RequestClassifier` in your policy crate. Its provider validates `RequestClassifierParameters` at startup and returns a `RequestClassifierFactory`: a closure that receives `RequestClassifierContext` and creates a fresh `Box` for each routed model. The shared router builder supplies the cached capacity context for both HTTP discovery and Python `KvRouter.generate`. Register that provider alongside your worker-selection providers: - -```rust -use std::sync::Arc; -use dynamo_kv_router::plugins::RouterPluginRegistry; - -pub fn register(registry: &mut RouterPluginRegistry) -> Result<(), Box> { - acme_worker_policy::register(registry)?; - registry.register_request_classifier("acme-admission", Arc::new(acme_admission::provider))?; - Ok(()) -} -``` - -The provider signature is: - -```rust -use dynamo_kv_router::plugins::request_classifier::{ - RequestClassifierFactory, RequestClassifierParameters, RequestClassifierProviderError, -}; - -pub fn provider( - parameters: &RequestClassifierParameters, -) -> Result; -``` - -Select the classifier in the same `--router-policy-config` YAML: - -```yaml -request_classifier: - type: acme-admission - parameters: - max_in_flight: 32 -``` - -The parameter mapping belongs to the plugin; `max_in_flight` above is an example plugin setting. Omit `request_classifier` to retain pass-through admission. Unknown types and invalid parameters fail startup. Build and link the catalog with the existing `custom-policy` workflow below; no second catalog export or build flag is required. - -The discovery-backed HTTP frontend and Python `KvRouter.generate` path install classifiers on their KV routers. Classification applies to decode/aggregated admission before queue ordering and worker selection. Selection is fixed at frontend startup and is process-wide; per-model queue profiles do not select different classifier types. A classifier can change the policy class, due time, and scheduling cost through the existing `ClassifyRequest` API. Dynamo retains queueing, worker selection, reservations, and lifecycle delivery. - -### Shared Plugin Construction +## Shared Plugin Construction The catalog registers plugin providers once. `RouterPluginRegistry::resolve_plugins` validates the selected configuration and returns a `RouterPlugins` bundle containing the configured factories. Both the HTTP frontend and Python router pass this bundle to the shared router construction path. Each router receives fresh plugin instances; factories run during construction, not for each request. @@ -403,14 +360,9 @@ flowchart LR HTTP --> Build[Shared router construction] Python --> Build Build --> Selector[Per-router worker selector] - Build --> Classifier[Per-router request classifier] ``` -The host's `RouterPluginBuilder` passes worker selection through the shared selection core and installs the remaining plugins before the router is exposed to requests. The worker policy is constructed once per partition, including the probe for its required inputs. Adding another plugin trait requires a typed registration slot, a factory in the bundle, and its construction hook. Entry points continue to carry one bundle. Execution and lifecycle callbacks stay on each trait; the bundle does not add per-request dispatch. - -The compatibility paths described above remain available until 1.7. An empty bundle retains the default selector and pass-through admission. - -The disaggregated prefill hop and query-only probes do not run classifiers. Stateful low-level `best_worker` calls also lack classifier lifecycle enrollment. Offline replay and standalone selection reject `request_classifier` configuration because they cannot run its lifecycle. +Each worker policy is constructed once per routing partition. An empty bundle retains the default selector. ## Available Signals @@ -590,16 +542,6 @@ Standalone EPP supplies `WorkerType::Aggregated` because it selects from one ful -## Request Classifier Signals and Placement - -Rust request classifiers can defer admission while retaining independently pollable futures. `KvRouter::request_classifier_context()` supplies the router's cached worker view when constructing a classifier: - -- `block_size()` is the number of tokens per KV block. Each entry in `workers()` identifies a registered worker/rank and its advertised `total_kv_blocks()`, or `None` if capacity was not advertised. Their product is total capacity, not currently free space. The view follows existing discovery updates and does not probe worker health. -- `ClassifyRequest::progress()` returns a read-only context high-water mark, initialized from `input_tokens()` and raised by host observations of prompt plus generated tokens. Clones share a monotonic counter across migration retries. A new request lifecycle gets a fresh counter, even if it reuses a request ID. This measures logical tokens, not physical KV occupancy. -- `set_worker_selection_target(worker)` replaces the request's soft affinity preference; `clear_worker_selection_target()` removes that preference. Hard pins and caller eligibility constraints remain authoritative. A custom picker can fall back when the preferred worker is ineligible. The classifier's `Sent` event reports the worker actually used. - -Holding a classification future pauses admission before dispatch; it does not interrupt an already running generation. Worker filters, scorers, and pickers continue to own placement decisions within Dynamo's eligibility rules. - ## Policy Contract - Return `true` from a filter to keep a worker and `false` to reject it. diff --git a/docs/fern/pages/use-cases/agents/priority-scheduling.md b/docs/fern/pages/use-cases/agents/priority-scheduling.md index e81785aaec9a..679967995425 100644 --- a/docs/fern/pages/use-cases/agents/priority-scheduling.md +++ b/docs/fern/pages/use-cases/agents/priority-scheduling.md @@ -52,7 +52,7 @@ The router queue only matters when requests are held before dispatch. If a reque (strict_priority, due_at, configured_policy_key) ``` -The strict tier is compared first. Within that tier, every request with a classifier-provided `due_at` is scheduled ahead of every request without one (an unset due time sorts last) before FCFS, LCFS, or Weighted Shortest Processing Time (WSPT) applies its configured policy key. Sustained deadline-bearing traffic can therefore starve non-deadline requests in the same class. A `due_at` is set only by a router request classifier; select a [linked request classifier](../../developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx#register-a-request-classifier) through router policy YAML. +The strict tier is compared first. Within that tier, every request with a classifier-provided `due_at` is scheduled ahead of every request without one (an unset due time sorts last) before FCFS, LCFS, or Weighted Shortest Processing Time (WSPT) applies its configured policy key. Sustained deadline-bearing traffic can therefore starve non-deadline requests in the same class. A `due_at` is set only by a router request classifier. The default policy is `fcfs`, which uses the priority value as a positive arrival-time bump. Higher values move the request earlier in the queue. Negative priority values are clamped to zero for router queueing, so a request cannot be pushed behind normal first-come, first-served ordering by sending a negative priority. From 1084384acf6baf398e314127dfc17d272e480492 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 22 Sep 2026 18:16:21 +0000 Subject: [PATCH 31/31] fix(router): version the builtin policy dependency Signed-off-by: Ishan Dhanani --- .github/workflows/pre-merge.yml | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pre-merge.yml b/.github/workflows/pre-merge.yml index 9a442fabb52c..470c972ca81b 100644 --- a/.github/workflows/pre-merge.yml +++ b/.github/workflows/pre-merge.yml @@ -541,7 +541,7 @@ jobs: working-directory: ${{ matrix.dir }} run: | # custom-policy is on by default; also lint it alongside the standalone selection - # service, and the build that links no worker-selection policies at all. + # service, and the build without the optional policy catalog. cargo clippy --locked --no-deps --features custom-policy,select-service --all-targets -- -D warnings cargo clippy --locked --no-deps --no-default-features --all-targets -- -D warnings - name: Lint feature-gated benchmark entrypoints diff --git a/Cargo.toml b/Cargo.toml index 5a53f27cca84..e2a1f91c991d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,7 +55,7 @@ repository = "https://github.com/ai-dynamo/dynamo.git" keywords = ["llm", "genai", "inference", "nvidia", "distributed"] [workspace.dependencies] -dynamo-custom-policy-builtin = { path = "lib/router-plugins/builtin" } +dynamo-custom-policy-builtin = { path = "lib/router-plugins/builtin", version = "1.6.0" } # Published separately by https://github.com/ai-dynamo/aisimulate. aisimulate-core = { version = "=0.12.0" }