minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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/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"
);
}