diff --git a/experimental/sgl-router/src/policies/buckets.rs b/experimental/sgl-router/src/policies/buckets.rs index 45f112d13..7ed9443c4 100644 --- a/experimental/sgl-router/src/policies/buckets.rs +++ b/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 { - 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 { + 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); diff --git a/experimental/sgl-router/src/policies/cache_aware.rs b/experimental/sgl-router/src/policies/cache_aware.rs index de8812d55..23b9c20ce 100644 --- a/experimental/sgl-router/src/policies/cache_aware.rs +++ b/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(); } diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index 50a9f3243..64605262a 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/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, diff --git a/experimental/sgl-router/tests/component/policies/bucket_domains.rs b/experimental/sgl-router/tests/component/policies/bucket_domains.rs index 5c75e82e6..2a7f4cf83 100644 --- a/experimental/sgl-router/tests/component/policies/bucket_domains.rs +++ b/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 { - input_tokens: 256, - expected_peak_sequence_tokens: None, - ttft_slo_ms: Some(100), - tps_slo: None, - }; + + let domains = selector.prefill_domains( + &[short, long], + BucketRequest { + input_tokens: 256, + expected_peak_sequence_tokens: None, + ttft_slo_ms: None, + tps_slo: None, + }, + ); assert_eq!( - selector - .prefill_domains(&workers, request) + domains .iter() .map(|domain| domain.id.as_str()) .collect::>(), - ["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" + ["p-long"] ); } #[test] -fn cache_candidate_without_bucket_configuration_keeps_global_metadata() { - let p = worker("p", WorkerMode::Prefill); - let selector = BucketSelector::new(None); +fn cache_candidate_uses_bucket_metadata_without_extend_range_filtering() { + let cached = worker("cached", WorkerMode::Prefill); + 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, - matched_prefix_tokens: 64, - uncached_tokens: 64, - candidate_range_id: "probe".into(), - max_pending_prefill_tokens: Some(1), + worker: cached, + matched_prefix_tokens: 128, + uncached_tokens: 128, + candidate_range_id: "global".into(), + 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!(bound.max_pending_prefill_tokens, None); + assert_eq!(prepared.candidate_range_id, "p-cached"); + 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) diff --git a/experimental/sgl-router/tests/proxy/bucket_routing.rs b/experimental/sgl-router/tests/proxy/bucket_routing.rs index c482bcf19..b1438b947 100644 --- a/experimental/sgl-router/tests/proxy/bucket_routing.rs +++ b/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(), - "the top Indexer hit is Bucket-incompatible and must not consume K=1" + lower_ranked.captured.lock().unwrap().last_body.is_none(), + "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 = 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 = 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 = 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 = 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" ); }