feat(router): ship opt-in soft affinity load guard

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

@@ -112,6 +112,7 @@ link no catalog and reject a configured policy type at startup.
 |---|---|
 | `default` | Dynamo's built-in selector and cost model. Reserved; always available. |
 | `dynamo-two-tier-cost-fn` | Ranks on two tiers instead of one additive cost: active-request load first, then device-KV prefix overlap. Prefers the worker holding the largest prefix overlap unless load is badly imbalanced. Thresholds and selection order ported from the experimental SGLang router's `cache_aware_zmq` policy. Thresholds are tunable; the defaults reproduce it exactly. |
+| `dynamo-soft-affinity-load-guard` | Retains an eligible soft-affinity target until its active-request count exceeds a threshold and another worker has a sufficient load advantage. Without a target, selects the least-loaded eligible worker. |
 
 Write the instance into the same YAML file that `--router-policy-config` already points at:
 
@@ -160,6 +161,36 @@ rather than being silently ignored. It selects the least-loaded worker once 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.
 
+#### Guard Soft Affinity Under Load
+
+Use `dynamo-soft-affinity-load-guard` with soft session affinity to retain a session's worker while
+its active-request count is at or below `max_active_requests`. Above that threshold, the policy
+moves only when the least-loaded alternative has at least `move_margin` fewer active requests.
+The same rule applies when requests carrying a parent session ID share the parent-group binding.
+
+```yaml
+worker_selection:
+  aggregated: soft-guard
+  instances:
+    - name: soft-guard
+      type: dynamo-soft-affinity-load-guard
+      parameters:
+        max_active_requests: 32
+        move_margin: 2
+```
+
+Run the frontend with `--router-mode kv`, `--router-session-affinity-ttl-secs 60`,
+`--router-session-affinity-mode soft`, and `--router-policy-config` pointing to that YAML file.
+Both parameters are optional; the values above are the defaults. `move_margin` must be at least
+one. When the affinity target is missing from the eligible candidate set, the policy falls back
+to the least-loaded worker. If no better alternative exists, it retains the target. Equal-load
+workers never trigger a move. Hard affinity and explicit worker targets remain exact and bypass
+this advisory choice.
+
+The active-request count is a routing load signal, not a prediction of queue wait or remaining
+work. This policy does not inspect cache residency. Moving to a cold worker may recompute a prefix;
+choose thresholds with a workload comparison of request latency and cached-token usage.
+
 #### Override the Selection
 
 `DYN_ROUTER_WORKER_SELECTION_POLICY` overrides every stage. `--router-prefill-policy` and

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

@@ -18,6 +18,7 @@
 //! 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.
 
+mod soft_affinity_load_guard;
 mod two_tier_cost_fn;
 
 use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegistryError};
@@ -30,7 +31,8 @@ use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegis
 pub fn register(
     registry: &mut RouterPluginRegistry,
 ) -> Result<(), WorkerSelectionPolicyRegistryError> {
-    two_tier_cost_fn::register(registry)
+    two_tier_cost_fn::register(registry)?;
+    soft_affinity_load_guard::register(registry)
 }
 
 #[cfg(test)]
@@ -90,6 +92,28 @@ worker_selection:
         }
     }
 
+    #[test]
+    fn resolves_soft_affinity_load_guard_yaml() {
+        let (config, resolved) = resolve(
+            r#"
+worker_selection:
+  aggregated: soft-guard
+  instances:
+    - name: soft-guard
+      type: dynamo-soft-affinity-load-guard
+      parameters:
+        max_active_requests: 7
+        move_margin: 2
+"#,
+        );
+        let factory = resolved.unwrap().expect("the shipped policy resolves");
+        factory(
+            &config,
+            WorkerType::Aggregated,
+            RoutingPartitionRef::new("model", "default"),
+        );
+    }
+
     /// 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]

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

@@ -0,0 +1,261 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Retain an eligible soft-affinity target until its active load warrants moving.
+//!
+//! Dynamo owns the affinity binding. This picker only selects an eligible worker; a successful
+//! dispatch updates a soft binding through the normal coordinator path.
+
+use std::sync::Arc;
+
+use dynamo_kv_router::plugins::worker_selection::WorkerSelectionPolicyFactory;
+use dynamo_kv_router::plugins::{
+    RouterPluginRegistry, WorkerSelectionPolicyParameters, WorkerSelectionPolicyProviderError,
+    WorkerSelectionPolicyRegistryError,
+};
+use dynamo_kv_router::protocols::WorkerAffinityTarget;
+use dynamo_kv_router::{
+    KvRouterConfig, ScoredWorkerCandidate, WorkerInputView, WorkerInputs, WorkerLoadInput,
+    WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError,
+};
+
+pub const POLICY_TYPE: &str = "dynamo-soft-affinity-load-guard";
+
+#[derive(Clone, Copy, serde::Deserialize)]
+#[serde(deny_unknown_fields, default)]
+struct Parameters {
+    max_active_requests: usize,
+    move_margin: usize,
+}
+
+impl Default for Parameters {
+    fn default() -> Self {
+        Self {
+            max_active_requests: 32,
+            move_margin: 2,
+        }
+    }
+}
+
+impl Parameters {
+    fn validate(self) -> Result<Self, WorkerSelectionPolicyProviderError> {
+        if self.move_margin == 0 {
+            return Err(WorkerSelectionPolicyProviderError::new(
+                "move_margin must be at least 1",
+            ));
+        }
+        Ok(self)
+    }
+}
+
+struct SoftAffinityLoadGuardPicker {
+    parameters: Parameters,
+}
+
+fn matches_target(candidate: &ScoredWorkerCandidate, target: WorkerAffinityTarget) -> bool {
+    let worker = candidate.worker();
+    worker.worker_id == target.worker_id && target.dp_rank.is_none_or(|rank| worker.dp_rank == rank)
+}
+
+fn least_loaded(
+    candidates: &[ScoredWorkerCandidate],
+    loads: &[WorkerLoadInput],
+    include: impl Fn(&ScoredWorkerCandidate) -> bool,
+) -> Option<usize> {
+    candidates
+        .iter()
+        .zip(loads)
+        .enumerate()
+        .filter(|(_, (candidate, _))| include(candidate))
+        .min_by_key(|(_, (candidate, load))| (load.active_requests(), candidate.worker()))
+        .map(|(row, _)| row)
+}
+
+impl WorkerPicker for SoftAffinityLoadGuardPicker {
+    fn required_worker_inputs(&self) -> WorkerInputs {
+        WorkerInputs::LOAD
+    }
+
+    fn pick(
+        &mut self,
+        context: &WorkerSelectionContext<'_>,
+        input: WorkerInputView<'_>,
+    ) -> Result<usize, WorkerSelectionPolicyError> {
+        let candidates = input.candidates();
+        let loads = input
+            .load()
+            .ok_or_else(|| WorkerSelectionPolicyError::failed("active load input unavailable"))?;
+
+        if let Some(target) = context.affinity_target()
+            && let Some(target_row) = least_loaded(candidates, loads, |candidate| {
+                matches_target(candidate, target)
+            })
+        {
+            let target_load = loads[target_row].active_requests();
+            if target_load <= self.parameters.max_active_requests {
+                return Ok(target_row);
+            }
+
+            if let Some(alternative_row) = least_loaded(candidates, loads, |candidate| {
+                !matches_target(candidate, target)
+            }) {
+                let alternative_load = loads[alternative_row].active_requests();
+                if alternative_load < target_load
+                    && target_load - alternative_load >= self.parameters.move_margin
+                {
+                    return Ok(alternative_row);
+                }
+            }
+            return Ok(target_row);
+        }
+
+        least_loaded(candidates, loads, |_| true)
+            .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker"))
+    }
+}
+
+fn provider(
+    parameters: &WorkerSelectionPolicyParameters,
+) -> Result<WorkerSelectionPolicyFactory, WorkerSelectionPolicyProviderError> {
+    let parameters: Parameters = parameters.deserialize()?;
+    let parameters = parameters.validate()?;
+
+    Ok(Arc::new(
+        move |config: &KvRouterConfig, worker_type, _partition| {
+            WorkerSelectionPolicy::new(
+                config.clone(),
+                worker_type.as_str(),
+                Vec::new(),
+                Box::new(SoftAffinityLoadGuardPicker { parameters }),
+            )
+        },
+    ))
+}
+
+pub fn register(
+    registry: &mut RouterPluginRegistry,
+) -> Result<(), WorkerSelectionPolicyRegistryError> {
+    registry.register_worker_selection(POLICY_TYPE, Arc::new(provider))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank};
+    use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode};
+    use dynamo_kv_router::{
+        SchedulingRequest, WorkerLoadProjection, WorkerSelectionInput, WorkerSelector,
+    };
+
+    use super::*;
+
+    struct TestWorker;
+
+    impl WorkerConfigLike for TestWorker {
+        fn data_parallel_start_rank(&self) -> u32 {
+            0
+        }
+
+        fn data_parallel_size(&self) -> u32 {
+            1
+        }
+
+        fn max_num_batched_tokens(&self) -> Option<u64> {
+            None
+        }
+
+        fn total_kv_blocks(&self) -> Option<u64> {
+            Some(1024)
+        }
+    }
+
+    fn worker(id: u64) -> WorkerWithDpRank {
+        WorkerWithDpRank::from_worker_id(id)
+    }
+
+    fn select(parameters: Parameters, target: Option<u64>, a_load: usize, b_load: usize) -> u64 {
+        let mut request = SchedulingRequest {
+            mode: ScheduleMode::QueryOnly { request_id: None },
+            token_seq: None,
+            isl_tokens: 16,
+            lora_name: None,
+            expected_output_tokens: None,
+            affinity_target: target.map(|id| worker(id).into()),
+            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,
+        };
+        for (id, active_requests) in [(7, a_load), (9, b_load)] {
+            request.worker_loads.insert(
+                worker(id),
+                WorkerLoadProjection {
+                    active_requests,
+                    ..Default::default()
+                },
+            );
+        }
+        let workers = HashMap::from([(7, TestWorker), (9, TestWorker)]);
+        WorkerSelectionPolicy::new(
+            KvRouterConfig::default(),
+            "test",
+            Vec::new(),
+            Box::new(SoftAffinityLoadGuardPicker { parameters }),
+        )
+        .select_worker(WorkerSelectionInput::configured(
+            &workers,
+            &request,
+            request.eligibility(),
+            16,
+        ))
+        .unwrap()
+        .worker
+        .worker_id
+    }
+
+    #[test]
+    fn retains_the_group_until_both_load_conditions_hold() {
+        let parameters = Parameters {
+            max_active_requests: 7,
+            move_margin: 2,
+        };
+        assert_eq!(select(parameters, Some(7), 7, 0), 7);
+        assert_eq!(select(parameters, Some(7), 8, 7), 7);
+        assert_eq!(select(parameters, Some(7), 8, 6), 9);
+        assert_eq!(select(parameters, Some(9), 6, 8), 7);
+        assert_eq!(select(parameters, Some(9), 7, 8), 9);
+    }
+
+    #[test]
+    fn selects_an_eligible_worker_when_no_group_target_exists() {
+        let parameters = Parameters::default();
+        assert_eq!(select(parameters, None, 3, 0), 9);
+        assert_eq!(select(parameters, Some(99), 3, 0), 9);
+        assert_eq!(select(parameters, Some(7), 33, 33), 7);
+        assert_eq!(select(parameters, Some(7), 33, 34), 7);
+    }
+
+    #[test]
+    fn rejects_zero_move_margin() {
+        assert!(
+            Parameters {
+                move_margin: 0,
+                ..Parameters::default()
+            }
+            .validate()
+            .is_err()
+        );
+    }
+}