diff --git a/docs/fern/pages/cli/kv-aware-routing/overview.mdx b/docs/fern/pages/cli/kv-aware-routing/overview.mdx index 6f483149b63a..0b7dbdf45301 100644 --- a/docs/fern/pages/cli/kv-aware-routing/overview.mdx +++ b/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). 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 0b8ec9d9c642..71354d5d1d8f 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 @@ -147,6 +147,7 @@ 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 | @@ -154,6 +155,7 @@ experimental router exactly. Add any subset to tune it: | `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, @@ -161,6 +163,12 @@ 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. +`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 diff --git a/lib/kv-router/src/scheduling/AGENTS.md b/lib/kv-router/src/scheduling/AGENTS.md index 22d6d5289467..9abf8f7a315f 100644 --- a/lib/kv-router/src/scheduling/AGENTS.md +++ b/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. diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index 1722ff618a91..a1b19f9698ae 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/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}" ); } 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 e77eda8b120c..cab66b8fb775 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/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) -> 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 { 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 { + 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, + ) -> 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]