From a09612e631006995544725c25c842e78836b9b35 Mon Sep 17 00:00:00 2001 From: Dundy Pasupuleti Date: Wed, 21 Oct 2026 01:40:01 -0500 Subject: [PATCH 1/4] feat(router): add prefix_affinity_weight and report non-max-overlap picks in aggregated mode The default cost function caps overlap credit at the request's own prefill, so a little extra decode load on the worker that holds a resumed session's prefix sends the next turn to a cold worker that recomputes it. Add an opt-in prefix_affinity_weight (default 0) that charges each prompt block a candidate is missing from its device cache, in both the plugin scorer and the reference selector. Also install the non-max-overlap observer on aggregated routers, so router_non_max_overlap_selections_total no longer reads 0 by construction in single-pool deployments. Fixes #16194 Signed-off-by: Dundy Pasupuleti --- .../router/configuration-and-tuning.md | 7 +- .../observability/metrics-catalog.mdx | 4 +- lib/bindings/python/rust/llm/entrypoint.rs | 2 + lib/kv-router/src/scheduling/config.rs | 22 +++++ .../src/scheduling/selector/reference.rs | 13 ++- lib/llm/src/kv_router/embedded.rs | 4 +- lib/llm/src/kv_router/metrics.rs | 4 +- .../builtin/src/default/parameters.rs | 6 ++ .../builtin/src/default/scorer.rs | 21 ++++- .../builtin/tests/default_policy.rs | 89 ++++++++++++++++++- 10 files changed, 159 insertions(+), 13 deletions(-) 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 cf80039dd199..05366ecf1ed2 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 @@ -148,6 +148,7 @@ worker_selection: overlap_score_credit_decay: 0.0 prefill_load_scale: 1.0 decode_active_request_weight: 0.0 + prefix_affinity_weight: 0.0 shared_cache_multiplier: 0.5 router_temperature: 0.0 ``` @@ -555,6 +556,8 @@ Use `decode_active_request_weight` when decode forward-pass time depends more on This setting is experimental and defaults to `0`, which preserves block-only decode scoring. A positive value trades some KV locality for batch-size balance and can help model, runtime, and hardware combinations near a compute-bound roofline knee, including some MLA or MTP configurations. It can regress throughput, TTFT, and ITL when decode remains primarily memory-bound, so benchmark representative traffic and start with a small weight before increasing it. +Use `prefix_affinity_weight` for multi-turn or agentic traffic where a resumed request should usually stay on the worker that already holds its prefix. The overlap credit only offsets the request's own prefill cost, so it can be worth at most about one request of load: one extra in-flight sequence on the cache-holding worker can send the next turn to a cold worker that recomputes the whole prefix. This parameter adds `prefix_affinity_weight * missing_prefix_blocks` to each candidate's score, where `missing_prefix_blocks` is the number of prompt blocks the worker does not hold on device. The term is not capped, so it keeps a session on its cache holder until the load gap exceeds the weighted cost of the missing blocks. It defaults to `0`, which preserves the existing scoring. Values between `1` and `4` worked well in a mocker replay of 16 sessions with 12 turns and a 3K-token prefix; benchmark representative traffic before raising it, since a large weight trades load balance for locality. Load-only decode routing ignores it. + Use `--router-host-cache-hit-weight` and `--router-disk-cache-hit-weight` when the backend exposes lower-tier prefix cache via a KV connector (for example, vLLM's `OffloadingConnector` for CPU offload, or a disk-backed tier). These multipliers control how much each lower-tier hit credits against the prefill load, mirroring the role of `overlap_score_credit` for the device tier. A worker holding a full prefix in CPU offload gets `host_cache_hit_weight * matched_blocks` credit against its prefill cost; raising the weight makes the router more willing to route prefix-matched requests to that worker even if a different worker has a partial device-local match. Use `--no-router-kv-events` when you are not confident that your backend engine emits KV events correctly. In this mode the router falls back to approximate routing. Keep the default TTL policy unless you are explicitly testing the experimental per-rank capacity-bounded LRU with workers that publish a positive `total_kv_blocks` value. @@ -630,8 +633,8 @@ Use a returned `router_id` and revision in `PATCH /router_policy`: The default policy accepts the following finite, non-negative parameters: `overlap_score_credit`, `overlap_score_credit_decay`, `prefill_load_scale`, -`decode_active_request_weight`, `host_cache_hit_weight`, `disk_cache_hit_weight`, -`shared_cache_multiplier`, and `router_temperature`. These change policy scoring only; +`decode_active_request_weight`, `prefix_affinity_weight`, `host_cache_hit_weight`, +`disk_cache_hit_weight`, `shared_cache_multiplier`, and `router_temperature`. These change policy scoring only; host-owned load accounting and queue settings remain unchanged. Omitted parameters retain their current values. A parameter set to `null` restores that diff --git a/docs/fern/pages/reference/observability/metrics-catalog.mdx b/docs/fern/pages/reference/observability/metrics-catalog.mdx index b9b6675da18f..32f7c37cb39b 100644 --- a/docs/fern/pages/reference/observability/metrics-catalog.mdx +++ b/docs/fern/pages/reference/observability/metrics-catalog.mdx @@ -367,11 +367,11 @@ The router exposes metrics for routing decisions and overhead. Not every metric - Admitted prefill scheduler selections routed to a worker with less KV cache overlap than another eligible worker. Pinned requests and equal-overlap ties are excluded. Labeled by `worker_type` (`prefill`). + Admitted prefill or aggregated scheduler selections routed to a worker with less KV cache overlap than another eligible worker. Pinned requests and equal-overlap ties are excluded. Decode routers do not report it. Labeled by `worker_type`. - Difference in effective KV cache overlap between the highest-overlap eligible prefill worker and the selected worker for each non-max-overlap selection. Labeled by `worker_type` (`prefill`). + Difference in effective KV cache overlap between the highest-overlap eligible worker and the selected worker for each non-max-overlap selection on a prefill or aggregated router. Labeled by `worker_type`. ### KV-cache reuse telemetry diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 3107184b5bf6..bc544e3a1cd2 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -187,6 +187,8 @@ impl KvRouterConfig { overlap_score_credit_decay, prefill_load_scale, decode_active_request_weight, + // Set through the default policy's `parameters` in router policy YAML. + prefix_affinity_weight: 0.0, host_cache_hit_weight, disk_cache_hit_weight, router_temperature, diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 9c6feb847289..306898eee1bd 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -183,6 +183,7 @@ fn log_env_config(config: &KvRouterConfig) { overlap_score_credit_decay = config.overlap_score_credit_decay, prefill_load_scale = config.prefill_load_scale, decode_active_request_weight = config.decode_active_request_weight, + prefix_affinity_weight = config.prefix_affinity_weight, router_temperature = config.router_temperature, use_kv_events = config.use_kv_events, router_replica_sync = config.router_replica_sync, @@ -698,6 +699,7 @@ struct KvRouterConfigSerde { overlap_score_credit_decay: f64, prefill_load_scale: f64, decode_active_request_weight: f64, + prefix_affinity_weight: f64, overlap_score_weight: Option, host_cache_hit_weight: f64, disk_cache_hit_weight: f64, @@ -750,6 +752,7 @@ impl Default for KvRouterConfigSerde { overlap_score_credit_decay: config.overlap_score_credit_decay, prefill_load_scale: config.prefill_load_scale, decode_active_request_weight: config.decode_active_request_weight, + prefix_affinity_weight: config.prefix_affinity_weight, overlap_score_weight: None, host_cache_hit_weight: config.host_cache_hit_weight, disk_cache_hit_weight: config.disk_cache_hit_weight, @@ -816,6 +819,14 @@ pub struct KvRouterConfig { #[serde(default, skip_serializing_if = "is_default")] pub decode_active_request_weight: f64, + /// Block-equivalent cost added for each prompt block a candidate worker is + /// missing from its device cache. Unlike `overlap_score_credit`, this term is + /// not capped by the request's own prefill cost, so a worker that holds a + /// resumed session's prefix keeps winning against a cold worker with slightly + /// lower transient decode load. Defaults to 0.0, which disables it. + #[serde(default, skip_serializing_if = "is_default")] + pub prefix_affinity_weight: f64, + #[serde(default = "default_host_cache_hit_weight")] pub host_cache_hit_weight: f64, @@ -1003,6 +1014,7 @@ impl Default for KvRouterConfig { overlap_score_credit_decay: default_overlap_score_credit_decay(), prefill_load_scale: default_prefill_load_scale(), decode_active_request_weight: default_decode_active_request_weight(), + prefix_affinity_weight: 0.0, host_cache_hit_weight: default_host_cache_hit_weight(), disk_cache_hit_weight: default_disk_cache_hit_weight(), router_temperature: 0.0, @@ -1068,6 +1080,7 @@ impl TryFrom for KvRouterConfig { overlap_score_credit_decay: compat.overlap_score_credit_decay, prefill_load_scale, decode_active_request_weight: compat.decode_active_request_weight, + prefix_affinity_weight: compat.prefix_affinity_weight, host_cache_hit_weight: compat.host_cache_hit_weight, disk_cache_hit_weight: compat.disk_cache_hit_weight, router_temperature: compat.router_temperature, @@ -1445,6 +1458,7 @@ impl KvRouterConfig { self.overlap_score_credit_decay, )?; validate_finite_non_negative("prefill_load_scale", self.prefill_load_scale)?; + validate_finite_non_negative("prefix_affinity_weight", self.prefix_affinity_weight)?; validate_finite_non_negative( "decode_active_request_weight", self.decode_active_request_weight, @@ -2490,6 +2504,7 @@ worker_selection: let value = serde_json::to_value(KvRouterConfig::default()).unwrap(); for post_v1_3_field in [ "decode_active_request_weight", + "prefix_affinity_weight", "router_tracking_hash", "router_tracking_key_file", "router_tracking_key_id", @@ -2898,6 +2913,13 @@ models: ..Default::default() }, ), + ( + "prefix_affinity_weight", + KvRouterConfig { + prefix_affinity_weight: value, + ..Default::default() + }, + ), ( "router_temperature", KvRouterConfig { diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index 8e8bbfb3da51..afaab5001039 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -378,7 +378,18 @@ impl> DefaultWorkerScorer { 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; + let mut logit = prefill_cost_blocks + decode_cost_blocks + active_request_cost_blocks; + // Load-only decode routing (the plugin's decode subtraction path) never charges it. + if kv_router_config.prefix_affinity_weight != 0.0 + && (self.worker_type != "decode" || context.track_prefill_tokens) + { + // Charge each prompt block this worker lacks. The overlap credit is capped at the + // request's own prefill, so it alone cannot outweigh transient decode load on the + // worker that already holds a resumed prefix. + let missing_prefix_blocks = + (context.request_blocks.max(1) as f64 - device_overlap_blocks).max(0.0); + logit += kv_router_config.prefix_affinity_weight * missing_prefix_blocks; + } // These rows are emitted from the `SchedulerQueueActor` task, which `scheduling::queue` // spawns without the caller's request span, so the logging layer cannot attach diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 5885b258e6d7..e11cbc1e6ce3 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -324,7 +324,9 @@ impl EmbeddedSelection { queue_metrics.clone(), cancellation_token.child_token(), ); - if worker_type == WorkerType::Prefill + // Decode routers ignore prompt overlap by design, so only routers that place + // prompt work (prefill pools and single-pool aggregated serving) report it. + if matches!(worker_type, WorkerType::Prefill | WorkerType::Aggregated) && !partition .scheduler() .set_non_max_overlap_selection_observer(non_max_overlap_observer( diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 02cbaa42f323..2509babc065a 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -1170,7 +1170,7 @@ impl RouterRequestMetrics { let non_max_overlap_selections_total = metrics .create_intcountervec( &router_metric(frontend_service::NON_MAX_OVERLAP_SELECTIONS_TOTAL), - "Total admitted prefill scheduler selections with less KV cache overlap than another eligible worker", + "Total admitted prefill or aggregated scheduler selections with less KV cache overlap than another eligible worker", &[labels::WORKER_TYPE], extra_labels, ) @@ -1178,7 +1178,7 @@ impl RouterRequestMetrics { let overlap_blocks_lost = metrics .create_histogramvec( &router_metric(frontend_service::OVERLAP_BLOCKS_LOST), - "Difference in effective KV cache overlap between the highest-overlap eligible prefill worker and selected worker", + "Difference in effective KV cache overlap between the highest-overlap eligible worker and selected worker", &[labels::WORKER_TYPE], extra_labels, Some(prometheus::exponential_buckets(0.25, 2.0, 16).unwrap()), diff --git a/lib/router-plugins/builtin/src/default/parameters.rs b/lib/router-plugins/builtin/src/default/parameters.rs index 077a6d592cd3..d040a93fb0ce 100644 --- a/lib/router-plugins/builtin/src/default/parameters.rs +++ b/lib/router-plugins/builtin/src/default/parameters.rs @@ -19,6 +19,7 @@ struct Parameters { overlap_score_credit_decay: Option, prefill_load_scale: Option, decode_active_request_weight: Option, + prefix_affinity_weight: Option, host_cache_hit_weight: Option, disk_cache_hit_weight: Option, shared_cache_multiplier: Option, @@ -32,6 +33,7 @@ pub(super) struct PolicyParameters { pub(super) overlap_score_credit_decay: f64, pub(super) prefill_load_scale: f64, pub(super) decode_active_request_weight: f64, + pub(super) prefix_affinity_weight: f64, pub(super) host_cache_hit_weight: f64, pub(super) disk_cache_hit_weight: f64, pub(super) shared_cache_multiplier: f64, @@ -57,6 +59,7 @@ impl Parameters { "decode_active_request_weight", self.decode_active_request_weight, ), + ("prefix_affinity_weight", self.prefix_affinity_weight), ("host_cache_hit_weight", self.host_cache_hit_weight), ("disk_cache_hit_weight", self.disk_cache_hit_weight), ("shared_cache_multiplier", self.shared_cache_multiplier), @@ -93,6 +96,9 @@ impl Parameters { decode_active_request_weight: self .decode_active_request_weight .unwrap_or(config.decode_active_request_weight), + prefix_affinity_weight: self + .prefix_affinity_weight + .unwrap_or(config.prefix_affinity_weight), host_cache_hit_weight: self .host_cache_hit_weight .unwrap_or(inherited_cache_weight(config.host_cache_hit_weight)), diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index d939021b84ec..2282c08dab59 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -34,6 +34,7 @@ struct DefaultScorer { shared_cache_multiplier: f64, decode_active_request_weight: f64, prefill_load_scale: f64, + prefix_affinity_weight: f64, is_decode: bool, prepared: PreparedRequest, } @@ -94,6 +95,7 @@ impl shared_cache_multiplier: config.shared_cache_multiplier, decode_active_request_weight: config.decode_active_request_weight, prefill_load_scale: config.prefill_load_scale, + prefix_affinity_weight: config.prefix_affinity_weight, is_decode: worker_label == "decode", prepared: PreparedRequest::default(), } @@ -139,7 +141,7 @@ impl let load = candidate .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; - let (cached_tokens, credit) = if let Some(cache) = candidate.cache() { + let (cached_tokens, credit, device_blocks) = 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() @@ -169,9 +171,9 @@ impl + self.host_cache_hit_weight * cache.host_overlap_blocks() + self.disk_cache_hit_weight * cache.disk_overlap_blocks() + shared_credit; - (cached_tokens, credit) + (cached_tokens, credit, device) } else { - (0, 0.0) + (0, 0.0, 0.0) }; let request_cost = if REQUEST_COST { self.decode_active_request_weight * load.active_requests() as f64 @@ -190,7 +192,17 @@ impl context.prompt_tokens() }; 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 logit = + self.prefill_load_scale * prefill + load.decode_cost_blocks() + request_cost; + if self.prefix_affinity_weight != 0.0 { + // Charge each prompt block this worker lacks. The overlap credit above is + // capped at the request's own prefill, so it cannot outweigh transient + // decode load on the worker that already holds a resumed prefix. + let missing = (self.prepared.request_blocks.value - device_blocks).max(0.0); + logit + self.prefix_affinity_weight * missing + } else { + logit + } }; let cost = logit * candidate.preferred_taint_multiplier().unwrap_or(1.0); Ok(cost) @@ -203,6 +215,7 @@ impl WorkerScorer fn required_worker_inputs(&self) -> WorkerInputs { let inputs = WorkerInputs::LOAD | WorkerInputs::PREFERRED_TAINT; if self.overlap_score_credit != 0.0 + || self.prefix_affinity_weight != 0.0 || self.host_cache_hit_weight != 0.0 || self.disk_cache_hit_weight != 0.0 || self.shared_cache_multiplier != 0.0 diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 5747b2bc8369..efc19b44fcd6 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -105,6 +105,92 @@ fn runtime_parameters_change_selection_and_reset_to_resolved_startup_values() { assert_eq!(select(&policy), 0); } +/// Issue 16194: a resumed turn whose prefix is cached on one worker. The overlap credit is +/// capped at the request's own prefill, so slightly higher decode load on the cache holder +/// sends the turn to a cold worker unless `prefix_affinity_weight` charges the missing blocks. +#[test] +fn prefix_affinity_weight_keeps_resumed_prefix_on_cache_holder() { + use dynamo_kv_router::plugins::worker_selection::RuntimePolicyUpdate; + use serde_json::json; + + let (workers, mut request) = fixture(2, 47 * 64); + for (worker, load) in &mut request.worker_loads { + let holder = worker.worker_id == 0; + load.active_prefill_tokens = 0; + load.active_requests = 0; + load.active_decode_blocks = if holder { 143 } else { 95 }; + load.additional_active_blocks = 0; + request + .overlap + .tier_overlap_blocks + .device + .insert(*worker, if holder { 46 } else { 0 }); + } + let select = |config: KvRouterConfig| { + let policy = default_policy(config, "prefill"); + policy + .select_worker(selection_input(&workers, &request, 64)) + .unwrap() + .worker + .worker_id + }; + let base = KvRouterConfig { + host_cache_hit_weight: 0.0, + disk_cache_hit_weight: 0.0, + ..Default::default() + }; + assert_eq!(select(base.clone()), 1, "default weights move the session"); + assert_eq!( + select(KvRouterConfig { + prefix_affinity_weight: 0.25, + ..base.clone() + }), + 0 + ); + // The weight is also a runtime policy parameter. + let mut policy = default_policy(base, "prefill"); + let revision = policy.runtime_config().unwrap().revision(); + let patch = json!({"prefix_affinity_weight": 0.25}) + .as_object() + .unwrap() + .clone(); + policy + .update_runtime_config(RuntimePolicyUpdate::patch(revision, patch)) + .unwrap(); + assert_eq!( + policy + .select_worker(selection_input(&workers, &request, 64)) + .unwrap() + .worker + .worker_id, + 0 + ); + let negative = json!({"prefix_affinity_weight": -1.0}) + .as_object() + .unwrap() + .clone(); + assert!( + policy + .update_runtime_config(RuntimePolicyUpdate::patch(revision + 1, negative)) + .is_err() + ); + + // A large enough load gap still wins, so the weight is a bias, not a pin. + for (worker, load) in &mut request.worker_loads { + if worker.worker_id == 0 { + load.active_decode_blocks = 400; + } + } + assert_eq!( + policy + .select_worker(selection_input(&workers, &request, 64)) + .unwrap() + .worker + .worker_id, + 1 + ); +} + #[test] fn runtime_parameters_cannot_enable_worker_signals_missing_at_startup() { use dynamo_kv_router::plugins::worker_selection::{RuntimePolicyError, RuntimePolicyUpdate}; @@ -251,7 +337,7 @@ fn runtime_cache_tier_weights_change_selection_without_device_credit() { 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..32 { + for mode in 0..64 { let (workers, mut request) = fixture(16, prompt); let config = KvRouterConfig { router_temperature: temperature, @@ -260,6 +346,7 @@ fn seeded_selection_matches_reference_across_cache_and_load_shapes() { disk_cache_hit_weight: 0.1, decode_active_request_weight: if mode & 8 == 0 { 0.0 } else { 0.7 }, shared_cache_multiplier: Some(if mode & 16 == 0 { 0.0 } else { 0.6 }), + prefix_affinity_weight: if mode & 32 == 0 { 0.0 } else { 0.9 }, ..Default::default() }; if mode % 8 >= 4 { From 5fe92ff9aeea89f75e189d7de47e26e80102ef19 Mon Sep 17 00:00:00 2001 From: Dundy Pasupuleti Date: Wed, 21 Oct 2026 22:30:57 -0400 Subject: [PATCH 2/4] fix(router): charge prefix affinity in cache-aware decode scoring Cache-aware decode scoring subtracts the overlap credit from decode load and returned before prefix_affinity_weight was applied, so the weight did nothing there. That path covers conditional disaggregation and aggregated routing without prefill tracking, since aggregated routers carry the decode label. Charge missing prefix blocks on it in both the plugin scorer and the reference selector; load-only decode routing with zero overlap credit still ignores the weight. Plain decode now zeroes the weight with the other cache weights, so it never asks for cache inputs. Also initialize the decode series of the non-max-overlap metrics, since aggregated routers emit under that label. Signed-off-by: Dundy Pasupuleti --- .../router/configuration-and-tuning.md | 2 +- .../src/scheduling/selector/reference.rs | 12 +- lib/llm/src/kv_router/metrics.rs | 9 +- lib/router-plugins/builtin/src/default.rs | 1 + .../builtin/src/default/scorer.rs | 12 +- .../builtin/tests/default_policy.rs | 158 +++++++++++++----- 6 files changed, 142 insertions(+), 52 deletions(-) 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 05366ecf1ed2..c12ae429ad4e 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 @@ -556,7 +556,7 @@ Use `decode_active_request_weight` when decode forward-pass time depends more on This setting is experimental and defaults to `0`, which preserves block-only decode scoring. A positive value trades some KV locality for batch-size balance and can help model, runtime, and hardware combinations near a compute-bound roofline knee, including some MLA or MTP configurations. It can regress throughput, TTFT, and ITL when decode remains primarily memory-bound, so benchmark representative traffic and start with a small weight before increasing it. -Use `prefix_affinity_weight` for multi-turn or agentic traffic where a resumed request should usually stay on the worker that already holds its prefix. The overlap credit only offsets the request's own prefill cost, so it can be worth at most about one request of load: one extra in-flight sequence on the cache-holding worker can send the next turn to a cold worker that recomputes the whole prefix. This parameter adds `prefix_affinity_weight * missing_prefix_blocks` to each candidate's score, where `missing_prefix_blocks` is the number of prompt blocks the worker does not hold on device. The term is not capped, so it keeps a session on its cache holder until the load gap exceeds the weighted cost of the missing blocks. It defaults to `0`, which preserves the existing scoring. Values between `1` and `4` worked well in a mocker replay of 16 sessions with 12 turns and a 3K-token prefix; benchmark representative traffic before raising it, since a large weight trades load balance for locality. Load-only decode routing ignores it. +Use `prefix_affinity_weight` for multi-turn or agentic traffic where a resumed request should usually stay on the worker that already holds its prefix. The overlap credit only offsets the request's own prefill cost, so it can be worth at most about one request of load: one extra in-flight sequence on the cache-holding worker can send the next turn to a cold worker that recomputes the whole prefix. This parameter adds `prefix_affinity_weight * missing_prefix_blocks` to each candidate's score, where `missing_prefix_blocks` is the number of prompt blocks the worker does not hold on device. The term is not capped, so it keeps a session on its cache holder until the load gap exceeds the weighted cost of the missing blocks. It defaults to `0`, which preserves the existing scoring. Values between `1` and `4` worked well in a mocker replay of 16 sessions with 12 turns and a 3K-token prefix; benchmark representative traffic before raising it, since a large weight trades load balance for locality. Cache-aware decode scoring, used by conditional disaggregation and by aggregated routing without prefill tracking, charges it the same way. Load-only decode routing ignores it. Use `--router-host-cache-hit-weight` and `--router-disk-cache-hit-weight` when the backend exposes lower-tier prefix cache via a KV connector (for example, vLLM's `OffloadingConnector` for CPU offload, or a disk-backed tier). These multipliers control how much each lower-tier hit credits against the prefill load, mirroring the role of `overlap_score_credit` for the device tier. A worker holding a full prefix in CPU offload gets `host_cache_hit_weight * matched_blocks` credit against its prefill cost; raising the weight makes the router more willing to route prefix-matched requests to that worker even if a different worker has a partial device-local match. diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index afaab5001039..af8fe10e99db 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -346,7 +346,14 @@ impl> DefaultWorkerScorer { // avoids inverting taint preference among negative-score workers. let overlap_adjusted_decode_blocks = (decode_cost_blocks - overlap_credit_blocks).max(0.0); - let logit = overlap_adjusted_decode_blocks + active_request_cost_blocks; + let mut logit = overlap_adjusted_decode_blocks + active_request_cost_blocks; + // The credit here only offsets decode load, so it has the same cap as the + // prefill path below; charge missing prefix blocks the same way. + if kv_router_config.prefix_affinity_weight != 0.0 { + let missing_prefix_blocks = + (context.request_blocks.max(1) as f64 - device_overlap_blocks).max(0.0); + logit += kv_router_config.prefix_affinity_weight * missing_prefix_blocks; + } // Stamped for the same reason as the two rows below: this row is emitted from the // `SchedulerQueueActor` task, so the logging layer cannot attach request identity to // it, and this branch returns early without reaching them. @@ -379,7 +386,8 @@ impl> DefaultWorkerScorer { 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 mut logit = prefill_cost_blocks + decode_cost_blocks + active_request_cost_blocks; - // Load-only decode routing (the plugin's decode subtraction path) never charges it. + // Load-only decode routing (decode without prefill tracking and with zero overlap + // credit) ignores the prompt, so it never charges missing prefix blocks. if kv_router_config.prefix_affinity_weight != 0.0 && (self.worker_type != "decode" || context.track_prefill_tokens) { diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 2509babc065a..a1ab19b5c1f1 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -67,7 +67,7 @@ use crate::http::service::metrics::generate_log_buckets; use crate::protocols::common::timing::RequestPhase; mod worker_registered; -use crate::protocols::common::timing::WORKER_TYPE_PREFILL; +use crate::protocols::common::timing::{WORKER_TYPE_DECODE, WORKER_TYPE_PREFILL}; use dynamo_kv_router::indexer::ApproximateLruStats; use dynamo_kv_router::sequences::LocalWorkerLoad; use worker_registered::RouterWorkerRegistered; @@ -1184,8 +1184,11 @@ impl RouterRequestMetrics { Some(prometheus::exponential_buckets(0.25, 2.0, 16).unwrap()), ) .expect("failed to create router_overlap_blocks_lost"); - non_max_overlap_selections_total.with_label_values(&[WORKER_TYPE_PREFILL]); - overlap_blocks_lost.with_label_values(&[WORKER_TYPE_PREFILL]); + // Aggregated routers emit under the decode label, so initialize both series. + for worker_type in [WORKER_TYPE_PREFILL, WORKER_TYPE_DECODE] { + non_max_overlap_selections_total.with_label_values(&[worker_type]); + overlap_blocks_lost.with_label_values(&[worker_type]); + } let kv_tokens_counter = |name: &str, help: &str| { metrics .create_intcountervec( diff --git a/lib/router-plugins/builtin/src/default.rs b/lib/router-plugins/builtin/src/default.rs index 34365f4ba5d6..b2f8117ff606 100644 --- a/lib/router-plugins/builtin/src/default.rs +++ b/lib/router-plugins/builtin/src/default.rs @@ -81,6 +81,7 @@ fn policy_for_role( host_cache_hit_weight: 0.0, disk_cache_hit_weight: 0.0, shared_cache_multiplier: 0.0, + prefix_affinity_weight: 0.0, ..parameters } } else { diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 2282c08dab59..516539718f29 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -181,7 +181,17 @@ impl self.decode_active_request_weight }; let logit = if self.prepared.needs_decode_subtraction { - (load.decode_cost_blocks() - credit).max(0.0) + request_cost + let logit = (load.decode_cost_blocks() - credit).max(0.0) + request_cost; + // Load-only decode routing (no overlap credit) ignores the prompt, so it never + // charges missing prefix blocks. Cache-aware decode routing (conditional + // disaggregation, or aggregated serving without prefill tracking) does: the + // credit here only offsets decode load, so it has the same cap as below. + if self.prefix_affinity_weight != 0.0 && self.prepared.overlap_credit > 0.0 { + let missing = (self.prepared.request_blocks.value - device_blocks).max(0.0); + logit + self.prefix_affinity_weight * missing + } else { + logit + } } else { let raw_tokens = if !context.tracks_prefill_tokens() { 0 diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index efc19b44fcd6..b86802f5112a 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -191,6 +191,70 @@ fn prefix_affinity_weight_keeps_resumed_prefix_on_cache_holder() { ); } +/// Cache-aware decode scoring subtracts the overlap credit from decode load instead of +/// prefill, so the credit has the same cap there. Conditional disaggregation and aggregated +/// serving without prefill tracking both take that path and must charge missing blocks too. +/// Plain decode routing stays load-only and never asks for cache inputs. +#[test] +fn prefix_affinity_weight_applies_to_cache_aware_decode_scoring() { + let (workers, mut request) = fixture(2, 47 * 64); + request.track_prefill_tokens = false; + for (worker, load) in &mut request.worker_loads { + let holder = worker.worker_id == 0; + *load = Default::default(); + load.active_decode_blocks = if holder { 50 } else { 0 }; + request + .overlap + .tier_overlap_blocks + .device + .insert(*worker, if holder { 46 } else { 0 }); + } + for (role, conditional_disagg_enabled) in + [(WorkerType::Decode, true), (WorkerType::Aggregated, false)] + { + let select = |prefix_affinity_weight: f64| { + let config = KvRouterConfig { + conditional_disagg_enabled, + host_cache_hit_weight: 0.0, + disk_cache_hit_weight: 0.0, + router_temperature: 0.0, + prefix_affinity_weight, + ..Default::default() + }; + dynamo_custom_policy_builtin::default_factory()( + &config, + role, + RoutingPartitionRef::new("model", "default"), + ) + .select_worker(selection_input(&workers, &request, 64)) + .unwrap() + .worker + .worker_id + }; + assert_eq!(select(0.0), 1, "{role:?}: default weights move the session"); + assert_eq!(select(1.0), 0, "{role:?}: affinity keeps the cache holder"); + } + + let plain_decode = dynamo_custom_policy_builtin::default_factory()( + &KvRouterConfig { + prefix_affinity_weight: 1.0, + ..Default::default() + }, + WorkerType::Decode, + RoutingPartitionRef::new("model", "default"), + ); + assert!( + !>::required_worker_inputs( + &plain_decode + ) + .contains(WorkerInputs::CACHE) + ); + assert_eq!( + plain_decode.runtime_config().unwrap().parameters()["prefix_affinity_weight"], + 0.0 + ); +} + #[test] fn runtime_parameters_cannot_enable_worker_signals_missing_at_startup() { use dynamo_kv_router::plugins::worker_selection::{RuntimePolicyError, RuntimePolicyUpdate}; @@ -231,6 +295,7 @@ fn runtime_parameters_cannot_enable_worker_signals_missing_at_startup() { "host_cache_hit_weight", "disk_cache_hit_weight", "shared_cache_multiplier", + "prefix_affinity_weight", ] { assert_eq!(initial.parameters()[name], 0.0); assert!( @@ -335,52 +400,55 @@ fn runtime_cache_tier_weights_change_selection_without_device_credit() { #[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..64 { - 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: Some(if mode & 16 == 0 { 0.0 } else { 0.6 }), - prefix_affinity_weight: if mode & 32 == 0 { 0.0 } else { 0.9 }, - ..Default::default() - }; - if mode % 8 >= 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, - _ => {} - } - 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 = support::selection_input(&workers, &request, 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 + // "decode" also covers the cache-aware decode subtraction path when prefill tracking is off. + for label in ["test", "decode"] { + for temperature in [0.0, 0.7] { + for prompt in [1, 17, 127, 2048] { + for mode in 0..64 { + 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: Some(if mode & 16 == 0 { 0.0 } else { 0.6 }), + prefix_affinity_weight: if mode & 32 == 0 { 0.0 } else { 0.9 }, + ..Default::default() + }; + if mode % 8 >= 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, + _ => {} + } + let reference = dynamo_kv_router::DefaultWorkerSelector::new_seeded( + Some(config.clone()), + label, + 42, ); + let plugin = DefaultWorkerSelector::new_seeded(Some(config), label, 42); + for _ in 0..64 { + let input = support::selection_input(&workers, &request, 16); + let expected = reference.select_worker(input).unwrap(); + let actual = plugin.select_worker(input).unwrap(); + assert_eq!( + actual.worker, expected.worker, + "label={label} temperature={temperature} prompt={prompt} mode={mode}" + ); + assert_eq!(actual.cached_tokens, expected.cached_tokens); + assert_eq!( + actual.potential_decode_blocks, + expected.potential_decode_blocks + ); + } } } } From 5e1ab89179cc9bcf559167fb50058f105475ae2a Mon Sep 17 00:00:00 2001 From: Dundy Pasupuleti Date: Thu, 22 Oct 2026 11:13:45 -0400 Subject: [PATCH 3/4] fix(router): skip prefix affinity charge for empty prompts The affinity penalty reused the request block count clamped to one, so an empty prompt was charged a phantom missing block that the preferred-taint multiplier then scaled per worker. Use the unclamped count in both the plugin and reference scorers and add a regression test. Restrict the decode parity loop to the cache-aware decode mode and drop a restating comment. Signed-off-by: Dundy Pasupuleti --- .../src/scheduling/selector/reference.rs | 4 +- .../builtin/src/default/scorer.rs | 7 +- .../builtin/tests/default_policy.rs | 79 ++++++++++++++++++- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index af8fe10e99db..e7d7310c54d7 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -351,7 +351,7 @@ impl> DefaultWorkerScorer { // prefill path below; charge missing prefix blocks the same way. if kv_router_config.prefix_affinity_weight != 0.0 { let missing_prefix_blocks = - (context.request_blocks.max(1) as f64 - device_overlap_blocks).max(0.0); + (context.request_blocks as f64 - device_overlap_blocks).max(0.0); logit += kv_router_config.prefix_affinity_weight * missing_prefix_blocks; } // Stamped for the same reason as the two rows below: this row is emitted from the @@ -395,7 +395,7 @@ impl> DefaultWorkerScorer { // request's own prefill, so it alone cannot outweigh transient decode load on the // worker that already holds a resumed prefix. let missing_prefix_blocks = - (context.request_blocks.max(1) as f64 - device_overlap_blocks).max(0.0); + (context.request_blocks as f64 - device_overlap_blocks).max(0.0); logit += kv_router_config.prefix_affinity_weight * missing_prefix_blocks; } diff --git a/lib/router-plugins/builtin/src/default/scorer.rs b/lib/router-plugins/builtin/src/default/scorer.rs index 516539718f29..f0a286a9c408 100644 --- a/lib/router-plugins/builtin/src/default/scorer.rs +++ b/lib/router-plugins/builtin/src/default/scorer.rs @@ -46,6 +46,8 @@ struct PreparedRequest { min_prefill: usize, block_size: IntegerDivisor, request_blocks: IntegerDivisor, + /// Unclamped prompt block count; the affinity penalty charges zero for an empty prompt. + prompt_blocks: f64, overlap_credit: f64, needs_decay: bool, needs_decode_subtraction: bool, @@ -116,6 +118,7 @@ impl // An empty prompt (e.g. embeddings-only input) has no overlap to decay; // clamp so a worker at the load floor does not compute 0/0. request_blocks: IntegerDivisor::new(context.request_blocks().max(1)), + prompt_blocks: context.request_blocks() as f64, overlap_credit: self.overlap_score_credit, needs_decay: context.tracks_prefill_tokens() && self.overlap_score_credit_decay > 0.0, needs_decode_subtraction: self.is_decode && !context.tracks_prefill_tokens(), @@ -187,7 +190,7 @@ impl // disaggregation, or aggregated serving without prefill tracking) does: the // credit here only offsets decode load, so it has the same cap as below. if self.prefix_affinity_weight != 0.0 && self.prepared.overlap_credit > 0.0 { - let missing = (self.prepared.request_blocks.value - device_blocks).max(0.0); + let missing = (self.prepared.prompt_blocks - device_blocks).max(0.0); logit + self.prefix_affinity_weight * missing } else { logit @@ -208,7 +211,7 @@ impl // Charge each prompt block this worker lacks. The overlap credit above is // capped at the request's own prefill, so it cannot outweigh transient // decode load on the worker that already holds a resumed prefix. - let missing = (self.prepared.request_blocks.value - device_blocks).max(0.0); + let missing = (self.prepared.prompt_blocks - device_blocks).max(0.0); logit + self.prefix_affinity_weight * missing } else { logit diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index b86802f5112a..d8efe8cf1a41 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -147,7 +147,6 @@ fn prefix_affinity_weight_keeps_resumed_prefix_on_cache_holder() { }), 0 ); - // The weight is also a runtime policy parameter. let mut policy = default_policy(base, "prefill"); let revision = policy.runtime_config().unwrap().revision(); let patch = json!({"prefix_affinity_weight": 0.25}) @@ -398,6 +397,81 @@ fn runtime_cache_tier_weights_change_selection_without_device_credit() { } } +#[derive(Clone)] +struct TaintedWorker(HashSet); +impl dynamo_kv_router::protocols::WorkerConfigLike for TaintedWorker { + 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) + } + fn taints(&self) -> &HashSet { + &self.0 + } +} + +/// An empty prompt has no prefix to miss, so the affinity weight must not charge a block. +/// The preferred-taint multiplier scales that phantom charge differently per worker, which +/// would flip the pick toward the worker with more decode load but a smaller multiplier. +#[test] +fn prefix_affinity_weight_charges_nothing_for_empty_prompt() { + let (_, mut request) = fixture(2, 0); + request + .routing_constraints + .preferred_taints + .insert("slow".into(), -1.0); + let workers: HashMap = HashMap::from([ + (0, TaintedWorker(HashSet::from(["slow".to_owned()]))), + (1, TaintedWorker(HashSet::new())), + ]); + 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 { 1 } else { 3 }; + load.additional_active_blocks = 0; + request + .overlap + .tier_overlap_blocks + .device + .insert(*worker, 0); + } + for label in ["prefill", "decode"] { + for track_prefill_tokens in [true, false] { + request.track_prefill_tokens = track_prefill_tokens; + for weight in [0.0, 4.0] { + let config = KvRouterConfig { + host_cache_hit_weight: 0.0, + disk_cache_hit_weight: 0.0, + router_temperature: 0.0, + prefix_affinity_weight: weight, + ..Default::default() + }; + let input = || { + WorkerSelectionInput::configured(&workers, &request, request.eligibility(), 16) + }; + let reference = dynamo_kv_router::DefaultWorkerSelector::new_seeded( + Some(config.clone()), + label, + 42, + ); + let plugin = DefaultWorkerSelector::new_seeded(Some(config), label, 42); + let expected = reference.select_worker(input()).unwrap().worker; + let actual = plugin.select_worker(input()).unwrap().worker; + let case = format!("label={label} track={track_prefill_tokens} weight={weight}"); + assert_eq!(actual, expected, "{case}"); + assert_eq!(actual.worker_id, 0, "{case}"); + } + } + } +} + #[test] fn seeded_selection_matches_reference_across_cache_and_load_shapes() { // "decode" also covers the cache-aware decode subtraction path when prefill tracking is off. @@ -405,6 +479,9 @@ 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..64 { + if label == "decode" && mode % 4 != 3 { + continue; + } let (workers, mut request) = fixture(16, prompt); let config = KvRouterConfig { router_temperature: temperature, From 360c6555b290b0b7f744a2d261a2c3884799f1f8 Mon Sep 17 00:00:00 2001 From: Dundy Pasupuleti Date: Tue, 27 Oct 2026 20:47:31 -0400 Subject: [PATCH 4/4] test(router): count prefix_affinity_weight in startup policy parameters Signed-off-by: Dundy Pasupuleti --- lib/router-plugins/builtin/tests/default_policy.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lib/router-plugins/builtin/tests/default_policy.rs b/lib/router-plugins/builtin/tests/default_policy.rs index 1c47cf45238f..4d02cd78e682 100644 --- a/lib/router-plugins/builtin/tests/default_policy.rs +++ b/lib/router-plugins/builtin/tests/default_policy.rs @@ -1121,6 +1121,7 @@ fn startup_policy_parameters_match_the_default_factory_for_every_role() { host_cache_hit_weight: 0.5, disk_cache_hit_weight: 0.2, shared_cache_multiplier: Some(0.3), + prefix_affinity_weight: 0.25, router_temperature: 0.7, ..Default::default() }, @@ -1152,7 +1153,7 @@ fn startup_policy_parameters_match_the_default_factory_for_every_role() { expected.parameters(), "role={role:?} conditional_disagg={conditional_disagg_enabled}" ); - assert_eq!(startup.len(), 8); + assert_eq!(startup.len(), 9); assert!(startup.values().all(serde_json::Value::is_number)); } }