feat(router): configure two-tier soft affinity

docs/fern/pages/cli/kv-aware-routing/overview.mdx

@@ -96,6 +96,8 @@ worker_selection:
   instances:
     - name: dynamo-two-tier-cost-fn
       type: dynamo-two-tier-cost-fn
+      parameters:
+        respect_soft_affinity: true
 ```
 
 Then point the frontend at it:
@@ -109,6 +111,11 @@ python3 -m dynamo.frontend \
 To compare against the built-in selector, restart with
 `DYN_ROUTER_WORKER_SELECTION_POLICY=default` — no config change needed.
 
+`respect_soft_affinity: true` makes the two-tier picker retain a soft session-affinity target when
+that target is among its eligible candidates. The default is `false`, which lets cache and load
+choose among the full eligible set. An absent target falls back to normal two-tier selection; hard
+affinity and explicit worker targets remain exact.
+
 For the available policy types and per-stage prefill/decode selection, see
 [Worker-Selection Policies](../../developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md#worker-selection-policies).
 

docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md

@@ -147,20 +147,28 @@ experimental router exactly. Add any subset to tune it:
         cache_threshold: 0.5
         balance_abs_threshold: 32
         balance_rel_threshold: 1.1
+        respect_soft_affinity: true
 ```
 
 | Parameter | Default | Meaning |
 |---|---|---|
 | `cache_threshold` | `0.5` | Fraction of the request's blocks that must be device-resident on the best worker before the cache tier applies. Compared strictly. Must be in `[0.0, 1.0]`. |
 | `balance_abs_threshold` | `32` | Minimum active-request spread before the load tier applies. |
 | `balance_rel_threshold` | `1.1` | Minimum ratio of largest to smallest active-request count before the load tier applies. Must be at least `1.0`. |
+| `respect_soft_affinity` | `false` | When `true`, the two-tier picker retains a soft session-affinity target found in its eligible candidates. When `false`, cache and load choose among the full eligible set. |
 
 Both load gates must hold before the load tier displaces the cache tier. Parameters are validated at
 startup, so an out-of-range value or an unknown key fails the process immediately, naming the key,
 rather than being silently ignored. It selects the least-loaded worker once the active-request spread is greater than 32 and the
 largest count is more than 1.1 times the smallest; otherwise it prefers the worker holding the
 largest device-KV overlap when that overlap covers more than 50% of the request's blocks.
 
+`respect_soft_affinity` affects soft affinity only. Custom policies always receive the full eligible
+candidate set plus the affinity target as context; this parameter tells the two-tier picker whether
+to retain a matching target before considering other workers. If the target is absent from that
+set, the policy falls back to normal two-tier selection. Hard session affinity and explicit request
+targets remain exact constraints for either value.
+
 #### Override the Selection
 
 `DYN_ROUTER_WORKER_SELECTION_POLICY` overrides every stage. `--router-prefill-policy` and

lib/kv-router/src/scheduling/AGENTS.md

@@ -122,6 +122,12 @@ PolicyClassQueue("agents")
 
 ## Public Worker-Selection API
 
+- `WorkerSelector::uses_exclusive_affinity_target()` is a one-off compatibility hook that
+  preserves the default selector's historical soft-affinity behavior. Custom
+  `WorkerSelectionPolicy` implementations must not override or make it configurable. They receive
+  the full eligible candidate set and use `WorkerSelectionContext::affinity_target()` to retain,
+  ignore, or replace a soft target.
+
 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.

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

@@ -90,6 +90,33 @@ worker_selection:
         }
     }
 
+    #[test]
+    fn resolves_soft_affinity_parameter() {
+        let yaml = |parameter: &str| {
+            format!(
+                r#"
+worker_selection:
+  aggregated: two-tier
+  instances:
+    - name: two-tier
+      type: dynamo-two-tier-cost-fn
+{parameter}
+"#
+            )
+        };
+
+        for parameter in [
+            "      parameters:\n        respect_soft_affinity: false",
+            "      parameters:\n        respect_soft_affinity: true",
+        ] {
+            let (_config, resolved) = resolve(&yaml(parameter));
+            assert!(
+                resolved.unwrap().is_some(),
+                "parameter variant should resolve: {parameter}"
+            );
+        }
+    }
+
     /// An unknown parameter key is a mistake, most often a misremembered threshold name. It must
     /// fail startup rather than silently leaving the default in place.
     #[test]
@@ -102,7 +129,7 @@ worker_selection:
     - name: dynamo-two-tier-cost-fn
       type: dynamo-two-tier-cost-fn
       parameters:
-        cache_affinity_threshold: 0.8
+        respect_soft_affinitty: true
 "#,
         );
 
@@ -115,7 +142,7 @@ worker_selection:
             "unexpected error: {error}"
         );
         assert!(
-            error.to_string().contains("cache_affinity_threshold"),
+            error.to_string().contains("respect_soft_affinitty"),
             "the error should name the offending key: {error}"
         );
     }

lib/router-plugins/builtin/src/two_tier_cost_fn.rs

@@ -48,6 +48,7 @@ pub const POLICY_TYPE: &str = "dynamo-two-tier-cost-fn";
 const DEFAULT_CACHE_THRESHOLD: f64 = 0.5;
 const DEFAULT_BALANCE_ABS_THRESHOLD: usize = 32;
 const DEFAULT_BALANCE_REL_THRESHOLD: f64 = 1.1;
+const DEFAULT_RESPECT_SOFT_AFFINITY: bool = false;
 
 /// Tunables for [`POLICY_TYPE`], named after their `sgl-router` counterparts.
 ///
@@ -63,6 +64,8 @@ struct Parameters {
     balance_abs_threshold: usize,
     /// Minimum ratio of largest to smallest active-request count before the load tier applies.
     balance_rel_threshold: f64,
+    /// Whether to retain an eligible soft-affinity target before considering other workers.
+    respect_soft_affinity: bool,
 }
 
 impl Default for Parameters {
@@ -71,6 +74,7 @@ impl Default for Parameters {
             cache_threshold: DEFAULT_CACHE_THRESHOLD,
             balance_abs_threshold: DEFAULT_BALANCE_ABS_THRESHOLD,
             balance_rel_threshold: DEFAULT_BALANCE_REL_THRESHOLD,
+            respect_soft_affinity: DEFAULT_RESPECT_SOFT_AFFINITY,
         }
     }
 }
@@ -95,28 +99,33 @@ fn least_loaded(load: &[WorkerLoadInput], rows: impl Iterator<Item = usize>) ->
     rows.min_by_key(|&row| load[row].active_requests())
 }
 
-fn select_row(
+fn select_row_where(
     parameters: &Parameters,
     cache: WorkerCacheInputs<'_>,
     load: &[WorkerLoadInput],
     request_blocks: u64,
+    include: impl Fn(usize) -> bool + Copy,
 ) -> Option<usize> {
     if cache.is_empty() || cache.len() != load.len() {
         return None;
     }
 
-    let min_load = load.iter().map(|item| item.active_requests()).min()?;
-    let max_load = load.iter().map(|item| item.active_requests()).max()?;
+    let device_overlap = |row| {
+        cache
+            .get(row)
+            .expect("cache and load lengths were checked")
+            .device_overlap_blocks()
+    };
+    let rows = || (0..load.len()).filter(|&row| include(row));
+    let min_load = rows().map(|row| load[row].active_requests()).min()?;
+    let max_load = rows().map(|row| load[row].active_requests()).max()?;
     if max_load.saturating_sub(min_load) > parameters.balance_abs_threshold
         && (max_load as f64) > parameters.balance_rel_threshold * (min_load as f64)
     {
-        return least_loaded(load, 0..load.len());
+        return least_loaded(load, rows());
     }
 
-    let max_overlap = cache
-        .iter()
-        .map(|item| item.device_overlap_blocks())
-        .max_by(f64::total_cmp)?;
+    let max_overlap = rows().map(device_overlap).max_by(f64::total_cmp)?;
     let cache_ratio = if request_blocks == 0 {
         0.0
     } else {
@@ -125,13 +134,20 @@ fn select_row(
     if cache_ratio > parameters.cache_threshold {
         return least_loaded(
             load,
-            cache.iter().enumerate().filter_map(|(row, item)| {
-                (item.device_overlap_blocks() == max_overlap).then_some(row)
-            }),
+            rows().filter(|&row| device_overlap(row) == max_overlap),
         );
     }
 
-    least_loaded(load, 0..load.len())
+    least_loaded(load, rows())
+}
+
+fn select_row(
+    parameters: &Parameters,
+    cache: WorkerCacheInputs<'_>,
+    load: &[WorkerLoadInput],
+    request_blocks: u64,
+) -> Option<usize> {
+    select_row_where(parameters, cache, load, request_blocks, |_| true)
 }
 
 struct TwoTierCostFnPicker {
@@ -154,6 +170,24 @@ impl WorkerPicker for TwoTierCostFnPicker {
         let load = input
             .load()
             .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?;
+
+        if self.parameters.respect_soft_affinity
+            && let Some(target) = context.affinity_target()
+            && let Some(row) = select_row_where(
+                &self.parameters,
+                cache,
+                load,
+                context.request_blocks(),
+                |row| {
+                    let worker = input.candidates()[row].worker();
+                    worker.worker_id == target.worker_id
+                        && target.dp_rank.is_none_or(|rank| worker.dp_rank == rank)
+                },
+            )
+        {
+            return Ok(row);
+        }
+
         select_row(&self.parameters, cache, load, context.request_blocks())
             .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker"))
     }
@@ -187,7 +221,9 @@ pub fn register(
 mod tests {
     use std::collections::HashMap;
 
-    use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank};
+    use dynamo_kv_router::protocols::{
+        RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, WorkerWithDpRank,
+    };
     use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode};
     use dynamo_kv_router::{
         SchedulingRequest, WorkerLoadProjection, WorkerSelectionInput, WorkerSelector,
@@ -224,17 +260,21 @@ mod tests {
 
     /// Select among workers given as `(worker_id, device_overlap_blocks, active_requests)`.
     fn select(workers: [(u64, usize, usize); 2]) -> WorkerWithDpRank {
-        select_with(Parameters::default(), workers)
+        select_with(Parameters::default(), workers, None)
     }
 
-    fn select_with(parameters: Parameters, workers: [(u64, usize, usize); 2]) -> WorkerWithDpRank {
+    fn select_with(
+        parameters: Parameters,
+        workers: [(u64, usize, usize); 2],
+        affinity_target: Option<WorkerAffinityTarget>,
+    ) -> WorkerWithDpRank {
         let mut request = SchedulingRequest {
             mode: ScheduleMode::QueryOnly { request_id: None },
             token_seq: None,
             isl_tokens: TEN_BLOCKS,
             lora_name: None,
             expected_output_tokens: None,
-            affinity_target: None,
+            affinity_target,
             pinned_worker: None,
             allowed_worker_ids: None,
             routing_constraints: RoutingConstraints::default(),
@@ -304,7 +344,58 @@ mod tests {
             cache_threshold: 0.2,
             ..Parameters::default()
         };
-        assert_eq!(select_with(tuned, workers), worker(B));
+        assert_eq!(select_with(tuned, workers, None), worker(B));
+    }
+
+    #[test]
+    fn soft_affinity_is_policy_owned() {
+        let workers = [(A, 0, 0), (B, 6, 4)];
+        let target = Some(WorkerAffinityTarget::new(A, None));
+
+        assert_eq!(
+            select_with(Parameters::default(), workers, target),
+            worker(B),
+            "the default ignores advisory soft affinity"
+        );
+
+        let respecting = Parameters {
+            respect_soft_affinity: true,
+            ..Parameters::default()
+        };
+        assert_eq!(
+            select_with(respecting, workers, target),
+            worker(A),
+            "an eligible target is retained by the picker"
+        );
+        assert_eq!(
+            select_with(
+                respecting,
+                workers,
+                Some(WorkerAffinityTarget::new(999, None)),
+            ),
+            worker(B),
+            "an absent target falls back to normal two-tier selection"
+        );
+
+        let rank_sensitive_workers = [(A, 0, 0), (B, 0, 4)];
+        assert_eq!(
+            select_with(
+                respecting,
+                rank_sensitive_workers,
+                Some(WorkerAffinityTarget::new(B, Some(0))),
+            ),
+            worker(B),
+            "a matching rank-specific target is retained"
+        );
+        assert_eq!(
+            select_with(
+                respecting,
+                rank_sensitive_workers,
+                Some(WorkerAffinityTarget::new(B, Some(1))),
+            ),
+            worker(A),
+            "a missing target rank falls back instead of retaining another rank"
+        );
     }
 
     #[test]