[sgl-router] refactor - cache-aware policy (#40366)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

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<CacheSource>` 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),
+Follow-up order: concrete admission (#40271), then bucket SLO ordering
-then bucket SLO ordering, remaining policies, and production configuration.
+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.

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<u64>,
+    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<Bucket>) -> Self {
+    /// Fails if any bucket installs a policy on a stage it cannot serve.
-        Self { buckets }
+    pub fn new(buckets: Vec<Bucket>) -> Result<Self, PickError> {
+        for bucket in &buckets {
+            bucket.validate()?;
+        }
+        Ok(Self { buckets })
     }
 
     /// Return all length-compatible buckets, ordered by input capacity, rank, and ID.

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<Arc<ExternalPrefixSignal>>;
+type Lookup = Arc<OnceCell<Signal>>;
+
+/// Local radix tree or remote indexer. Groups sharing an index namespace share
+/// one `Arc<CacheSource>`.
+pub enum CacheSource {
+    Local(RadixTreePrefixProvider),
+    Remote {
+        index: Arc<dyn PrefixIndex>,
+        block_size: Arc<BlockSizeOracle>,
+    },
+}
+
+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<Signal, PickError> {
+        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<Vec<(Arc<CacheSource>, 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<CacheSource>) -> 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<Worker>,
+    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<CacheSource>,
+    engine_load: Arc<EngineReportedLoadTable>,
+    config: AffinityConfig,
+    pub admission: Arc<dyn EngineAdmission>,
+    /// Runs on a miss; this policy checks admission on its pick.
+    pub fallback: Arc<dyn Policy>,
+}
+
+impl CacheAwarePolicy {
+    pub fn new(
+        source: Arc<CacheSource>,
+        engine_load: Arc<EngineReportedLoadTable>,
+        config: AffinityConfig,
+    ) -> Result<Self, PickError> {
+        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<Worker>],
+        request: &PickRequest<'_>,
+        signal: Option<&ExternalPrefixSignal>,
+        load: &EngineReportedLoadSnapshot,
+    ) -> Vec<Candidate<'e>> {
+        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<Option<Rejection>, 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<Rejection>,
+    ) -> Result<Vec<Candidate<'e>>, 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<Ordering> {
+        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<Worker>],
+        request: &PickRequest<'_>,
+        load: &EngineReportedLoadSnapshot,
+    ) -> Result<Option<Pick>, PickError> {
+        let limit = self.config.worker_queue_limit;
+        let (mut evaluated, gated): (Vec<Candidate<'_>>, 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<Worker>],
+        request: &'a PickRequest<'a>,
+    ) -> BoxFuture<'a, Result<Pick, PickError>> {
+        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())
+    }
+}

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<u64>,
+    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<Pick, PickError>>;
 
+    /// 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

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,

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;

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),

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<Worker> {
+    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<BlockSizeOracle> {
+    let oracle = BlockSizeOracle::new();
+    oracle.try_set(1).unwrap();
+    oracle
+}
+
+fn local(entries: &[(&Arc<Worker>, usize)]) -> Arc<CacheSource> {
+    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<Vec<(String, Option<u64>)>>,
+}
+
+impl Reject {
+    fn new(id: &'static str) -> Arc<Self> {
+        Arc::new(Self {
+            id,
+            calls: Mutex::new(Vec::new()),
+        })
+    }
+}
+
+impl EngineAdmission for Reject {
+    fn check(
+        &self,
+        engine: &Worker,
+        _: &PickRequest<'_>,
+        load: Option<&EngineReportedWorkerLoad>,
+    ) -> Result<Decision, PickError> {
+        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<PrefixOutcome, PrefixIndexError>,
+    calls: AtomicUsize,
+    hashes: Mutex<Vec<Vec<i64>>>,
+}
+
+#[tonic::async_trait]
+impl PrefixIndex for Index {
+    async fn match_prefix(&self, hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
+        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<PrefixOutcome, PrefixIndexError>,
+    oracle: Arc<BlockSizeOracle>,
+) -> (Arc<CacheSource>, Arc<Index>) {
+    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<dyn Policy>, decode: Arc<dyn Policy>| {
+        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());
+}

experimental/sgl-router/tests/proxy/chat_routing/reorg.rs

@@ -124,9 +124,12 @@ fn context(workers: &[(&str, Stage, &MockWorker)], buckets: Vec<Bucket>) -> Arc<
         Arc::new(PolicyRegistry::default()),
     );
     ctx.chat_routing = ChatRouting::Reorg(
-        [(ModelId("tiny".into()), BucketResolver::new(buckets))]
+        [(
-            .into_iter()
+            ModelId("tiny".into()),
-            .collect(),
+            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());
+}