From 0edc5c2ed3557653fd1855d3a96fa6225c3cb0e3 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Wed, 23 Sep 2026 08:50:10 -0700 Subject: [PATCH 1/2] perf(sglang): reduce admission checkpoint copies Signed-off-by: jthomson04 --- crates/core/src/engine/cache/radix_cache.rs | 57 +++++++------- .../src/engine/kv_manager/sglang_backend.rs | 74 ++++++++++++------- .../core/src/engine/scheduler/sglang/tests.rs | 36 +++++++++ 3 files changed, 114 insertions(+), 53 deletions(-) diff --git a/crates/core/src/engine/cache/radix_cache.rs b/crates/core/src/engine/cache/radix_cache.rs index 07202b36d..ab67d6109 100644 --- a/crates/core/src/engine/cache/radix_cache.rs +++ b/crates/core/src/engine/cache/radix_cache.rs @@ -13,6 +13,7 @@ use rustc_hash::FxHashMap; use slotmap::{SlotMap, new_key_type}; use std::cmp::Reverse; use std::collections::BTreeSet; +use std::sync::Arc; use std::time::Instant; new_key_type! { @@ -169,9 +170,11 @@ pub struct TreeNode { /// page identity so completed radix state does not retain token IDs. /// Consequently, as in router-side indexing, hash collisions are treated /// as identical pages rather than guarded by an exact-token comparison. - pub key: Vec, + /// Admission snapshots share edge contents. Copy only an edge that changes + /// while a snapshot is live; normal decode extensions retain their capacity. + pub key: Arc>, /// One physical page ID per key. Length = `key.len()`. - pub value: Vec, + pub value: Arc>, /// Walk-to-root reference count (protected when > 0). pub lock_ref: usize, /// Monotonic timestamp for LRU eviction. @@ -231,8 +234,8 @@ impl RadixCache { let root = nodes.insert(TreeNode { children: FxHashMap::default(), parent: None, - key: Vec::new(), - value: Vec::new(), + key: Arc::default(), + value: Arc::default(), lock_ref: 0, last_access_time: Instant::now(), }); @@ -810,8 +813,8 @@ impl RadixCache { node.last_access_time = now; debug_assert!(node.children.is_empty()); debug_assert!(node.lock_ref <= 1); - node.key.extend_from_slice(key); - node.value.extend_from_slice(value); + Arc::make_mut(&mut node.key).extend_from_slice(key); + Arc::make_mut(&mut node.value).extend_from_slice(value); if node.lock_ref == 0 { self.evictable_size += key.len() * self.page_size; } else { @@ -830,10 +833,12 @@ impl RadixCache { let child = &mut self.nodes[child_id]; let child_parent = child.parent; let original_ck = child.key[0]; - let suffix_key = child.key.split_off(split_pos); - let mut prefix_key = std::mem::replace(&mut child.key, suffix_key); - let suffix_value = child.value.split_off(split_pos); - let mut prefix_value = std::mem::replace(&mut child.value, suffix_value); + let key = Arc::make_mut(&mut child.key); + let suffix_key = key.split_off(split_pos); + let mut prefix_key = std::mem::replace(key, suffix_key); + let value = Arc::make_mut(&mut child.value); + let suffix_value = value.split_off(split_pos); + let mut prefix_value = std::mem::replace(value, suffix_value); // The short prefix must not retain the original edge's allocation. prefix_key.shrink_to_fit(); prefix_value.shrink_to_fit(); @@ -855,8 +860,8 @@ impl RadixCache { let intermediate = TreeNode { children: inter_children, parent: child_parent, - key: prefix_key, - value: prefix_value, + key: Arc::new(prefix_key), + value: Arc::new(prefix_value), lock_ref, last_access_time: accessed, }; @@ -886,8 +891,8 @@ impl RadixCache { let new_node = TreeNode { children: FxHashMap::default(), parent: Some(parent_id), - key: key.to_vec(), - value: value.to_vec(), + key: Arc::new(key.to_vec()), + value: Arc::new(value.to_vec()), lock_ref: 0, last_access_time: now, }; @@ -1015,7 +1020,7 @@ impl RadixCache { let split_pos = victim_pages - eviction_pages; let (nodes, page_pool) = (&mut self.nodes, &mut self.page_pool); let victim_node = &mut nodes[victim_id]; - victim_node.key.truncate(split_pos); + Arc::make_mut(&mut victim_node.key).truncate(split_pos); let evicted_values = &victim_node.value[split_pos..]; if let Some(belady) = &mut self.belady { for page in evicted_values { @@ -1024,7 +1029,7 @@ impl RadixCache { } page_pool.free_pages(evicted_values); evicted_indices.extend_from_slice(evicted_values); - victim_node.value.truncate(split_pos); + Arc::make_mut(&mut victim_node.value).truncate(split_pos); self.evictable_size -= eviction_len; evicted += eviction_len; @@ -1044,7 +1049,7 @@ impl RadixCache { evicted_indices.extend_from_slice(&victim_node.value); if let Some(belady) = &mut self.belady { - for page in &victim_node.value { + for page in victim_node.value.iter() { belady.page_hashes[page.index()] = None; } } @@ -1239,9 +1244,9 @@ mod tests { let (len, node) = cache.match_prefix(&[1, 2, 3, 4, 5, 9, 9]); assert_eq!(len, 5); let n = cache.node(node); - assert_eq!(n.key, cache.page_hashes(&[1, 2, 3, 4, 5])); + assert_eq!(*n.key, cache.page_hashes(&[1, 2, 3, 4, 5])); assert_eq!( - n.value, + *n.value, vec![ KvPageId(10), KvPageId(20), @@ -1253,7 +1258,7 @@ mod tests { let suffix_key = cache.page_hashes(&[6])[0]; let &suffix_id = n.children.get(&suffix_key).unwrap(); assert_eq!( - cache.node(suffix_id).value, + *cache.node(suffix_id).value, vec![KvPageId(60), KvPageId(70)] ); } @@ -1331,7 +1336,7 @@ mod tests { assert_eq!(extended, tail); assert_eq!(cache.num_nodes(), nodes_before); assert_eq!( - cache.node(tail).key, + *cache.node(tail).key, cache.page_hashes(&[1, 2, 3, 4, 5, 6, 7, 8]) ); assert_eq!(cache.protected_size, 8); @@ -1356,8 +1361,8 @@ mod tests { assert_ne!(extended, tail); assert_eq!(cache.num_nodes(), nodes_before + 1); - assert_eq!(cache.node(tail).key, cache.page_hashes(&[1, 2, 3, 4])); - assert_eq!(cache.node(extended).key, cache.page_hashes(&[5, 6, 7, 8])); + assert_eq!(*cache.node(tail).key, cache.page_hashes(&[1, 2, 3, 4])); + assert_eq!(*cache.node(extended).key, cache.page_hashes(&[5, 6, 7, 8])); } #[test] @@ -1368,7 +1373,7 @@ mod tests { cache.insert(&[1, 2, 3, 4, 5, 6, 7], &[0, 1, 2, 3, 4, 5, 6]); assert_eq!(cache.match_prefix(&[1, 2, 3, 4]).0, 4); let (_, node) = cache.match_prefix(&[1, 2, 3, 4]); - assert_eq!(cache.node(node).value, vec![KvPageId(0)]); + assert_eq!(*cache.node(node).value, vec![KvPageId(0)]); cache.insert(&[1, 2, 3, 4, 5, 6, 7, 8], &[0, 1, 2, 3, 4, 5, 6, 7]); assert_eq!(cache.match_prefix(&[1, 2, 3, 4, 5, 6, 7, 8]).0, 8); @@ -1386,7 +1391,7 @@ mod tests { cache.insert(&[1, 2, 3, 4, 5, 6, 7, 8], &[0, 1, 2, 3, 4, 5, 6, 7]); cache.match_prefix(&[1, 2, 3, 4, 9, 9, 9, 9]); let (_, node) = cache.match_prefix(&[1, 2, 3, 4]); - assert_eq!(cache.node(node).value, vec![KvPageId(0)]); + assert_eq!(*cache.node(node).value, vec![KvPageId(0)]); } #[test] @@ -1398,7 +1403,7 @@ mod tests { assert_eq!(cache.node(node).key.len(), 2); assert_eq!(cache.node(node).value.len(), 2); - assert_eq!(cache.node(node).value, vec![KvPageId(0), KvPageId(1)]); + assert_eq!(*cache.node(node).value, vec![KvPageId(0), KvPageId(1)]); assert_eq!(cache.match_prefix(&tokens).0, tokens.len()); } diff --git a/crates/core/src/engine/kv_manager/sglang_backend.rs b/crates/core/src/engine/kv_manager/sglang_backend.rs index 0445e766a..3393dccf8 100644 --- a/crates/core/src/engine/kv_manager/sglang_backend.rs +++ b/crates/core/src/engine/kv_manager/sglang_backend.rs @@ -166,11 +166,18 @@ pub struct SglangKvManager { /// block hash so router events reflect logical block visibility, not /// transient slot ownership. block_hash_refcounts: FxHashMap, - pending_admission_events: Option>, + pending_admission: Option, } #[must_use = "admission must be committed or rolled back"] -pub(crate) struct SglangAdmissionCheckpoint { +pub(crate) struct SglangAdmission; + +struct PendingAdmission { + checkpoint: Option, + events: Vec, +} + +struct SglangAdmissionCheckpoint { cache: RadixCache, next_event_id: u64, page_to_block_hash: Vec>, @@ -267,7 +274,7 @@ impl SglangKvManager { next_event_id: 0, page_to_block_hash, block_hash_refcounts: FxHashMap::default(), - pending_admission_events: None, + pending_admission: None, } } @@ -285,44 +292,53 @@ impl SglangKvManager { } /// Hold speculative events until the timing provider accepts the admitted batch. - pub(crate) fn begin_admission(&mut self) -> SglangAdmissionCheckpoint { - assert!( - self.pending_admission_events.is_none(), - "nested KV admission" - ); - self.pending_admission_events = Some(Vec::new()); - SglangAdmissionCheckpoint { + pub(crate) fn begin_admission(&mut self) -> SglangAdmission { + assert!(self.pending_admission.is_none(), "nested KV admission"); + self.pending_admission = Some(PendingAdmission { + checkpoint: None, + events: Vec::new(), + }); + SglangAdmission + } + + /// Admission can reject all waiting requests without changing the cache. + /// Snapshot only before the first allocation or extension can mutate it. + fn checkpoint_admission(&mut self) { + let Some(admission) = self.pending_admission.as_mut() else { + return; + }; + if admission.checkpoint.is_some() { + return; + } + admission.checkpoint = Some(SglangAdmissionCheckpoint { cache: self.cache.admission_checkpoint(), next_event_id: self.next_event_id, page_to_block_hash: self.page_to_block_hash.clone(), block_hash_refcounts: self.block_hash_refcounts.clone(), - } + }); } - pub(crate) fn commit_admission(&mut self, _checkpoint: SglangAdmissionCheckpoint) { - for event in self - .pending_admission_events - .take() - .expect("active KV admission") - { + pub(crate) fn commit_admission(&mut self, _admission: SglangAdmission) { + let admission = self.pending_admission.take().expect("active KV admission"); + for event in admission.events { self.publish_event(event); } } - pub(crate) fn rollback_admission(&mut self, checkpoint: SglangAdmissionCheckpoint) { - self.pending_admission_events - .take() - .expect("active KV admission"); - self.cache = checkpoint.cache; - self.next_event_id = checkpoint.next_event_id; - self.page_to_block_hash = checkpoint.page_to_block_hash; - self.block_hash_refcounts = checkpoint.block_hash_refcounts; + pub(crate) fn rollback_admission(&mut self, _admission: SglangAdmission) { + let admission = self.pending_admission.take().expect("active KV admission"); + if let Some(checkpoint) = admission.checkpoint { + self.cache = checkpoint.cache; + self.next_event_id = checkpoint.next_event_id; + self.page_to_block_hash = checkpoint.page_to_block_hash; + self.block_hash_refcounts = checkpoint.block_hash_refcounts; + } self.log_trace("admission_rollback", 0); } fn publish_event(&mut self, event: KvEvent) { - if let Some(events) = self.pending_admission_events.as_mut() { - events.push(event); + if let Some(admission) = self.pending_admission.as_mut() { + admission.events.push(event); } else if let Err(error) = self.kv_event_publishers.publish(event, None) { tracing::warn!("Failed to publish SGLang KV event: {error}"); } @@ -362,6 +378,9 @@ impl SglangKvManager { let materialized_hashes = lease.page_hashes_through(token_ids.len(), page_size); let matchable_hashes = &materialized_hashes[..(max_prefix_tokens.min(token_ids.len()) / page_size)]; + // Prefix matching can split, touch and lock radix nodes even when the + // following capacity check rejects the allocation. + self.checkpoint_admission(); let (prefix_len, last_node) = self.match_prefix_hashes_and_lock(matchable_hashes); let required_pages = token_ids.len().div_ceil(page_size) - prefix_len / page_size; let required_tokens = required_pages * page_size; @@ -442,6 +461,7 @@ impl SglangKvManager { if required_tokens > reservable { return None; } + self.checkpoint_admission(); let available = self.cache.available_tokens(); if required_tokens > available { self.evict(required_tokens - available); diff --git a/crates/core/src/engine/scheduler/sglang/tests.rs b/crates/core/src/engine/scheduler/sglang/tests.rs index 7c3d1dbda..510f6237f 100644 --- a/crates/core/src/engine/scheduler/sglang/tests.rs +++ b/crates/core/src/engine/scheduler/sglang/tests.rs @@ -2703,6 +2703,42 @@ mod admission_validation_rollback { ) } + #[test] + fn failed_prefill_validation_restores_queue_when_no_request_fits() { + let mut core = SglangCore::new_with_kv_capture(test_args(32, 4, 32), 0); + core.receive(direct_request((0..8).collect(), 0)); + core.try_execute_hidden_pass(0.0).unwrap(); + core.config.max_running_requests = 1; + let running = core.receive(direct_request((20..24).collect(), 4)); + core.try_execute_hidden_pass(1.0).unwrap(); + let cold = core.receive(direct_request((100..112).collect(), 1)); + let warm = core.receive(direct_request((0..12).collect(), 1)); + core.config.schedule_policy = SchedulePolicy::Lpm; + core.drain_kv_events(); + + let waiting_before = format!("{:?}", core.waiting); + let capacity_before = capacity(&core); + let fail = Arc::new(AtomicBool::new(true)); + install_timing(&mut core, &fail, false); + // LPM changes queue order and computes lease hashes, but the running + // request fills the batch, so admission never changes the cache. + for now_ms in [2.0, 3.0] { + assert!(core.try_execute_hidden_pass(now_ms).is_err()); + assert_eq!(format!("{:?}", core.waiting), waiting_before); + assert_eq!(capacity(&core), capacity_before); + assert_eq!(core.running[0].uuid, running); + assert!(core.drain_kv_events().is_empty()); + } + + fail.store(false, Ordering::Relaxed); + let pass = core.try_execute_hidden_pass(4.0).unwrap(); + assert!(pass.admissions.is_empty()); + assert_eq!(core.waiting[0].uuid, warm); + assert_eq!(core.waiting[1].uuid, cold); + assert_eq!(pass.output_signals.len(), 1); + assert_eq!(pass.output_signals[0].uuid, running); + } + #[rstest::rstest] fn overlength_rejections_survive_failed_prefill_passes( #[values(false, true)] fail_in_prediction: bool, From 294e3a131138e42915be12e9ea5e1d867dad2db4 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Wed, 23 Sep 2026 09:27:39 -0700 Subject: [PATCH 2/2] perf(sglang): reduce remaining admission copies and test rollback Signed-off-by: jthomson04 --- crates/core/src/engine/cache/radix_cache.rs | 88 +++++++++++++++---- .../src/engine/kv_manager/sglang_backend.rs | 62 +++++++++++-- .../core/src/engine/scheduler/sglang/tests.rs | 58 ++++++++++++ 3 files changed, 185 insertions(+), 23 deletions(-) diff --git a/crates/core/src/engine/cache/radix_cache.rs b/crates/core/src/engine/cache/radix_cache.rs index ab67d6109..77561c7b6 100644 --- a/crates/core/src/engine/cache/radix_cache.rs +++ b/crates/core/src/engine/cache/radix_cache.rs @@ -210,6 +210,30 @@ struct BeladyLeaves { indexed: FxHashMap, } +/// Leave the suffix in `edge` and return a compact prefix. When a snapshot +/// shares the edge, copy each half once instead of cloning before splitting. +fn split_edge_prefix(edge: &mut Arc>, at: usize) -> Vec { + if let Some(values) = Arc::get_mut(edge) { + let suffix = values.split_off(at); + let mut prefix = std::mem::replace(values, suffix); + // The short prefix must not retain the original edge's allocation. + prefix.shrink_to_fit(); + prefix + } else { + let prefix = edge[..at].to_vec(); + *edge = Arc::new(edge[at..].to_vec()); + prefix + } +} + +fn truncate_edge(edge: &mut Arc>, len: usize) { + if let Some(values) = Arc::get_mut(edge) { + values.truncate(len); + } else { + *edge = Arc::new(edge[..len].to_vec()); + } +} + impl RadixCache { /// Snapshot allocator and radix metadata for a fallible admission transaction. /// This is not a second owner of live request leases; only one cache state may commit. @@ -833,15 +857,8 @@ impl RadixCache { let child = &mut self.nodes[child_id]; let child_parent = child.parent; let original_ck = child.key[0]; - let key = Arc::make_mut(&mut child.key); - let suffix_key = key.split_off(split_pos); - let mut prefix_key = std::mem::replace(key, suffix_key); - let value = Arc::make_mut(&mut child.value); - let suffix_value = value.split_off(split_pos); - let mut prefix_value = std::mem::replace(value, suffix_value); - // The short prefix must not retain the original edge's allocation. - prefix_key.shrink_to_fit(); - prefix_value.shrink_to_fit(); + let prefix_key = split_edge_prefix(&mut child.key, split_pos); + let prefix_value = split_edge_prefix(&mut child.value, split_pos); let suffix_ck = child.key[0]; ( child_parent, @@ -1020,7 +1037,7 @@ impl RadixCache { let split_pos = victim_pages - eviction_pages; let (nodes, page_pool) = (&mut self.nodes, &mut self.page_pool); let victim_node = &mut nodes[victim_id]; - Arc::make_mut(&mut victim_node.key).truncate(split_pos); + truncate_edge(&mut victim_node.key, split_pos); let evicted_values = &victim_node.value[split_pos..]; if let Some(belady) = &mut self.belady { for page in evicted_values { @@ -1029,7 +1046,7 @@ impl RadixCache { } page_pool.free_pages(evicted_values); evicted_indices.extend_from_slice(evicted_values); - Arc::make_mut(&mut victim_node.value).truncate(split_pos); + truncate_edge(&mut victim_node.value, split_pos); self.evictable_size -= eviction_len; evicted += eviction_len; @@ -1102,20 +1119,29 @@ mod tests { .unwrap() } - #[test] - fn belady_reranks_compressed_tail_before_evicting_its_prefix() { + #[rstest::rstest] + fn belady_reranks_compressed_tail_before_evicting_its_prefix( + #[values(false, true)] shared: bool, + ) { let mut cache = RadixCache::new(3, 1); cache.set_belady_oracle(oracle_for_prompts(&[&[1], &[3], &[1, 2]])); let ab = cache.page_pool.allocate(2).unwrap(); cache.insert(&[1, 2], &ab); let c = cache.page_pool.allocate(1).unwrap(); cache.insert(&[3], &c); + let snapshot = shared.then(|| cache.admission_checkpoint()); let (tokens, pages) = cache.evict(2); assert_eq!((tokens, pages), (2, vec![KvPageId(1), KvPageId(2)])); assert_eq!(cache.prefix_match_len(&[1, 2]), 1); assert_eq!(cache.prefix_match_len(&[3]), 0); assert_eq!(cache.available_tokens(), 2); + if let Some(mut saved) = snapshot { + assert_eq!(saved.prefix_match_len(&[1, 2]), 2); + assert_eq!(saved.prefix_match_len(&[3]), 1); + assert_eq!(saved.available_tokens(), 0); + assert_eq!(saved.evict(2).1, vec![KvPageId(1), KvPageId(2)]); + } } #[test] @@ -1291,15 +1317,31 @@ mod tests { assert_eq!(cache.match_prefix(&[1, 2, 3, 4, 5, 6, 7, 8]).0, 8); } - #[test] - fn repeated_prefix_splits_keep_retained_capacity_proportional_to_live_pages() { + #[rstest::rstest] + fn repeated_prefix_splits_keep_retained_capacity_proportional_to_live_pages( + #[values(false, true)] shared: bool, + ) { let mut cache = RadixCache::new(512, 1); let tokens: Vec = (0..512).collect(); let indices: Vec = (0..512).collect(); cache.insert(&tokens, &indices); for prefix_len in (1..512).step_by(4) { + let snapshot = shared.then(|| cache.admission_checkpoint()); + let saved_edges = snapshot.as_ref().map(|saved| { + saved + .nodes + .iter() + .map(|(id, node)| (id, (node.key.to_vec(), node.value.to_vec()))) + .collect::>() + }); assert_eq!(cache.match_prefix(&tokens[..prefix_len]).0, prefix_len); + if let Some(saved) = snapshot { + for (id, (key, pages)) in saved_edges.unwrap() { + assert_eq!(*saved.node(id).key, key); + assert_eq!(*saved.node(id).value, pages); + } + } } assert_eq!(cache.match_prefix(&tokens).0, tokens.len()); @@ -1516,8 +1558,10 @@ mod tests { assert_eq!(cache.match_prefix(&[4, 5, 6]).0, 3); } - #[test] - fn partial_eviction_and_parent_promotion_keep_index_consistent() { + #[rstest::rstest] + fn partial_eviction_and_parent_promotion_keep_index_consistent( + #[values(false, true)] shared: bool, + ) { let mut cache = RadixCache::new(100, 1); let base = Instant::now(); cache.set_test_now(base); @@ -1525,6 +1569,7 @@ mod tests { let (_, leaf) = cache.match_prefix(&[1, 2, 3, 4]); cache.set_test_now(base + std::time::Duration::from_secs(1)); cache.insert(&[5, 6, 7], &[4, 5, 6]); + let snapshot = shared.then(|| cache.admission_checkpoint()); assert_eq!(cache.evict(2).0, 2); assert_eq!(cache.node(leaf).key.len(), 2); @@ -1532,6 +1577,15 @@ mod tests { assert_eq!(cache.evict(2).0, 2); assert_eq!(cache.match_prefix(&[1, 2, 3, 4]).0, 0); assert_eq!(cache.match_prefix(&[5, 6, 7]).0, 3); + if let Some(mut saved) = snapshot { + assert_eq!(*saved.node(leaf).key, saved.page_hashes(&[1, 2, 3, 4])); + assert_eq!( + *saved.node(leaf).value, + vec![KvPageId(0), KvPageId(1), KvPageId(2), KvPageId(3)] + ); + assert!(saved.evictable_leaf_is_indexed(leaf)); + assert_eq!(saved.evict(2).1, vec![KvPageId(2), KvPageId(3)]); + } let mut branched = RadixCache::new(100, 1); branched.set_test_now(base); diff --git a/crates/core/src/engine/kv_manager/sglang_backend.rs b/crates/core/src/engine/kv_manager/sglang_backend.rs index 3393dccf8..1542d5da6 100644 --- a/crates/core/src/engine/kv_manager/sglang_backend.rs +++ b/crates/core/src/engine/kv_manager/sglang_backend.rs @@ -34,7 +34,7 @@ pub(crate) struct RadixRequestLease { /// Metadata used only with the matching KV-manager admission checkpoint. pub(crate) struct RadixLeaseCheckpoint { - pages: Vec, + pages_len: usize, materialized_tokens: usize, cached_tokens: usize, admission_reused_tokens: usize, @@ -45,7 +45,7 @@ pub(crate) struct RadixLeaseCheckpoint { impl RadixRequestLease { pub(crate) fn admission_checkpoint(&self) -> RadixLeaseCheckpoint { RadixLeaseCheckpoint { - pages: self.pages.clone(), + pages_len: self.pages.len(), materialized_tokens: self.materialized_tokens, cached_tokens: self.cached_tokens, admission_reused_tokens: self.admission_reused_tokens, @@ -55,8 +55,11 @@ impl RadixRequestLease { } /// Restore only after the matching manager checkpoint has been rolled back. + /// Admission only appends pages or fills an empty, inactive lease; it never + /// changes the existing page IDs. Prefix canonicalization happens after commit. pub(crate) fn restore_admission(&mut self, checkpoint: RadixLeaseCheckpoint) { - self.pages = checkpoint.pages; + debug_assert!(self.pages.len() >= checkpoint.pages_len); + self.pages.truncate(checkpoint.pages_len); self.materialized_tokens = checkpoint.materialized_tokens; self.cached_tokens = checkpoint.cached_tokens; self.admission_reused_tokens = checkpoint.admission_reused_tokens; @@ -170,7 +173,7 @@ pub struct SglangKvManager { } #[must_use = "admission must be committed or rolled back"] -pub(crate) struct SglangAdmission; +pub(crate) struct SglangAdmission(()); struct PendingAdmission { checkpoint: Option, @@ -298,11 +301,12 @@ impl SglangKvManager { checkpoint: None, events: Vec::new(), }); - SglangAdmission + SglangAdmission(()) } /// Admission can reject all waiting requests without changing the cache. /// Snapshot only before the first allocation or extension can mutate it. + /// Any new mutating path inside admission must checkpoint before its first change. fn checkpoint_admission(&mut self) { let Some(admission) = self.pending_admission.as_mut() else { return; @@ -332,12 +336,18 @@ impl SglangKvManager { self.next_event_id = checkpoint.next_event_id; self.page_to_block_hash = checkpoint.page_to_block_hash; self.block_hash_refcounts = checkpoint.block_hash_refcounts; + } else { + debug_assert!( + admission.events.is_empty(), + "events without a KV checkpoint" + ); } self.log_trace("admission_rollback", 0); } fn publish_event(&mut self, event: KvEvent) { if let Some(admission) = self.pending_admission.as_mut() { + debug_assert!(admission.checkpoint.is_some(), "event before KV checkpoint"); admission.events.push(event); } else if let Err(error) = self.kv_event_publishers.publish(event, None) { tracing::warn!("Failed to publish SGLang KV event: {error}"); @@ -373,6 +383,7 @@ impl SglangKvManager { max_prefix_tokens: usize, ) -> Option { assert!(!lease.is_active(), "request KV lease is already active"); + debug_assert!(lease.pages.is_empty(), "inactive KV lease retains pages"); let page_size = self.cache.page_size(); lease.ensure_page_hashes(token_ids, page_size); let materialized_hashes = lease.page_hashes_through(token_ids.len(), page_size); @@ -380,7 +391,9 @@ impl SglangKvManager { &materialized_hashes[..(max_prefix_tokens.min(token_ids.len()) / page_size)]; // Prefix matching can split, touch and lock radix nodes even when the // following capacity check rejects the allocation. - self.checkpoint_admission(); + if self.enable_prefix_caching { + self.checkpoint_admission(); + } let (prefix_len, last_node) = self.match_prefix_hashes_and_lock(matchable_hashes); let required_pages = token_ids.len().div_ceil(page_size) - prefix_len / page_size; let required_tokens = required_pages * page_size; @@ -393,6 +406,10 @@ impl SglangKvManager { self.cache.dec_lock_ref(last_node); return None; } + // Without prefix caching, matching and the capacity check are read-only. + if !self.enable_prefix_caching { + self.checkpoint_admission(); + } let available = self.cache.available_tokens(); if required_tokens > available { self.evict(required_tokens - available); @@ -1192,6 +1209,39 @@ mod tests { } } + #[test] + fn prefix_disabled_admission_checkpoints_only_allocations_that_fit() { + let mut mgr = SglangKvManager::new_with_prefix_caching( + 8, + 4, + KvEventPublishers::default(), + 0, + false, + false, + ); + let mut lease = RadixRequestLease::default(); + let before = lease.admission_checkpoint(); + let admission = mgr.begin_admission(); + assert_eq!(mgr.allocate_for_request_lease(&[1; 9], &mut lease), None); + assert!(mgr.pending_admission.as_ref().unwrap().checkpoint.is_none()); + mgr.rollback_admission(admission); + lease.restore_admission(before); + assert_eq!(mgr.cache.available_tokens(), 8); + assert!(!lease.is_active()); + assert!(lease.pages.is_empty()); + + let before = lease.admission_checkpoint(); + let admission = mgr.begin_admission(); + assert_eq!(mgr.allocate_for_request_lease(&[1; 8], &mut lease), Some(0)); + assert!(mgr.pending_admission.as_ref().unwrap().checkpoint.is_some()); + assert_eq!(mgr.cache.available_tokens(), 0); + mgr.rollback_admission(admission); + lease.restore_admission(before); + assert_eq!(mgr.cache.available_tokens(), 8); + assert!(!lease.is_active()); + assert!(lease.pages.is_empty()); + } + #[test] fn dense_page_metadata_is_allocated_only_for_kv_events() { let disabled = SglangKvManager::new(64, 4, KvEventPublishers::default(), 0); diff --git a/crates/core/src/engine/scheduler/sglang/tests.rs b/crates/core/src/engine/scheduler/sglang/tests.rs index 510f6237f..86f7b4ee2 100644 --- a/crates/core/src/engine/scheduler/sglang/tests.rs +++ b/crates/core/src/engine/scheduler/sglang/tests.rs @@ -2739,6 +2739,64 @@ mod admission_validation_rollback { assert_eq!(pass.output_signals[0].uuid, running); } + #[test] + fn failed_partial_prefix_match_restores_radix_state_and_eviction_order() { + let mut core = SglangCore::new_with_kv_capture(test_args(32, 4, 32), 0); + core.receive(direct_request((0..8).collect(), 0)); + core.try_execute_hidden_pass(0.0).unwrap(); + core.receive(direct_request((100..108).collect(), 0)); + core.try_execute_hidden_pass(1.0).unwrap(); + let request = core.receive(direct_request((0..4).chain(50..58).collect(), 1)); + core.drain_kv_events(); + + // Own the saved arrays so the assertion does not share the live cache's + // copy-on-write buffers. Compare topology and timestamps as well as counts. + let radix_state = |core: &SglangCore| { + let cache = core.kv_manager.cache(); + let mut nodes = rustc_hash::FxHashMap::default(); + let mut pending = vec![cache.root()]; + while let Some(id) = pending.pop() { + let node = cache.node(id); + pending.extend(node.children.values().copied()); + nodes.insert( + id, + ( + node.parent, + node.children.clone(), + node.key.to_vec(), + node.value.to_vec(), + node.lock_ref, + node.last_access_time, + ), + ); + } + let mut probe = cache.admission_checkpoint(); + let eviction_order = probe.evict(probe.evictable_size).1; + (nodes, eviction_order) + }; + let state_before = radix_state(&core); + let waiting_before = format!("{:?}", core.waiting); + let capacity_before = capacity(&core); + let fail = Arc::new(AtomicBool::new(true)); + install_timing(&mut core, &fail, false); + for now_ms in [2.0, 3.0] { + assert!(core.try_execute_hidden_pass(now_ms).is_err()); + assert_eq!(radix_state(&core), state_before); + assert_eq!(capacity(&core), capacity_before); + assert_eq!(format!("{:?}", core.waiting), waiting_before); + assert!(core.running.is_empty()); + assert!(core.drain_kv_events().is_empty()); + } + + fail.store(false, Ordering::Relaxed); + let pass = core.try_execute_hidden_pass(4.0).unwrap(); + assert_eq!(pass.admissions.len(), 1); + assert_eq!(pass.admissions[0].uuid, request); + assert_eq!(pass.admissions[0].reused_input_tokens, 4); + assert_eq!(pass.completed_requests, 1); + assert!(core.is_empty()); + } + #[rstest::rstest] fn overlength_rejections_survive_failed_prefill_passes( #[values(false, true)] fail_in_prediction: bool,