From fcb080bd4007d88c76fd901a07248066b460684a Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Sun, 20 Sep 2026 22:53:57 -0700 Subject: [PATCH] [sgl-router] refactor - cache-aware policy (#40366) Co-authored-by: Claude Fable 5.1 --- experimental/sgl-router/POLICY_DESIGN.md | 20 +- experimental/sgl-router/src/buckets_reorg.rs | 29 +- .../src/policies_reorg/cache_aware.rs | 462 ++++++++++++++ .../sgl-router/src/policies_reorg/mod.rs | 8 + .../src/server/routes/chat/reorg.rs | 2 + .../sgl-router/tests/component/main.rs | 1 + .../tests/component/policies_reorg.rs | 13 +- .../component/policies_reorg_cache_aware.rs | 563 ++++++++++++++++++ .../tests/proxy/chat_routing/reorg.rs | 84 ++- 9 files changed, 1168 insertions(+), 14 deletions(-) create mode 100644 experimental/sgl-router/src/policies_reorg/cache_aware.rs create mode 100644 experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs diff --git a/experimental/sgl-router/POLICY_DESIGN.md b/experimental/sgl-router/POLICY_DESIGN.md index de1f79ad6..d0077d2cb 100644 --- a/experimental/sgl-router/POLICY_DESIGN.md +++ b/experimental/sgl-router/POLICY_DESIGN.md @@ -172,7 +172,7 @@ bucket, both groups share this one request-length decision. Policy fallback on a cache/affinity miss stays within that group's candidates. There is no second pass with relaxed admission and no post-policy substitution. -SLO ordering, cache lookup, global session modes, and sticky policies are follow-ups. Their +SLO ordering, global session modes, and sticky policies are follow-ups. Their integration must preserve bucket-first selection and the same-bucket PD rule. Cross-bucket affinity probing is not part of this interface. Session/routing keys still pass through `PickRequest` for policies operating inside the selected @@ -568,7 +568,17 @@ Implemented here: dispatches only after one complete selection. Exhaustion retains admission reasons. - `AppContext::chat_routing` configures legacy versus reorg routing on the same endpoint and carries the reorg model-resolver map. - +- `CacheAwarePolicy` reads local radix-tree or remote indexer prefixes, intersects + exact worker URLs with the current group, applies hit thresholds and candidate + bounds, and preserves the soft queue gate, saturation pin and pressure guard. +- `PrefixMemo` shares lookup results (including misses and unavailable backends) + across bucket attempts for one prepared request. Entries are keyed by the shared + `Arc` so different index namespaces remain independent. Each pick + reruns its own candidate filtering and admission after obtaining a fresh snapshot. +- Cache selection checks bounded candidates explicitly; hard rejection cannot + become a cold fallback or bypass admission through saturation pinning. A miss + defaults to power-of-two within the group's soft queue tier, then the cache + policy checks its fallback winner. Cache policies require plain/prefill groups. - `SessionAwarePolicy` reuses admitted model/bucket/role-scoped bindings from a shared `AffinityStore`, falling back to power-of-two for new or keyless sessions. Assignments follow admission; concurrent binding winners are rechecked. @@ -576,8 +586,8 @@ Implemented here: The caller owns expiry and sweeper lifecycle. A binding may remain after a later PD group fails, because it records placement rather than dispatch. -Follow-up order: cache-aware selection (#40366), concrete admission (#40271), -then bucket SLO ordering, remaining policies, and production configuration. +Follow-up order: concrete admission (#40271), then bucket SLO ordering +in a separate PR, followed by remaining policies and production configuration. Not yet implemented in the reorg path: @@ -586,7 +596,7 @@ Not yet implemented in the reorg path: - CLI/configuration parsing, validation, and model-specific construction. The YAML above is illustrative; reorg resolvers are installed in code. - Global session modes and sticky routing-key affinity. -- Prefix memoization and cache-aware selection. +- Power-of-k cache-miss fallback configuration and cache decision metrics. - Shared load interpretation, dispatch correction, and policy-specific dispatch-timestamp requirements. - PD compatibility filtering, retry integration, and legacy-route switchover. diff --git a/experimental/sgl-router/src/buckets_reorg.rs b/experimental/sgl-router/src/buckets_reorg.rs index 2bf39509c..ed7caa9dd 100644 --- a/experimental/sgl-router/src/buckets_reorg.rs +++ b/experimental/sgl-router/src/buckets_reorg.rs @@ -103,6 +103,7 @@ pub struct BucketRequest<'a> { pub model: &'a ModelId, pub input_tokens: u64, pub expected_peak_tokens: Option, + pub prefix: Option<&'a crate::policies_reorg::cache_aware::PrefixMemo>, pub token_ids: Option<&'a [u32]>, pub session_key: Option<&'a str>, pub routing_key: Option<&'a str>, @@ -139,6 +140,25 @@ impl Bucket { } } + /// Reject a policy installed on a stage it cannot serve before any request reaches it. + pub fn validate(&self) -> Result<(), PickError> { + let groups: &[(Stage, &EngineGroup)] = match &self.groups { + BucketGroups::Plain(group) => &[(Stage::Plain, group)], + BucketGroups::Pd { prefill, decode } => { + &[(Stage::Prefill, prefill), (Stage::Decode, decode)] + } + }; + for (stage, group) in groups { + if !group.policy.supports(*stage) { + return Err(PickError::InvalidConfiguration(format!( + "bucket {} installs a policy that cannot serve the {stage:?} stage", + self.id + ))); + } + } + Ok(()) + } + /// Select this bucket's plain engine or complete P/D pair, without dispatching. /// A failed group reports its stage; the caller may then try another bucket. pub async fn pick_engines( @@ -179,6 +199,7 @@ impl Bucket { bucket: &self.id, input_tokens: request.input_tokens, expected_peak_tokens: request.expected_peak_tokens, + prefix: request.prefix, token_ids: request.token_ids, session_key: request.session_key, routing_key: request.routing_key, @@ -211,8 +232,12 @@ pub struct BucketResolver { } impl BucketResolver { - pub fn new(buckets: Vec) -> Self { - Self { buckets } + /// Fails if any bucket installs a policy on a stage it cannot serve. + pub fn new(buckets: Vec) -> Result { + for bucket in &buckets { + bucket.validate()?; + } + Ok(Self { buckets }) } /// Return all length-compatible buckets, ordered by input capacity, rank, and ID. diff --git a/experimental/sgl-router/src/policies_reorg/cache_aware.rs b/experimental/sgl-router/src/policies_reorg/cache_aware.rs new file mode 100644 index 000000000..199317be1 --- /dev/null +++ b/experimental/sgl-router/src/policies_reorg/cache_aware.rs @@ -0,0 +1,462 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Cache-aware selection within one engine group, following the legacy +//! `policies::cache_aware` proposal and `resolve_cache_candidates` rules. +//! Prefix I/O is memoized per request; candidate bounding, the queue gate and +//! admission run per pick against a fresh load snapshot. + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use futures::future::BoxFuture; +use sgl_kv_indexer::{PrefixIndex, PrefixIndexError, PrefixOutcome}; +use tokio::sync::OnceCell; + +use crate::config::AffinityConfig; +use crate::policies::admission::{fleet_is_all_queued, queue_gate_admits, FreshLoadLookup}; +use crate::policies::prefix_provider::RadixTreePrefixProvider; +use crate::policies::ExternalPrefixSignal; +use crate::state::kv_events::{compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle}; +use crate::state::load_monitor::engine_reported_load::{ + EngineReportedLoadSnapshot, EngineReportedLoadTable, EngineReportedSchedulingLoad, +}; +use crate::workers::Worker; + +use super::admission::{AllowAll, Decision, EngineAdmission}; +use super::power_of_two::PowerOfTwoPolicy; +use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; + +type Signal = Option>; +type Lookup = Arc>; + +/// Local radix tree or remote indexer. Groups sharing an index namespace share +/// one `Arc`. +pub enum CacheSource { + Local(RadixTreePrefixProvider), + Remote { + index: Arc, + block_size: Arc, + }, +} + +impl fmt::Debug for CacheSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Local(_) => "CacheSource::Local", + Self::Remote { .. } => "CacheSource::Remote", + }) + } +} + +impl CacheSource { + async fn lookup(&self, tokens: Option<&[u32]>) -> Result { + let Some(tokens) = tokens else { + return Ok(None); + }; + let (index, oracle) = match self { + Self::Local(provider) => { + return Ok(provider.match_request_tokens(tokens).map(Arc::new)) + } + Self::Remote { index, block_size } => (index, block_size), + }; + let Some(block_size) = oracle.get() else { + return Ok(None); + }; + let hashes = if oracle.is_bigram() { + compute_block_hashes_bigram(tokens, block_size as usize) + } else { + compute_block_hashes(tokens, block_size as usize) + }; + let query_blocks = hashes.len(); + if query_blocks == 0 { + return Ok(None); + } + match index.match_prefix(hashes).await { + Ok(outcome) => Ok(Some(Arc::new(ExternalPrefixSignal { + outcome, + query_blocks, + }))), + Err(PrefixIndexError::Rejected(code)) => Err(PickError::InvalidSignal(format!( + "KV Indexer rejected the query: {code}" + ))), + Err(error) => { + tracing::warn!(%error, "KV Indexer unavailable; using cache policy fallback"); + Ok(None) + } + } + } +} + +/// One per prepared request. Keyed by source identity so distinct index +/// namespaces never share an answer. +#[derive(Default)] +pub struct PrefixMemo { + cells: Mutex, Lookup)>>, +} + +impl fmt::Debug for PrefixMemo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("PrefixMemo") + } +} + +impl PrefixMemo { + fn cell(&self, source: &Arc) -> Lookup { + let mut cells = self.cells.lock().unwrap_or_else(|e| e.into_inner()); + match cells.iter().find(|(s, _)| Arc::ptr_eq(s, source)) { + Some((_, cell)) => Arc::clone(cell), + None => { + let cell = Arc::default(); + cells.push((Arc::clone(source), Arc::clone(&cell))); + cell + } + } + } +} + +#[derive(Clone, Copy)] +struct Candidate<'a> { + engine: &'a Arc, + uncached_tokens: u64, +} + +/// Less uncached work first, then lower prefill pressure, then worker id. +fn rank(loads: &FreshLoadLookup<'_>, left: &Candidate<'_>, right: &Candidate<'_>) -> Ordering { + left.uncached_tokens + .cmp(&right.uncached_tokens) + .then_with(|| loads.compare_prefill_pressure(left.engine, right.engine)) + .then_with(|| left.engine.id.0.cmp(&right.engine.id.0)) +} + +#[derive(Debug)] +pub struct CacheAwarePolicy { + source: Arc, + engine_load: Arc, + config: AffinityConfig, + pub admission: Arc, + /// Runs on a miss; this policy checks admission on its pick. + pub fallback: Arc, +} + +impl CacheAwarePolicy { + pub fn new( + source: Arc, + engine_load: Arc, + config: AffinityConfig, + ) -> Result { + let unit = |ratio: f64| ratio.is_finite() && (0.0..=1.0).contains(&ratio); + let nonnegative = |ms: f64| ms.is_finite() && ms >= 0.0; + let floor_fits = |floor| { + config + .worker_queue_limit + .is_some_and(|limit| floor <= limit) + }; + let valid = (1..=config.cache_candidate_max_workers) + .contains(&config.cache_candidate_min_workers) + && unit(config.cache_candidate_ratio) + && config.cache_affinity_min_match_ratio.is_none_or(unit) + && config.pressure_rel_threshold.is_finite() + && config.pressure_rel_threshold > 1.0 + && config.pressure_abs_threshold_ms.is_none_or(nonnegative) + && config.saturation_queue_floor.is_none_or(floor_fits); + if !valid { + return Err(PickError::InvalidConfiguration( + "invalid cache candidate bounds, thresholds or saturation floor".into(), + )); + } + Ok(Self { + source, + fallback: Arc::new(PowerOfTwoPolicy::new(Arc::clone(&engine_load))), + engine_load, + config, + admission: Arc::new(AllowAll), + }) + } + + /// Prefix holders in `engines` (matched by exact URL) that pass the hit + /// thresholds, ranked and bounded to the configured candidate count. + fn candidates<'e>( + &self, + engines: &'e [Arc], + request: &PickRequest<'_>, + signal: Option<&ExternalPrefixSignal>, + load: &EngineReportedLoadSnapshot, + ) -> Vec> { + let Some(ExternalPrefixSignal { + outcome: PrefixOutcome::Matched { matches, .. }, + query_blocks, + }) = signal.filter(|signal| signal.query_blocks > 0) + else { + return Vec::new(); + }; + let query_blocks = *query_blocks as u64; + let mut depths = HashMap::<&str, u64>::new(); + for entry in matches { + let depth = depths.entry(entry.address.as_str()).or_default(); + *depth = (*depth).max(u64::from(entry.matched_prefix_blocks)); + } + let config = &self.config; + let input = request.input_tokens; + let mut candidates: Vec<_> = engines + .iter() + .filter_map(|engine| { + let blocks = depths.get(engine.url.as_str())?.min(&query_blocks); + let matched = input.saturating_mul(*blocks) / query_blocks; + let ratio = matched as f64 / input.max(1) as f64; + let hit = *blocks > 0 + && config + .cache_affinity_min_matched_tokens + .is_none_or(|min| matched >= min) + && config + .cache_affinity_min_match_ratio + .is_none_or(|min| ratio >= min); + let uncached_tokens = input - matched; + hit.then_some(Candidate { + engine, + uncached_tokens, + }) + }) + .collect(); + let proportional = (config.cache_candidate_ratio * engines.len() as f64).ceil() as usize; + let limit = engines + .len() + .min(config.cache_candidate_max_workers) + .min(config.cache_candidate_min_workers.max(proportional)); + let loads = FreshLoadLookup::new(Some(load), candidates.iter().map(|c| c.engine)); + candidates.sort_by(|left, right| rank(&loads, left, right)); + candidates.truncate(limit); + candidates + } + + /// `None` when admitted. + fn check( + &self, + engine: &Worker, + request: &PickRequest<'_>, + load: &EngineReportedLoadSnapshot, + ) -> Result, PickError> { + let load = load.fresh_load_for_url(&engine.url); + Ok(match self.admission.check(engine, request, load)? { + Decision::Allow => None, + Decision::Reject(reason) => Some(Rejection { + engine: engine.id.clone(), + reason, + }), + }) + } + + fn admit<'e>( + &self, + candidates: &[Candidate<'e>], + request: &PickRequest<'_>, + load: &EngineReportedLoadSnapshot, + rejections: &mut Vec, + ) -> Result>, PickError> { + let mut admitted = Vec::new(); + for &candidate in candidates { + match self.check(candidate.engine, request, load)? { + None => admitted.push(candidate), + Some(rejection) => rejections.push(rejection), + } + } + Ok(admitted) + } + + fn more_pressured( + &self, + left: &EngineReportedSchedulingLoad, + right: &EngineReportedSchedulingLoad, + ) -> bool { + let config = &self.config; + let queue_ms = |load: &EngineReportedSchedulingLoad| load.estimated_prefill_queue_ms; + match ( + config.pressure_abs_threshold_ms, + queue_ms(left), + queue_ms(right), + ) { + (Some(abs), Some(left), Some(right)) => { + left - right > abs && left > right * config.pressure_rel_threshold + } + _ => { + let (left, right) = ( + left.num_waiting_uncached_tokens, + right.num_waiting_uncached_tokens, + ); + left.saturating_sub(right) > config.pressure_abs_threshold_tokens + && left as f64 > right as f64 * config.pressure_rel_threshold + } + } + } + + /// Pressure guard between near-tied candidates; `None` defers to `rank`. + fn guard( + &self, + left: &Worker, + right: &Worker, + load: &EngineReportedLoadSnapshot, + ) -> Option { + let left = load.fresh_native_cache_load_for_url(&left.url)?; + let right = load.fresh_native_cache_load_for_url(&right.url)?; + if self.more_pressured(left, right) { + Some(Ordering::Greater) + } else if self.more_pressured(right, left) { + Some(Ordering::Less) + } else { + None + } + } + + /// Soft queue gate, saturation rules and hard admission over the bounded + /// candidates. `Ok(None)` is a miss; a rejection never becomes a cold fallback. + fn resolve( + &self, + candidates: &[Candidate<'_>], + engines: &[Arc], + request: &PickRequest<'_>, + load: &EngineReportedLoadSnapshot, + ) -> Result, PickError> { + let limit = self.config.worker_queue_limit; + let (mut evaluated, gated): (Vec>, Vec<_>) = candidates + .iter() + .partition(|c| queue_gate_admits(load, c.engine, limit)); + // Diverting off an all-queued group buys nothing, so keep the prefix. + // A configured floor replaces this tier with the pressure-ranked pin below. + let saturated = evaluated.is_empty() + && !gated.is_empty() + && fleet_is_all_queued(load, engines, limit) + && self.config.saturation_queue_floor.is_none(); + if saturated { + evaluated.extend(&gated); + } + let mut rejections = Vec::new(); + let admitted = self.admit(&evaluated, request, load, &mut rejections)?; + if let Some(&least) = admitted.iter().min_by_key(|c| c.uncached_tokens) { + let loads = FreshLoadLookup::new(Some(load), evaluated.iter().map(|c| c.engine)); + let guarded = self.config.pressure_guard + && evaluated.iter().all(|c| { + load.fresh_native_cache_load_for_url(&c.engine.url) + .is_some() + }); + let ceiling = least + .uncached_tokens + .saturating_add(self.config.cache_switch_margin_tokens); + let winner = admitted + .iter() + .filter(|c| c.uncached_tokens <= ceiling) + .fold(least, |winner, &candidate| { + // The guard applies only to pairs within the margin of + // each other, not merely of the work floor. + let near_tie = winner.uncached_tokens.abs_diff(candidate.uncached_tokens) + <= self.config.cache_switch_margin_tokens; + let guard = (guarded && near_tie) + .then(|| self.guard(winner.engine, candidate.engine, load)) + .flatten(); + match guard.unwrap_or_else(|| rank(&loads, &winner, &candidate)) { + Ordering::Greater => candidate, + _ => winner, + } + }); + return Ok(Some(Pick { + engine: Arc::clone(winner.engine), + reason: if saturated { + "saturation_pin" + } else { + "cache_candidate" + }, + })); + } + // Saturation pin: with no engine below the floor the request waits + // anywhere, so wait at the least-pressured admitted prefix owner. + let pinned = self.config.saturation_queue_floor.is_some_and(|floor| { + !load.any_fresh_queue_below(engines.iter().map(|e| e.url.as_str()), floor) + }); + if pinned { + let loads = FreshLoadLookup::new(Some(load), gated.iter().map(|c| c.engine)); + let owner = self + .admit(&gated, request, load, &mut rejections)? + .into_iter() + .min_by(|left, right| { + loads + .compare_prefill_pressure(left.engine, right.engine) + .then_with(|| left.engine.id.0.cmp(&right.engine.id.0)) + }); + if let Some(owner) = owner { + return Ok(Some(Pick { + engine: Arc::clone(owner.engine), + reason: "saturation_pin", + })); + } + } + if rejections.is_empty() { + Ok(None) + } else { + Err(PickError::NoAdmissibleEngine(rejections)) + } + } +} + +impl Policy for CacheAwarePolicy { + fn supports(&self, stage: Stage) -> bool { + stage != Stage::Decode + } + + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + if engines.is_empty() { + return Err(PickError::NoCandidates); + } + if request.stage == Stage::Decode { + return Err(PickError::InvalidConfiguration( + "cache-aware selection requires a plain or prefill group".into(), + )); + } + let lookup = || self.source.lookup(request.token_ids); + let signal = match request.prefix { + Some(memo) => memo + .cell(&self.source) + .get_or_try_init(lookup) + .await? + .clone(), + None => lookup().await?, + }; + // Capture load after remote I/O; selection and admission share it. + let load = self.engine_load.capture_snapshot(Instant::now()); + let candidates = self.candidates(engines, request, signal.as_deref(), &load); + if let Some(pick) = self.resolve(&candidates, engines, request, &load)? { + return Ok(pick); + } + // Miss: fall back within the unqueued tier when one exists. + let unqueued: Vec<_> = engines + .iter() + .filter(|e| queue_gate_admits(&load, e, self.config.worker_queue_limit)) + .cloned() + .collect(); + let pool = if unqueued.is_empty() { + engines + } else { + &unqueued + }; + let mut pick = self.pick_fallback(pool, request).await?; + if !pool.iter().any(|e| Arc::ptr_eq(e, &pick.engine)) { + return Err(PickError::OutsideCandidates(pick.engine.id.clone())); + } + if let Some(rejection) = self.check(&pick.engine, request, &load)? { + return Err(PickError::AdmissionRejected(rejection)); + } + pick.reason = "no_cache_candidate"; + Ok(pick) + }) + } + + fn fallback(&self) -> Option<&dyn Policy> { + Some(self.fallback.as_ref()) + } +} diff --git a/experimental/sgl-router/src/policies_reorg/mod.rs b/experimental/sgl-router/src/policies_reorg/mod.rs index 1fa209bdf..eb4af4b31 100644 --- a/experimental/sgl-router/src/policies_reorg/mod.rs +++ b/experimental/sgl-router/src/policies_reorg/mod.rs @@ -5,6 +5,7 @@ //! this interface through AppContext; `policies` remains the default. pub mod admission; +pub mod cache_aware; pub mod power_of_two; pub mod session_aware; @@ -26,6 +27,7 @@ pub struct PickRequest<'a> { pub bucket: &'a str, pub input_tokens: u64, pub expected_peak_tokens: Option, + pub prefix: Option<&'a cache_aware::PrefixMemo>, pub token_ids: Option<&'a [u32]>, pub session_key: Option<&'a str>, pub routing_key: Option<&'a str>, @@ -39,6 +41,7 @@ impl<'a> PickRequest<'a> { bucket: "", input_tokens, expected_peak_tokens: None, + prefix: None, token_ids: None, session_key: None, routing_key: None, @@ -87,6 +90,11 @@ pub trait Policy: Send + Sync + Debug { request: &'a PickRequest<'a>, ) -> BoxFuture<'a, Result>; + /// Whether this policy may serve `stage`; checked when a resolver is built. + fn supports(&self, _stage: Stage) -> bool { + true + } + /// Runs on a miss within the same candidates; never on an admission rejection. fn fallback(&self) -> Option<&dyn Policy> { None diff --git a/experimental/sgl-router/src/server/routes/chat/reorg.rs b/experimental/sgl-router/src/server/routes/chat/reorg.rs index eb872d7d9..0822ebf07 100644 --- a/experimental/sgl-router/src/server/routes/chat/reorg.rs +++ b/experimental/sgl-router/src/server/routes/chat/reorg.rs @@ -55,7 +55,9 @@ pub(super) async fn chat_completions( None, )); } + let prefix = crate::policies_reorg::cache_aware::PrefixMemo::default(); let bucket_request = BucketRequest { + prefix: Some(&prefix), model: &request.model, input_tokens, expected_peak_tokens, diff --git a/experimental/sgl-router/tests/component/main.rs b/experimental/sgl-router/tests/component/main.rs index 24d4dc09d..d6acf13a7 100644 --- a/experimental/sgl-router/tests/component/main.rs +++ b/experimental/sgl-router/tests/component/main.rs @@ -11,6 +11,7 @@ mod discovery; mod health; mod policies; mod policies_reorg; +mod policies_reorg_cache_aware; mod policies_reorg_load; mod policies_reorg_power_of_two; mod policies_reorg_session_aware; diff --git a/experimental/sgl-router/tests/component/policies_reorg.rs b/experimental/sgl-router/tests/component/policies_reorg.rs index 65ac82cf9..fa49c61ec 100644 --- a/experimental/sgl-router/tests/component/policies_reorg.rs +++ b/experimental/sgl-router/tests/component/policies_reorg.rs @@ -184,7 +184,8 @@ fn resolve_orders_all_length_fits_by_capacity_rank_and_id() { later, a, min, - ]); + ]) + .unwrap(); assert_eq!( resolver .resolve(10, None) @@ -207,7 +208,7 @@ fn context_capacity_checks_peak_when_known_and_input_otherwise() { short.max_context_tokens = Some(20); let mut long = bucket("long", None, policy); long.max_context_tokens = Some(30); - let resolver = BucketResolver::new(vec![long, short]); + let resolver = BucketResolver::new(vec![long, short]).unwrap(); assert_eq!(resolver.resolve(10, None).unwrap()[0].id, "short"); assert_eq!(resolver.resolve(10, Some(20)).unwrap()[0].id, "short"); assert_eq!(resolver.resolve(10, Some(21)).unwrap()[0].id, "long"); @@ -237,9 +238,11 @@ async fn selected_pd_bucket_owns_both_memberships_and_policies() { prefill: group(&["p2", "d", "a"], prefill_policy.clone()), decode: group(&["d2", "p", "other"], decode_policy.clone()), }, - )]); + )]) + .unwrap(); let bucket = resolver.resolve(10, Some(20)).unwrap()[0]; let request = BucketRequest { + prefix: None, model: &model, input_tokens: 10, expected_peak_tokens: Some(20), @@ -267,7 +270,8 @@ async fn resolver_includes_empty_groups_without_invoking_policies() { min: None, max: Some(10), }; - let resolver = BucketResolver::new(vec![empty, bucket("available", Some(20), policy.clone())]); + let resolver = + BucketResolver::new(vec![empty, bucket("available", Some(20), policy.clone())]).unwrap(); let buckets = resolver.resolve(10, None).unwrap(); assert_eq!( buckets @@ -406,6 +410,7 @@ async fn bucket_scopes_plain_pick_and_preserves_request_facts() { BucketGroups::Plain(group(&["b"], Arc::new(InspectRequest))), ); let request = BucketRequest { + prefix: None, model: &model, input_tokens: 2, expected_peak_tokens: Some(12), diff --git a/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs b/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs new file mode 100644 index 000000000..a002100dd --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs @@ -0,0 +1,563 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use sgl_kv_indexer::{PrefixIndex, PrefixIndexError, PrefixMatch, PrefixOutcome}; +use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup}; +use sgl_router::config::AffinityConfig; +use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; +use sgl_router::policies::prefix_provider::RadixTreePrefixProvider; +use sgl_router::policies_reorg::admission::{Decision, EngineAdmission}; +use sgl_router::policies_reorg::cache_aware::{CacheAwarePolicy, CacheSource, PrefixMemo}; +use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; +use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; +use sgl_router::state::kv_events::{ + compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree, KvWorkerId, +}; +use sgl_router::state::load_monitor::engine_reported_load::{ + EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, NativeCacheRankLoad, +}; +use sgl_router::workers::Worker; + +const TOKENS: [u32; 8] = [1, 2, 3, 4, 5, 6, 7, 8]; + +fn engine(id: &str, active: usize) -> Arc { + let engine = Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}"), + mode: Stage::Plain, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + })); + engine.active_requests.store(active, Ordering::Relaxed); + engine +} + +fn config() -> AffinityConfig { + AffinityConfig { + cache_affinity_min_matched_tokens: Some(1), + cache_switch_margin_tokens: 0, + ..Default::default() + } +} + +fn oracle() -> Arc { + let oracle = BlockSizeOracle::new(); + oracle.try_set(1).unwrap(); + oracle +} + +fn local(entries: &[(&Arc, usize)]) -> Arc { + let tree = Arc::new(HashTree::new()); + let hashes = compute_block_hashes(&TOKENS, 1); + for (worker, depth) in entries { + tree.insert( + &KvWorkerId::new(worker.url.clone(), 0), + None, + &hashes[..*depth], + ); + } + Arc::new(CacheSource::Local(RadixTreePrefixProvider::new( + tree, + oracle(), + ))) +} + +fn request(model: &ModelId) -> PickRequest<'_> { + PickRequest { + token_ids: Some(&TOKENS), + ..PickRequest::new(model, Stage::Plain, 8) + } +} + +fn report( + table: &EngineReportedLoadTable, + engine: &Worker, + waiting: u64, + pending: u64, + at: Instant, +) { + table.set( + &engine.url, + 0, + LoadStat { + num_running_reqs: 1, + num_waiting_reqs: waiting, + num_tokens: 10, + max_total_num_tokens: 100, + native_cache: Some(NativeCacheRankLoad { + num_waiting_uncached_tokens: pending, + num_total_tokens: 10, + max_running_requests: 100, + total_prefill_uncached_tokens: 0, + total_prefill_busy_us: 0, + }), + }, + at, + ); +} + +#[derive(Debug)] +struct Reject { + id: &'static str, + calls: Mutex)>>, +} + +impl Reject { + fn new(id: &'static str) -> Arc { + Arc::new(Self { + id, + calls: Mutex::new(Vec::new()), + }) + } +} + +impl EngineAdmission for Reject { + fn check( + &self, + engine: &Worker, + _: &PickRequest<'_>, + load: Option<&EngineReportedWorkerLoad>, + ) -> Result { + self.calls + .lock() + .unwrap() + .push((engine.id.0.clone(), load.map(|load| load.num_waiting_reqs))); + Ok(if engine.id.0 == self.id { + Decision::Reject("full".into()) + } else { + Decision::Allow + }) + } +} + +struct Index { + result: Result, + calls: AtomicUsize, + hashes: Mutex>>, +} + +#[tonic::async_trait] +impl PrefixIndex for Index { + async fn match_prefix(&self, hashes: Vec) -> Result { + self.calls.fetch_add(1, Ordering::Relaxed); + self.hashes.lock().unwrap().push(hashes); + tokio::task::yield_now().await; + self.result.clone() + } +} + +fn remote( + result: Result, + oracle: Arc, +) -> (Arc, Arc) { + let index = Arc::new(Index { + result, + calls: AtomicUsize::new(0), + hashes: Mutex::new(Vec::new()), + }); + ( + Arc::new(CacheSource::Remote { + index: index.clone(), + block_size: oracle, + }), + index, + ) +} + +fn matches(entries: &[(&str, u32)]) -> PrefixOutcome { + PrefixOutcome::Matched { + matches: entries + .iter() + .map(|(url, depth)| PrefixMatch { + address: (*url).into(), + worker_id: "not-a-routing-identity".into(), + matched_prefix_blocks: *depth, + }) + .collect(), + best_prefix_blocks: entries.iter().map(|(_, depth)| *depth).max().unwrap_or(0), + } +} + +#[tokio::test] +async fn local_prefix_wins_within_group_and_threshold_misses_use_load() { + let engines = [ + engine("deep", 9), + engine("shallow", 0), + engine("outside", 0), + ]; + let source = local(&[(&engines[0], 7), (&engines[1], 4), (&engines[2], 8)]); + let model = ModelId("m".into()); + for (minimum, ratio, expected) in [(Some(1), None, 0), (Some(8), None, 1), (None, Some(0.9), 1)] + { + let policy = CacheAwarePolicy::new( + source.clone(), + EngineReportedLoadTable::new(), + AffinityConfig { + cache_affinity_min_matched_tokens: minimum, + cache_affinity_min_match_ratio: ratio, + ..config() + }, + ) + .unwrap(); + for stage in [Stage::Plain, Stage::Prefill] { + let request = PickRequest { + stage, + ..request(&model) + }; + let pick = policy.pick(&engines[..2], &request).await.unwrap(); + assert!(Arc::ptr_eq(&pick.engine, &engines[expected])); + } + } +} + +#[tokio::test] +async fn remote_urls_are_exact_duplicate_depths_merge_and_block_counts_are_capped() { + let engines = [engine("a", 9), engine("b", 0)]; + let (source, _) = remote( + Ok(matches(&[ + ("http://a", 1), + ("http://a", u32::MAX), + ("http://b/", u32::MAX), + ])), + oracle(), + ); + let policy = CacheAwarePolicy::new( + source, + EngineReportedLoadTable::new(), + AffinityConfig { + cache_affinity_min_match_ratio: Some(1.0), + ..config() + }, + ) + .unwrap(); + let model = ModelId("m".into()); + assert_eq!( + policy + .pick(&engines, &request(&model)) + .await + .unwrap() + .engine + .id + .0, + "a" + ); +} + +#[tokio::test] +async fn memo_reuses_io_but_reruns_admission_and_group_selection() { + let engines = [engine("a", 0), engine("b", 0)]; + let (source, index) = remote(Ok(matches(&[("http://a", 8), ("http://b", 7)])), oracle()); + let mut policy = + CacheAwarePolicy::new(source.clone(), EngineReportedLoadTable::new(), config()).unwrap(); + let admission = Reject::new("a"); + policy.admission = admission.clone(); + let memo = PrefixMemo::default(); + let model = ModelId("m".into()); + let request = PickRequest { + prefix: Some(&memo), + ..request(&model) + }; + assert!(matches!( + policy.pick(&engines[..1], &request).await, + Err(PickError::NoAdmissibleEngine(_)) + )); + assert_eq!( + policy + .pick(&engines[1..], &request) + .await + .unwrap() + .engine + .id + .0, + "b" + ); + assert_eq!( + policy.pick(&engines, &request).await.unwrap().engine.id.0, + "b" + ); + assert_eq!(index.calls.load(Ordering::Relaxed), 1); + assert_eq!(admission.calls.lock().unwrap().len(), 4); + // A second namespace in the same request must query its own backend. + let (other, other_index) = remote(Ok(matches(&[("http://a", 8)])), oracle()); + let policy = CacheAwarePolicy::new(other, EngineReportedLoadTable::new(), config()).unwrap(); + assert_eq!( + policy.pick(&engines, &request).await.unwrap().engine.id.0, + "a" + ); + assert_eq!(other_index.calls.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn concurrent_picks_share_one_query_and_new_requests_query_again() { + let engines = [engine("a", 0)]; + let (source, index) = remote(Ok(PrefixOutcome::Empty), oracle()); + let policy = CacheAwarePolicy::new(source, EngineReportedLoadTable::new(), config()).unwrap(); + let model = ModelId("m".into()); + let memo = PrefixMemo::default(); + let request = PickRequest { + prefix: Some(&memo), + ..request(&model) + }; + let (left, right) = tokio::join!( + policy.pick(&engines, &request), + policy.pick(&engines, &request) + ); + assert!(left.is_ok() && right.is_ok()); + assert_eq!(index.calls.load(Ordering::Relaxed), 1); + let next = PrefixMemo::default(); + policy + .pick( + &engines, + &PickRequest { + prefix: Some(&next), + ..request + }, + ) + .await + .unwrap(); + assert_eq!(index.calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn unavailable_index_falls_back_but_rejected_queries_fail() { + let engines = [engine("a", 9), engine("b", 0)]; + let model = ModelId("m".into()); + for error in [ + PrefixIndexError::Timeout, + PrefixIndexError::Unreachable, + PrefixIndexError::Overloaded, + PrefixIndexError::QueryTooLarge, + PrefixIndexError::Rejected(sgl_kv_indexer::RpcCode::InvalidArgument), + ] { + let rejected = matches!(error, PrefixIndexError::Rejected(_)); + let (source, index) = remote(Err(error), oracle()); + let policy = + CacheAwarePolicy::new(source, EngineReportedLoadTable::new(), config()).unwrap(); + let memo = PrefixMemo::default(); + let request = PickRequest { + prefix: Some(&memo), + ..request(&model) + }; + let result = policy.pick(&engines, &request).await; + if rejected { + assert!(matches!(result, Err(PickError::InvalidSignal(_)))); + } else { + assert_eq!(result.unwrap().engine.id.0, "b"); + policy.pick(&engines, &request).await.unwrap(); + } + assert_eq!(index.calls.load(Ordering::Relaxed), 1); + } +} + +#[tokio::test] +async fn missing_tokens_or_block_size_skip_io_and_bigram_hashes_match_workers() { + let engines = [engine("a", 0)]; + let model = ModelId("m".into()); + let (source, index) = remote(Ok(PrefixOutcome::Empty), BlockSizeOracle::new()); + let policy = CacheAwarePolicy::new(source, EngineReportedLoadTable::new(), config()).unwrap(); + policy.pick(&engines, &request(&model)).await.unwrap(); + assert_eq!(index.calls.load(Ordering::Relaxed), 0); + let oracle = oracle(); + oracle.set_bigram(true); + let (source, index) = remote(Ok(PrefixOutcome::Empty), oracle); + let policy = CacheAwarePolicy::new(source, EngineReportedLoadTable::new(), config()).unwrap(); + policy + .pick(&engines, &PickRequest::new(&model, Stage::Plain, 8)) + .await + .unwrap(); + assert_eq!(index.calls.load(Ordering::Relaxed), 0); + policy.pick(&engines, &request(&model)).await.unwrap(); + assert_eq!( + index.hashes.lock().unwrap()[0], + compute_block_hashes_bigram(&TOKENS, 1) + ); +} + +#[tokio::test] +async fn queue_diversion_and_saturation_use_only_this_group() { + let engines = [engine("owner", 0), engine("cold", 9), engine("outside", 0)]; + let model = ModelId("m".into()); + for (cold_waiting, floor, expected, reason) in [ + (0, None, "cold", "no_cache_candidate"), + (5, None, "owner", "saturation_pin"), + (3, Some(2), "owner", "saturation_pin"), + ] { + let table = EngineReportedLoadTable::new(); + report(&table, &engines[0], 5, 100, Instant::now()); + report(&table, &engines[1], cold_waiting, 1, Instant::now()); + report(&table, &engines[2], 0, 0, Instant::now()); + let policy = CacheAwarePolicy::new( + local(&[(&engines[0], 8)]), + table, + AffinityConfig { + worker_queue_limit: Some(4), + saturation_queue_floor: floor, + ..config() + }, + ) + .unwrap(); + let pick = policy.pick(&engines[..2], &request(&model)).await.unwrap(); + assert_eq!(pick.engine.id.0, expected); + assert_eq!(pick.reason, reason); + } +} + +#[tokio::test] +async fn hard_rejection_never_becomes_cold_fallback_or_saturation_bypass() { + let engines = [engine("owner", 0), engine("cold", 9)]; + let model = ModelId("m".into()); + for floor in [None, Some(2)] { + let table = EngineReportedLoadTable::new(); + for engine in &engines { + report(&table, engine, 5, 100, Instant::now()); + } + let mut policy = CacheAwarePolicy::new( + local(&[(&engines[0], 8)]), + table, + AffinityConfig { + worker_queue_limit: Some(4), + saturation_queue_floor: floor, + ..config() + }, + ) + .unwrap(); + let admission = Reject::new("owner"); + policy.admission = admission.clone(); + assert!(matches!( + policy.pick(&engines, &request(&model)).await, + Err(PickError::NoAdmissibleEngine(_)) + )); + assert_eq!( + *admission.calls.lock().unwrap(), + vec![("owner".into(), Some(5))] + ); + } + let mut policy = + CacheAwarePolicy::new(local(&[]), EngineReportedLoadTable::new(), config()).unwrap(); + policy.admission = Reject::new("owner"); + assert!(matches!( + policy.pick(&engines, &request(&model)).await, + Err(PickError::AdmissionRejected(_)) + )); +} + +#[tokio::test] +async fn guard_switches_near_ties_only_with_complete_fresh_telemetry() { + let engines = [engine("deep", 0), engine("shallow", 9)]; + let model = ModelId("m".into()); + for (margin, stale, expected) in [(0, false, "deep"), (1, false, "shallow"), (1, true, "deep")] + { + let table = EngineReportedLoadTable::new(); + report(&table, &engines[0], 5, 100, Instant::now()); + report( + &table, + &engines[1], + 1, + 1, + Instant::now() + - if stale { + Duration::from_secs(3600) + } else { + Duration::ZERO + }, + ); + let policy = CacheAwarePolicy::new( + local(&[(&engines[0], 8), (&engines[1], 7)]), + table, + AffinityConfig { + cache_switch_margin_tokens: margin, + pressure_abs_threshold_tokens: 10, + ..config() + }, + ) + .unwrap(); + assert_eq!( + policy + .pick(&engines, &request(&model)) + .await + .unwrap() + .engine + .id + .0, + expected + ); + } +} + +#[tokio::test] +async fn candidate_cap_is_applied_before_admission() { + let engines = [engine("deep", 9), engine("shallow", 0)]; + let model = ModelId("m".into()); + let mut policy = CacheAwarePolicy::new( + local(&[(&engines[0], 8), (&engines[1], 7)]), + EngineReportedLoadTable::new(), + AffinityConfig { + cache_candidate_min_workers: 1, + cache_candidate_max_workers: 1, + ..config() + }, + ) + .unwrap(); + let admission = Reject::new("deep"); + policy.admission = admission.clone(); + assert!(matches!( + policy.pick(&engines, &request(&model)).await, + Err(PickError::NoAdmissibleEngine(_)) + )); + assert_eq!(admission.calls.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn decode_group_and_invalid_configuration_are_rejected() { + let source = local(&[]); + for invalid in [ + AffinityConfig { + cache_candidate_min_workers: 0, + ..config() + }, + AffinityConfig { + cache_candidate_ratio: f64::NAN, + ..config() + }, + AffinityConfig { + saturation_queue_floor: Some(1), + ..config() + }, + ] { + assert!(matches!( + CacheAwarePolicy::new(source.clone(), EngineReportedLoadTable::new(), invalid), + Err(PickError::InvalidConfiguration(_)) + )); + } + let policy = + Arc::new(CacheAwarePolicy::new(source, EngineReportedLoadTable::new(), config()).unwrap()); + let model = ModelId("m".into()); + assert!(matches!( + policy + .pick( + &[engine("a", 0)], + &PickRequest::new(&model, Stage::Decode, 8) + ) + .await, + Err(PickError::InvalidConfiguration(_)) + )); + let pd = |prefill: Arc, decode: Arc| { + Bucket::new( + "pd", + BucketGroups::Pd { + prefill: EngineGroup::new(prefill), + decode: EngineGroup::new(decode), + }, + ) + }; + let load = Arc::new(PowerOfTwoPolicy::new(EngineReportedLoadTable::new())); + assert!(matches!( + BucketResolver::new(vec![pd(load.clone(), policy.clone())]), + Err(PickError::InvalidConfiguration(_)) + )); + assert!(BucketResolver::new(vec![pd(policy, load)]).is_ok()); +} diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs index b9040a2c6..b00ea76ad 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -124,9 +124,12 @@ fn context(workers: &[(&str, Stage, &MockWorker)], buckets: Vec) -> Arc< Arc::new(PolicyRegistry::default()), ); ctx.chat_routing = ChatRouting::Reorg( - [(ModelId("tiny".into()), BucketResolver::new(buckets))] - .into_iter() - .collect(), + [( + ModelId("tiny".into()), + BucketResolver::new(buckets).unwrap(), + )] + .into_iter() + .collect(), ); Arc::new(ctx) } @@ -523,3 +526,78 @@ async fn invalid_policy_signal_stops_bucket_iteration() { assert!(later.calls.lock().unwrap().is_empty()); assert!(worker.captured.lock().unwrap().last_body.is_none()); } + +#[tokio::test] +async fn cache_aware_routes_tokenized_prompt_and_rechecks_the_next_bucket() { + use sgl_router::config::AffinityConfig; + use sgl_router::policies::prefix_provider::RadixTreePrefixProvider; + use sgl_router::policies_reorg::cache_aware::{CacheAwarePolicy, CacheSource}; + use sgl_router::state::kv_events::{ + compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId, + }; + use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable; + + let rejected = MockWorker::start(vec![]).await; + let owner = MockWorker::start(vec![]).await; + let cold = MockWorker::start(vec![]).await; + let value = body("hello world"); + let config = config_for(""); + let tokenizers = TokenizerRegistry::load_from_config(&config).unwrap(); + let ids = + sgl_router::policies::request_tokens_for(&tokenizers, &ModelId("tiny".into()), &value) + .unwrap() + .ids; + let hashes = compute_block_hashes(&ids, 1); + let tree = Arc::new(HashTree::new()); + for worker in [&rejected, &owner] { + tree.insert(&KvWorkerId::new(worker.url.clone(), 0), None, &hashes); + } + let oracle = BlockSizeOracle::new(); + oracle.try_set(1).unwrap(); + let source = Arc::new(CacheSource::Local(RadixTreePrefixProvider::new( + tree, oracle, + ))); + let table = EngineReportedLoadTable::new(); + let config = AffinityConfig { + cache_affinity_min_matched_tokens: Some(1), + ..Default::default() + }; + let mut rejecting = + CacheAwarePolicy::new(source.clone(), table.clone(), config.clone()).unwrap(); + rejecting.admission = Arc::new(RejectAll); + let mut first = Bucket::new( + "first", + BucketGroups::Plain(EngineGroup { + worker_ids: Some([WorkerId("rejected".into())].into_iter().collect()), + policy: Arc::new(rejecting), + }), + ); + first.rank = 0; + let mut second = Bucket::new( + "second", + BucketGroups::Plain(EngineGroup { + worker_ids: Some( + [WorkerId("owner".into()), WorkerId("cold".into())] + .into_iter() + .collect(), + ), + policy: Arc::new(CacheAwarePolicy::new(source, table, config).unwrap()), + }), + ); + second.rank = 1; + let ctx = context( + &[ + ("rejected", Stage::Plain, &rejected), + ("owner", Stage::Plain, &owner), + ("cold", Stage::Plain, &cold), + ], + vec![first, second], + ); + let app = build_router(ctx); + let response = app.oneshot(request(value)).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + assert!(rejected.captured.lock().unwrap().last_body.is_none()); + assert!(cold.captured.lock().unwrap().last_body.is_none()); + assert!(owner.captured.lock().unwrap().last_body.is_some()); +}