minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Vincent Gao <vincentbo@linux.alibaba.com>
This commit is contained in:
Vincent Gao
2026-09-11 17:41:27 +08:00
committed by GitHub
parent 822e73ccdd
commit 8e7deb329e
5 changed files with 272 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/src/policies/buckets.rs
+29 -30
View File
@@ -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
+3 -1
View File
@@ -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
+1 -1
View File
@@ -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
+36 -70
View File
@@ -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::<Vec<_>>(),
["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)
experimental/sgl-router/tests/proxy/bucket_routing.rs
+203 -9
View File
@@ -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<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"
);
}