[Router] Preserve global cache affinity with bucket routing (#38814)

Signed-off-by: Vincent Gao <vincentbo@linux.alibaba.com>

experimental/sgl-router/src/policies/buckets.rs

@@ -118,36 +118,6 @@ impl BucketSelector {
         .collect()
     }
 
-    /// Maps global Indexer candidates to prefill buckets using `E` as the workload.
-    pub fn bind_prefill_cache_candidate(
-        &self,
-        mut candidate: CacheCandidate,
-        request: BucketRequest,
-    ) -> Option<CacheCandidate> {
-        let Some(config) = &self.config else {
-            candidate.candidate_range_id = "global".to_string();
-            candidate.max_pending_prefill_tokens = None;
-            return Some(candidate);
-        };
-        let spec = config.buckets.iter().find(|spec| {
-            spec.stage == BucketStage::Prefill
-                && self.contains(spec, &candidate.worker.id.0)
-                && within(
-                    candidate.uncached_tokens,
-                    spec.min_extend_tokens,
-                    spec.max_extend_tokens,
-                )
-                && spec
-                    .max_context_tokens
-                    .is_none_or(|max_context| request.input_tokens <= max_context)
-                && (config.ttft_slo_policy != SloBucketPolicy::SloFirst
-                    || ttft_eligible(spec, request.ttft_slo_ms))
-        })?;
-        candidate.candidate_range_id = spec.id.clone();
-        candidate.max_pending_prefill_tokens = spec.max_pending_prefill_tokens;
-        Some(candidate)
-    }
-
     /// Finds the prefill bucket containing a global session primary.
     pub fn prefill_affinity_domain(
         &self,
@@ -178,6 +148,35 @@ impl BucketSelector {
             })
     }
 
+    /// Applies a cache candidate's Bucket metadata and hard limits.
+    /// Unbucketed candidates remain globally eligible, and extend ranges do not
+    /// constrain global cache affinity.
+    pub fn prepare_prefill_cache_candidate(
+        &self,
+        mut candidate: CacheCandidate,
+        request: BucketRequest,
+    ) -> Option<CacheCandidate> {
+        let Some(config) = &self.config else {
+            return Some(candidate);
+        };
+        let Some(spec) = config.buckets.iter().find(|spec| {
+            spec.stage == BucketStage::Prefill && self.contains(spec, &candidate.worker.id.0)
+        }) else {
+            return Some(candidate);
+        };
+        if !spec
+            .max_context_tokens
+            .is_none_or(|max_context| request.input_tokens <= max_context)
+            || (config.ttft_slo_policy == SloBucketPolicy::SloFirst
+                && !ttft_eligible(spec, request.ttft_slo_ms))
+        {
+            return None;
+        }
+        candidate.candidate_range_id = spec.id.clone();
+        candidate.max_pending_prefill_tokens = spec.max_pending_prefill_tokens;
+        Some(candidate)
+    }
+
     fn contains(&self, spec: &BucketSpec, worker_id: &str) -> bool {
         if spec.worker_ids.len() <= MEMBER_SCAN_MAX {
             return spec.worker_ids.iter().any(|id| id == worker_id);

experimental/sgl-router/src/policies/cache_aware.rs

@@ -75,7 +75,9 @@ impl CacheAwarePolicy {
         if let Some((selector, request)) = ctx.prefill_cache_bucket() {
             candidates = candidates
                 .into_iter()
-                .filter_map(|candidate| selector.bind_prefill_cache_candidate(candidate, request))
+                .filter_map(|candidate| {
+                    selector.prepare_prefill_cache_candidate(candidate, request)
+                })
                 .collect();
         }
 

experimental/sgl-router/src/policies/mod.rs

@@ -239,7 +239,7 @@ impl<'a> SelectionContext<'a> {
         self
     }
 
-    /// Cache-Aware uses this binding before Top-K truncation so an
+    /// Cache-Aware applies Bucket constraints before Top-K truncation so an
     /// incompatible cache holder cannot displace a lower-ranked usable one.
     pub fn with_prefill_cache_bucket(
         mut self,

experimental/sgl-router/tests/component/policies/bucket_domains.rs

@@ -107,91 +107,71 @@ fn prefill_best_effort_tries_non_slo_bucket_before_reserved_slo_capacity() {
 }
 
 #[test]
-fn cache_candidate_uses_uncached_work_range_but_full_context_and_own_ttft_profile() {
+fn prefill_domain_uses_full_input_as_extend_work_without_cache() {
     let short = worker("short", WorkerMode::Prefill);
     let long = worker("long", WorkerMode::Prefill);
     let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, &["short"]);
     short_bucket.max_extend_tokens = Some(64);
-    short_bucket.max_context_tokens = Some(4_096);
-    short_bucket.ttft_p95_at_capacity_ms = Some(80);
     let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, &["long"]);
     long_bucket.min_extend_tokens = Some(65);
-    long_bucket.max_context_tokens = Some(4_096);
-    long_bucket.ttft_p95_at_capacity_ms = Some(300);
     let selector = BucketSelector::new(Some(BucketConfig {
         buckets: vec![short_bucket, long_bucket],
-        ttft_slo_policy: SloBucketPolicy::SloFirst,
+        ttft_slo_policy: SloBucketPolicy::Disabled,
         tps_slo_policy: SloBucketPolicy::Disabled,
     }));
-    let workers = vec![Arc::clone(&short), Arc::clone(&long)];
+
-    let request = BucketRequest {
+    let domains = selector.prefill_domains(
+        &[short, long],
+        BucketRequest {
             input_tokens: 256,
             expected_peak_sequence_tokens: None,
-        ttft_slo_ms: Some(100),
+            ttft_slo_ms: None,
             tps_slo: None,
-    };
+        },
+    );
 
     assert_eq!(
-        selector
+        domains
-            .prefill_domains(&workers, request)
             .iter()
             .map(|domain| domain.id.as_str())
             .collect::<Vec<_>>(),
-        ["p-long"],
+        ["p-long"]
-        "no-hit target selection uses E=L for extend-work compatibility"
-    );
-    let short_hit = CacheCandidate {
-        worker: Arc::clone(&short),
-        matched_prefix_tokens: 224,
-        uncached_tokens: 32,
-        candidate_range_id: "global".into(),
-        max_pending_prefill_tokens: None,
-    };
-    let bound = selector
-        .bind_prefill_cache_candidate(short_hit, request)
-        .expect("E=32 fits short work range and the full L=256 fits max context");
-    assert_eq!(bound.candidate_range_id, "p-short");
-
-    let long_hit = CacheCandidate {
-        worker: Arc::clone(&long),
-        matched_prefix_tokens: 0,
-        uncached_tokens: 256,
-        candidate_range_id: "global".into(),
-        max_pending_prefill_tokens: None,
-    };
-    assert!(
-        selector
-            .bind_prefill_cache_candidate(long_hit, request)
-            .is_none(),
-        "a cache candidate whose own Hard TTFT profile misses the request SLO is rejected"
     );
 }
 
 #[test]
-fn cache_candidate_without_bucket_configuration_keeps_global_metadata() {
+fn cache_candidate_uses_bucket_metadata_without_extend_range_filtering() {
-    let p = worker("p", WorkerMode::Prefill);
+    let cached = worker("cached", WorkerMode::Prefill);
-    let selector = BucketSelector::new(None);
+    let mut cached_bucket = bucket("p-cached", BucketStage::Prefill, 10, &["cached"]);
+    cached_bucket.max_extend_tokens = Some(8);
+    cached_bucket.max_pending_prefill_tokens = Some(64);
+    let selector = BucketSelector::new(Some(BucketConfig {
+        buckets: vec![cached_bucket],
+        ttft_slo_policy: SloBucketPolicy::Disabled,
+        tps_slo_policy: SloBucketPolicy::Disabled,
+    }));
     let candidate = CacheCandidate {
-        worker: p,
+        worker: cached,
-        matched_prefix_tokens: 64,
+        matched_prefix_tokens: 128,
-        uncached_tokens: 64,
+        uncached_tokens: 128,
-        candidate_range_id: "probe".into(),
+        candidate_range_id: "global".into(),
-        max_pending_prefill_tokens: Some(1),
+        max_pending_prefill_tokens: None,
     };
-    let bound = selector
+
-        .bind_prefill_cache_candidate(
+    let prepared = selector
+        .prepare_prefill_cache_candidate(
             candidate,
             BucketRequest {
-                input_tokens: 128,
+                input_tokens: 256,
                 expected_peak_sequence_tokens: None,
                 ttft_slo_ms: None,
                 tps_slo: None,
             },
         )
-        .expect("Step 1 always has a catch-all domain");
+        .expect("extend range does not constrain global cache affinity");
 
-    assert_eq!(bound.candidate_range_id, "global");
+    assert_eq!(prepared.candidate_range_id, "p-cached");
-    assert_eq!(bound.max_pending_prefill_tokens, None);
+    assert_eq!(prepared.max_pending_prefill_tokens, Some(64));
 }
 
 #[test]
@@ -332,20 +312,6 @@ fn membership_index_preserves_exact_matching_and_fleet_order() {
     assert_eq!(ids(0), ["w1", "w3"]);
     assert_eq!(ids(1), ["w1", "w3", "w9"]);
 
-    let candidate = CacheCandidate {
-        worker: Arc::clone(&workers[9]),
-        matched_prefix_tokens: 0,
-        uncached_tokens: 128,
-        candidate_range_id: "global".into(),
-        max_pending_prefill_tokens: None,
-    };
-    assert_eq!(
-        selector
-            .bind_prefill_cache_candidate(candidate, request)
-            .expect("w9 belongs to the hash-indexed bucket")
-            .candidate_range_id,
-        "set"
-    );
     assert_eq!(
         selector
             .prefill_affinity_domain(&workers, &workers[9], request)

experimental/sgl-router/tests/proxy/bucket_routing.rs

@@ -222,6 +222,16 @@ fn set_native_load(
     worker_url: &str,
     num_total_tokens: u64,
     max_total_num_tokens: u64,
+) {
+    set_native_load_with_waiting(ctx, worker_url, num_total_tokens, max_total_num_tokens, 0);
+}
+
+fn set_native_load_with_waiting(
+    ctx: &AppContext,
+    worker_url: &str,
+    num_total_tokens: u64,
+    max_total_num_tokens: u64,
+    num_waiting_uncached_tokens: u64,
 ) {
     ctx.engine_load.set(
         worker_url,
@@ -232,7 +242,7 @@ fn set_native_load(
             num_tokens: num_total_tokens,
             max_total_num_tokens,
             native_cache: Some(NativeCacheRankLoad {
-                num_waiting_uncached_tokens: 0,
+                num_waiting_uncached_tokens,
                 num_total_tokens,
                 max_running_requests: 64,
                 total_prefill_uncached_tokens: 1,
@@ -694,17 +704,14 @@ async fn cache_winner_uses_target_uncached_work_before_prompt_length_bucket() {
 }
 
 #[tokio::test]
-async fn cache_candidate_bucket_binding_happens_before_candidate_limit() {
+async fn unbucketed_global_cache_winner_remains_eligible() {
     let best = crate::common::mock_worker::MockWorker::start(vec![]).await;
     let lower_ranked = crate::common::mock_worker::MockWorker::start(vec![]).await;
     let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
-    let mut best_bucket = bucket("p-best", BucketStage::Prefill, 10, "p-best");
-    best_bucket.min_extend_tokens = Some(32);
     let mut lower_ranked_bucket = bucket("p-lower", BucketStage::Prefill, 20, "p-lower");
     lower_ranked_bucket.min_extend_tokens = Some(32);
     let bucket_config = BucketConfig {
         buckets: vec![
-            best_bucket,
             lower_ranked_bucket,
             bucket("d-catch-all", BucketStage::Decode, 30, "d"),
         ],
@@ -736,16 +743,203 @@ async fn cache_candidate_bucket_binding_happens_before_candidate_limit() {
         .unwrap();
 
     assert_eq!(response.status(), StatusCode::OK);
-    wait_for_prefill(&lower_ranked).await;
+    wait_for_prefill(&best).await;
     assert!(
-        best.captured.lock().unwrap().last_body.is_none(),
+        lower_ranked.captured.lock().unwrap().last_body.is_none(),
-        "the top Indexer hit is Bucket-incompatible and must not consume K=1"
+        "an admitted unbucketed cache holder must remain globally eligible"
     );
     assert!(
         ctx.metrics.render().contains(
             r#"sgl_router_policy_decisions_total{policy="cache_aware",reason="cache_candidate"} 1"#
         ),
-        "the compatible lower-ranked cache holder must remain a cache candidate"
+        "the global cache winner must remain a cache candidate"
+    );
+}
+
+#[tokio::test]
+async fn cache_candidate_over_context_limit_falls_back_to_compatible_bucket() {
+    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-cached");
+    short_bucket.max_context_tokens = Some(32);
+    let bucket_config = BucketConfig {
+        buckets: vec![
+            short_bucket,
+            bucket("p-long", BucketStage::Prefill, 20, "p-fallback"),
+            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
+        ],
+        ttft_slo_policy: SloBucketPolicy::Disabled,
+        tps_slo_policy: SloBucketPolicy::Disabled,
+    };
+    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
+    let ctx = build_cache_ctx(
+        vec![
+            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
+            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
+            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
+        ],
+        bucket_config,
+        prefix_index,
+    );
+
+    let content = "context limit ".repeat(128);
+    let response = build_router(ctx)
+        .oneshot(chat_request_with_content(&content, None, Some(8), None))
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    wait_for_prefill(&fallback).await;
+    assert!(
+        cached.captured.lock().unwrap().last_body.is_none(),
+        "a cache holder that cannot serve the full context must not receive the request"
+    );
+}
+
+#[tokio::test]
+async fn cache_candidate_outside_slo_first_tier_falls_back_to_eligible_bucket() {
+    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let mut slow_bucket = bucket("p-slow", BucketStage::Prefill, 10, "p-cached");
+    slow_bucket.ttft_p95_at_capacity_ms = Some(400);
+    let mut fast_bucket = bucket("p-fast", BucketStage::Prefill, 20, "p-fallback");
+    fast_bucket.ttft_p95_at_capacity_ms = Some(100);
+    let bucket_config = BucketConfig {
+        buckets: vec![
+            slow_bucket,
+            fast_bucket,
+            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
+        ],
+        ttft_slo_policy: SloBucketPolicy::SloFirst,
+        tps_slo_policy: SloBucketPolicy::Disabled,
+    };
+    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
+    let ctx = build_cache_ctx(
+        vec![
+            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
+            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
+            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
+        ],
+        bucket_config,
+        prefix_index,
+    );
+
+    let content = "ttft tier ".repeat(128);
+    let response = build_router(ctx)
+        .oneshot(chat_request_with_content(
+            &content,
+            Some(200),
+            Some(8),
+            None,
+        ))
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    wait_for_prefill(&fallback).await;
+    assert!(
+        cached.captured.lock().unwrap().last_body.is_none(),
+        "SloFirst must exclude a cache holder outside the request's TTFT tier"
+    );
+}
+
+#[tokio::test]
+async fn cache_candidate_respects_its_bucket_pending_prefill_budget() {
+    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let mut limited_bucket = bucket("p-limited", BucketStage::Prefill, 10, "p-cached");
+    limited_bucket.max_pending_prefill_tokens = Some(64);
+    let bucket_config = BucketConfig {
+        buckets: vec![
+            limited_bucket,
+            bucket("p-fallback", BucketStage::Prefill, 20, "p-fallback"),
+            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
+        ],
+        ttft_slo_policy: SloBucketPolicy::Disabled,
+        tps_slo_policy: SloBucketPolicy::Disabled,
+    };
+    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
+    let ctx = build_cache_ctx(
+        vec![
+            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
+            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
+            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
+        ],
+        bucket_config,
+        prefix_index,
+    );
+    set_native_load_with_waiting(&ctx, &cached.url, 0, 100_000, 64);
+
+    let content = "pending budget ".repeat(128);
+    let response = build_router(Arc::clone(&ctx))
+        .oneshot(chat_request_with_content(&content, None, Some(8), None))
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    wait_for_prefill(&fallback).await;
+    assert!(
+        cached.captured.lock().unwrap().last_body.is_none(),
+        "a cache holder beyond its Bucket pending-prefill budget must not receive the request"
+    );
+    assert!(ctx
+        .metrics
+        .render()
+        .contains("sgl_router_cache_admission_rejected_total 1"));
+}
+
+#[tokio::test]
+async fn rejected_global_cache_candidate_falls_back_to_prompt_length_bucket() {
+    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-cached");
+    short_bucket.max_extend_tokens = Some(8);
+    let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-fallback");
+    long_bucket.min_extend_tokens = Some(9);
+    let bucket_config = BucketConfig {
+        buckets: vec![
+            short_bucket,
+            long_bucket,
+            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
+        ],
+        ttft_slo_policy: SloBucketPolicy::Disabled,
+        tps_slo_policy: SloBucketPolicy::Disabled,
+    };
+    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
+    let ctx = build_cache_ctx(
+        vec![
+            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
+            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
+            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
+        ],
+        bucket_config,
+        prefix_index,
+    );
+    set_native_load(&ctx, &cached.url, 1, 1);
+
+    let content = "admission fallback ".repeat(128);
+    let response = build_router(Arc::clone(&ctx))
+        .oneshot(chat_request_with_content(&content, None, Some(8), None))
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    wait_for_prefill(&fallback).await;
+    assert!(
+        cached.captured.lock().unwrap().last_body.is_none(),
+        "an admission-rejected cache holder must not receive the request"
+    );
+    let metrics = ctx.metrics.render();
+    assert!(metrics.contains("sgl_router_cache_admission_rejected_total 1"));
+    assert!(
+        !metrics.contains(
+            r#"sgl_router_policy_decisions_total{policy="cache_aware",reason="cache_candidate"} 1"#
+        ),
+        "Bucket fallback must not be reported as a cache-candidate decision"
     );
 }