feat(router): add prefix_affinity_weight and report non-max-overlap picks in aggregated mode

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. 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.
 
 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

docs/fern/pages/reference/observability/metrics-catalog.mdx

@@ -367,11 +367,11 @@ The router exposes metrics for routing decisions and overhead. Not every metric
 </ParamField>
 
 <ParamField path="dynamo_component_router_non_max_overlap_selections_total" type="counter">
-  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`.
 </ParamField>
 
 <ParamField path="dynamo_component_router_overlap_blocks_lost" type="histogram">
-  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`.
 </ParamField>
 
 ### KV-cache reuse telemetry

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,

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<f64>,
     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<KvRouterConfigSerde> 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 {

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

@@ -346,7 +346,14 @@ impl<C: Borrow<KvRouterConfig>> DefaultWorkerScorer<C> {
             // 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.
@@ -378,7 +385,19 @@ impl<C: Borrow<KvRouterConfig>> DefaultWorkerScorer<C> {
         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 (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)
+        {
+            // 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

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(

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;
@@ -1170,22 +1170,25 @@ 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,
             )
             .expect("failed to create router_non_max_overlap_selections_total");
         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()),
             )
             .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(

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 {

lib/router-plugins/builtin/src/default/parameters.rs

@@ -19,6 +19,7 @@ struct Parameters {
     overlap_score_credit_decay: Option<f64>,
     prefill_load_scale: Option<f64>,
     decode_active_request_weight: Option<f64>,
+    prefix_affinity_weight: Option<f64>,
     host_cache_hit_weight: Option<f64>,
     disk_cache_hit_weight: Option<f64>,
     shared_cache_multiplier: Option<f64>,
@@ -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)),