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

[Router] Add bucket-aware policy domains and native cache indexing (#38108)

Browse Source 
Signed-off-by: Vincent Gao <vincentbo@linux.alibaba.com>
Co-authored-by: inkcherry <mingzhi.liu@amd.com>
Co-authored-by: yangbodong22011 <13137470+yangbodong22011@users.noreply.github.com>
This commit is contained in:
Vincent Gao
2026-09-06 19:47:51 +08:00
committed by GitHub
co-authored by inkcherry yangbodong22011
parent a176ba2f7b
commit 5bebe7a033
76 changed files with 5842 additions and 3639 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/sgl-kv-indexer/tests/common/kv.rs
+22
View File
@@ -14,12 +14,22 @@ pub fn dram() -> i32 {
}
pub fn action(kind: ExternalKvActionType, tier: i32, hashes: &[i64]) -> ExternalKvAction {
action_with_parent(kind, tier, None, hashes)
}
pub fn action_with_parent(
kind: ExternalKvActionType,
tier: i32,
parent_block_hash: Option<i64>,
hashes: &[i64],
) -> ExternalKvAction {
ExternalKvAction {
r#type: kind as i32,
tier,
hashes: hashes.to_vec(),
component_masks: Vec::new(),
block_sizes: Vec::new(),
parent_block_hash,
}
}
@@ -31,6 +41,17 @@ pub fn component_report(
hashes: &[i64],
masks: &[u32],
block_sizes: &[u32],
) -> ExternalKvAction {
component_report_with_parent(tier, None, hashes, masks, block_sizes)
}
#[allow(dead_code)]
pub fn component_report_with_parent(
tier: i32,
parent_block_hash: Option<i64>,
hashes: &[i64],
masks: &[u32],
block_sizes: &[u32],
) -> ExternalKvAction {
ExternalKvAction {
r#type: ExternalKvActionType::ActionReport as i32,
@@ -38,6 +59,7 @@ pub fn component_report(
hashes: hashes.to_vec(),
component_masks: masks.to_vec(),
block_sizes: block_sizes.to_vec(),
parent_block_hash,
}
}
experimental/sgl-router/sgl-kv-indexer/tests/grpc_contract.rs
+99 -5
View File
@@ -34,7 +34,7 @@ use sgl_kv_indexer::{
PrefixIndex, PrefixIndexConfig, MAX_GRPC_DECODING_MESSAGE_SIZE,
};
use test_id::nanos;
use test_kv::{action, apply_request, hbm};
use test_kv::{action, action_with_parent, apply_request, hbm};
use test_net::free_addr;
async fn start_backend(
@@ -211,6 +211,27 @@ fn apply_report(
)
}
fn apply_report_with_parent(
worker: &str,
addr: &str,
seq: u64,
tier: i32,
parent_block_hash: Option<i64>,
hashes: &[i64],
) -> ApplyExternalKvBatchRequest {
apply_request(
worker,
addr,
seq,
vec![action_with_parent(
ExternalKvActionType::ActionReport,
tier,
parent_block_hash,
hashes,
)],
)
}
#[tokio::test]
async fn multiple_workers_share_one_indexer_server() {
let mut indexer = start().await;
@@ -225,7 +246,7 @@ async fn multiple_workers_share_one_indexer_server() {
"10.0.0.1:9000",
1,
hbm(),
&[hash_0, shared_hash],
&[shared_hash, hash_0],
))
.await
.expect("apply worker-0");
@@ -235,7 +256,7 @@ async fn multiple_workers_share_one_indexer_server() {
"10.0.0.2:9000",
1,
hbm(),
&[hash_1, shared_hash],
&[shared_hash, hash_1],
))
.await
.expect("apply worker-1");
@@ -309,6 +330,7 @@ async fn validation_errors_map_to_invalid_argument_over_grpc() {
hashes: vec![1],
component_masks: Vec::new(),
block_sizes: Vec::new(),
parent_block_hash: None,
}],
};
let err = c
@@ -318,6 +340,76 @@ async fn validation_errors_map_to_invalid_argument_over_grpc() {
assert_eq!(err.code(), Code::InvalidArgument);
}
#[tokio::test]
async fn rejected_batch_is_atomic_over_grpc() {
let mut c = start().await;
c.apply_external_kv_batch(apply_report("w", "old-address", 1, hbm(), &[1, 2]))
.await
.expect("seed chain");
let err = c
.apply_external_kv_batch(apply_request(
"w",
"new-address",
2,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[3]),
action_with_parent(ExternalKvActionType::ActionReport, hbm(), Some(9), &[2]),
],
))
.await
.expect_err("conflicting parent must reject the whole batch");
assert_eq!(err.code(), Code::InvalidArgument);
let old = c
.match_external_kv(MatchExternalKvRequest {
hashes: vec![1],
count_as_hit: false,
})
.await
.expect("query original state")
.into_inner();
assert_eq!(old.matches.len(), 1);
assert_eq!(old.matches[0].address, "old-address");
let leaked = c
.match_external_kv(MatchExternalKvRequest {
hashes: vec![3],
count_as_hit: false,
})
.await
.expect("query rejected action")
.into_inner();
assert!(leaked.matches.is_empty());
}
#[tokio::test]
async fn cyclic_report_is_rejected_over_grpc() {
let mut c = start().await;
let err = c
.apply_external_kv_batch(apply_report_with_parent(
"w",
"address",
1,
hbm(),
Some(2),
&[1, 2],
))
.await
.expect_err("cyclic report must be rejected");
assert_eq!(err.code(), Code::InvalidArgument);
let response = c
.match_external_kv(MatchExternalKvRequest {
hashes: vec![1, 2],
count_as_hit: false,
})
.await
.expect("query rejected report")
.into_inner();
assert!(response.matches.is_empty());
}
#[tokio::test]
async fn match_prefix_over_grpc() {
let mut c = start().await;
@@ -358,12 +450,14 @@ async fn prefix_query_scans_more_than_one_apply_chunk_over_grpc() {
let mut indexer = start().await;
let hashes: Vec<i64> = (0..=APPLY_CHUNK_SIZE as i64).collect();
for (seq, chunk) in hashes.chunks(APPLY_CHUNK_SIZE).enumerate() {
let parent_block_hash = (seq > 0).then_some(chunk[0] - 1);
indexer
.apply_external_kv_batch(apply_report(
.apply_external_kv_batch(apply_report_with_parent(
"large-prefix-worker",
"10.0.0.1:9000",
seq as u64,
hbm(),
parent_block_hash,
chunk,
))
.await
@@ -487,7 +581,7 @@ async fn start_recording_deadlines(
/// the only thing letting the indexer shed a query whose caller gave up.
#[tokio::test]
async fn router_client_publishes_its_deadline_on_the_wire() {
let (index, seen) = start_recording_deadlines(Duration::from_millis(250)).await;
let (index, seen) = start_recording_deadlines(Duration::from_secs(2)).await;
index
.match_prefix(vec![1, 2, 3])
experimental/sgl-router/sgl-kv-indexer/tests/memory_integration.rs
+266 -10
View File
@@ -17,7 +17,10 @@ use sgl_kv_indexer::pb::{
use sgl_kv_indexer::{
InMemoryKvIndexerBackend, KvIndexerBackend, WorkerPrefixInput, COMPONENT_FULL, COMPONENT_SWA,
};
use test_kv::{action, apply_request as apply_req, component_report, dram, hbm};
use test_kv::{
action, action_with_parent, apply_request as apply_req, component_report,
component_report_with_parent, dram, hbm,
};
use tonic::Status;
fn backend() -> InMemoryKvIndexerBackend {
@@ -81,6 +84,38 @@ itest!(report_then_match_returns_worker_and_address, b, {
assert!(tiers_for(&resp, "w1", 3).is_empty());
});
itest!(rejected_batch_does_not_publish_earlier_actions, b, {
b.apply_external_kv_batch(apply_req(
"w1",
"old-address",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])],
))
.await
.unwrap();
let error = b
.apply_external_kv_batch(apply_req(
"w1",
"new-address",
2,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[3]),
action_with_parent(ExternalKvActionType::ActionReport, hbm(), Some(9), &[2]),
],
))
.await
.unwrap_err();
assert_eq!(error.code(), tonic::Code::InvalidArgument);
let old = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert_eq!(old.matches.len(), 1);
assert_eq!(old.matches[0].address, "old-address");
let leaked = b.match_external_kv(match_req(&[3], false)).await.unwrap();
assert!(leaked.matches.is_empty());
});
itest!(large_request_preserves_complete_ordered_results, b, {
// Exercise a large write and read while preserving complete ordered results.
let expected_hashes: Vec<i64> = (0..300).collect();
@@ -587,10 +622,19 @@ async fn prefix_fast_path_matches_default_impl() {
fast.apply_external_kv_batch(report("w-short", "10.0.0.2:1", 1, &[1, 2]))
.await
.unwrap();
// w-hole holds 1, 3, 4 but not 2: strict prefix must be 1.
fast.apply_external_kv_batch(report("w-hole", "10.0.0.3:1", 1, &[1, 3, 4]))
// w-hole first learns the same chain, then loses block 2 while descendants
// remain placed: strict prefix must be 1.
fast.apply_external_kv_batch(report("w-hole", "10.0.0.3:1", 1, &[1, 2, 3, 4]))
.await
.unwrap();
fast.apply_external_kv_batch(apply_req(
"w-hole",
"10.0.0.3:1",
2,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[2])],
))
.await
.unwrap();
// w-noaddr is unroutable and must be excluded by both paths.
fast.apply_external_kv_batch(report("w-noaddr", "", 1, &[1, 2]))
.await
@@ -626,6 +670,64 @@ async fn prefix_fast_path_matches_default_impl() {
assert!(fast_resp.blocks_read >= 1);
}
#[tokio::test]
async fn prefix_fast_path_returns_worker_depths_with_uncached_suffix() {
let (fast, reference) = shared_state_pair();
fast.apply_external_kv_batch(report("w-long", "10.0.0.1:1", 1, &[1, 2, 3]))
.await
.unwrap();
fast.apply_external_kv_batch(report("w-short", "10.0.0.2:1", 1, &[1, 2]))
.await
.unwrap();
// Blocks 4 and 5 are the newly appended turn and are not cached anywhere.
// They must cap the maximum prefix without disabling the known-prefix path.
let query = [1, 2, 3, 4, 5];
let fast_response = fast
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
let reference_response = reference
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
assert_eq!(
prefix_pairs(&fast_response),
prefix_pairs(&reference_response)
);
assert_eq!(
prefix_pairs(&fast_response),
vec![("w-long".to_string(), 3), ("w-short".to_string(), 2)]
);
}
#[tokio::test]
async fn prefix_fast_path_falls_back_on_existing_parent_conflict() {
let (fast, reference) = shared_state_pair();
fast.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[1, 2]))
.await
.unwrap();
// Hash 9 is an independent root, not a child of hash 1.
fast.apply_external_kv_batch(report("w1", "10.0.0.1:1", 2, &[9]))
.await
.unwrap();
let query = [1, 9];
let fast_response = fast
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
let reference_response = reference
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
assert_eq!(
prefix_pairs(&fast_response),
prefix_pairs(&reference_response)
);
}
#[tokio::test]
async fn prefix_first_block_miss_reads_one_block() {
let b = backend();
@@ -662,6 +764,156 @@ async fn prefix_max_blocks_caps_the_scan() {
assert_eq!(resp.matches[0].matched_prefix_blocks, 2);
}
#[tokio::test]
async fn prefix_complete_revoke_and_restore_propagates_to_descendants() {
let (fast, reference) = shared_state_pair();
fast.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[1, 2, 3, 4]))
.await
.unwrap();
fast.apply_external_kv_batch(apply_req(
"w1",
"10.0.0.1:1",
2,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[2])],
))
.await
.unwrap();
let after_revoke = fast
.match_external_kv_prefix(prefix_req(&[1, 2, 3, 4]))
.await
.unwrap();
assert_eq!(prefix_pairs(&after_revoke), vec![("w1".to_string(), 1)]);
fast.apply_external_kv_batch(apply_req(
"w1",
"10.0.0.1:1",
3,
vec![action_with_parent(
ExternalKvActionType::ActionReport,
hbm(),
Some(1),
&[2],
)],
))
.await
.unwrap();
let restored = fast
.match_external_kv_prefix(prefix_req(&[1, 2, 3, 4]))
.await
.unwrap();
let expected = reference
.match_external_kv_prefix(prefix_req(&[1, 2, 3, 4]))
.await
.unwrap();
assert_eq!(prefix_pairs(&restored), vec![("w1".to_string(), 4)]);
assert_eq!(prefix_pairs(&restored), prefix_pairs(&expected));
}
#[tokio::test]
async fn prefix_complete_fast_path_preserves_cache_hit_rate() {
let (fast, reference) = shared_state_pair();
let query: Vec<i64> = (1..=64).collect();
let worker_count = 32usize;
let mut expected_prefix_sum = 0u64;
for worker in 0..worker_count {
let depth = 1 + (worker * 7 % query.len());
expected_prefix_sum += depth as u64;
fast.apply_external_kv_batch(report(
&format!("w-{worker:02}"),
&format!("http://worker-{worker:02}"),
1,
&query[..depth],
))
.await
.unwrap();
}
let fast_response = fast
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
let reference_response = reference
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
assert_eq!(
prefix_pairs(&fast_response),
prefix_pairs(&reference_response)
);
assert_eq!(fast_response.matches.len(), worker_count);
let fast_prefix_sum: u64 = fast_response
.matches
.iter()
.map(|item| item.matched_prefix_blocks as u64)
.sum();
let reference_prefix_sum: u64 = reference_response
.matches
.iter()
.map(|item| item.matched_prefix_blocks as u64)
.sum();
assert_eq!(fast_prefix_sum, expected_prefix_sum);
assert_eq!(fast_prefix_sum, reference_prefix_sum);
let hit_rate = fast_prefix_sum as f64 / (worker_count * query.len()) as f64;
let reference_hit_rate = reference_prefix_sum as f64 / (worker_count * query.len()) as f64;
assert!((hit_rate - reference_hit_rate).abs() < f64::EPSILON);
}
#[tokio::test]
async fn full_only_prefix_complete_fast_path_matches_reference() {
let (fast, reference) = shared_state_pair();
let full_spec = WorkerCacheSpec {
version: 1,
components: COMPONENT_FULL,
swa_window_tokens: 0,
full_tier_mask: (1 << hbm()) | (1 << dram()),
swa_tier_mask: 0,
mamba_tier_mask: 0,
};
fast.apply_external_kv_batch(apply_with_spec(
"w-full",
"10.0.0.1:1",
1,
full_spec,
vec![component_report(
hbm(),
&[11, 12, 13, 14],
&[COMPONENT_FULL; 4],
&[16; 4],
)],
))
.await
.unwrap();
fast.apply_external_kv_batch(apply_with_spec(
"w-full",
"10.0.0.1:1",
2,
full_spec,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[13])],
))
.await
.unwrap();
let query = [11, 12, 13, 14];
let fast_response = fast
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
let reference_response = reference
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
assert_eq!(
prefix_pairs(&fast_response),
prefix_pairs(&reference_response)
);
assert_eq!(
prefix_pairs(&fast_response),
vec![("w-full".to_string(), 2)]
);
}
// --- component-aware placement & prefix -------------------------------------
/// A hybrid-SWA spec: full servable from HBM+DRAM, swa a 100-token trailing
@@ -766,7 +1018,13 @@ async fn partial_eviction_replace_shrinks_component_set() {
"10.0.0.1:1",
2,
swa_spec(),
vec![component_report(hbm(), &[2], &[COMPONENT_FULL], &[80])],
vec![component_report_with_parent(
hbm(),
Some(1),
&[2],
&[COMPONENT_FULL],
&[80],
)],
))
.await
.unwrap();
@@ -827,12 +1085,10 @@ async fn duplicate_hash_in_one_report_keeps_last_snapshot() {
"10.0.0.1:1",
1,
swa_spec(),
vec![component_report(
hbm(),
&[1, 1],
&[COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL],
&[80, 80],
)],
vec![
component_report(hbm(), &[1], &[COMPONENT_FULL | COMPONENT_SWA], &[80]),
component_report(hbm(), &[1], &[COMPONENT_FULL], &[80]),
],
))
.await
.unwrap();