From c15ba51be5b498e8f657f642b488c4de49ffbf41 Mon Sep 17 00:00:00 2001 From: Rudy Pei Date: Fri, 23 Oct 2026 22:43:38 +0000 Subject: [PATCH] fix: refresh SGLang radix recency on splits and cached re-inserts SGLang refreshes every radix node that a match or insert walks, including a child whose edge is only partially matched and then split, so the retained suffix and the new prefix node are both most recently used (legacy `_match_prefix_helper` and `_insert_helper` in python/sglang/srt/mem_cache/radix_cache.py L731-L739 and L794-L813 at v0.5.20; the default unified tree core touches the walked child and refreshes it again in `_split_node`, unified_cache/unified_tree_core.py L1196-L1200 and L1396). Re-inserting a sequence whose pages are all cached walks and refreshes the same path. The simulator split a partially matched child without refreshing it, so its suffix kept the old timestamp, and an insert with no new pages returned before refreshing the path. Refresh the child before splitting it on the match and insert paths, and refresh the start path even when the inserted suffix is empty. Read-only admission matching and Belady's page-wise eviction are unchanged. Signed-off-by: Rudy Pei Co-authored-by: Isaac --- THIRD_PARTY_NOTICES.md | 19 +++ crates/core/src/engine/cache/radix_cache.rs | 136 +++++++++++++++++++- python/aisimulate/THIRD_PARTY_NOTICES.md | 19 +++ 3 files changed, 167 insertions(+), 7 deletions(-) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 086e1cf3f..30d14900b 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -1443,3 +1443,22 @@ https://github.com/alibaba/tair-kvcache/blob/85402f1a012fb4279be5bb4b57613bb0a06 The AISimulate files are modified re-implementations, not verbatim copies. Both derived files carry the source URLs and revisions in their headers. + +## SGLang v0.5.20 partial-edge access recency and reinsertion of an already cached prefix + +The following repository-root files independently model partial-edge access recency and reinsertion of an already cached prefix: + +- `crates/core/src/engine/cache/radix_cache.rs` + +The modified Rust behavior and synthetic regression fixtures adapt the contract +observed in SGLang v0.5.20, commit `94602c9c2b7cbdb8efd5c52802dac6a1c180089e`. +Original source paths: python/sglang/srt/mem_cache/radix_cache.py (_match_prefix_helper and _insert_helper). + +Source: https://github.com/sgl-project/sglang/tree/94602c9c2b7cbdb8efd5c52802dac6a1c180089e + +Copyright 2023-2024 SGLang Team and SGLang contributors. Apache-2.0; +the upstream license is at +https://github.com/sgl-project/sglang/blob/94602c9c2b7cbdb8efd5c52802dac6a1c180089e/LICENSE +and its full text is included in this repository as `LICENSE`. +Modified for AISimulate request leases, physical page accounting and replay. +No serving implementation, production request content or upstream test fixture is vendored. diff --git a/crates/core/src/engine/cache/radix_cache.rs b/crates/core/src/engine/cache/radix_cache.rs index 56b24e0db..fa3a5271f 100644 --- a/crates/core/src/engine/cache/radix_cache.rs +++ b/crates/core/src/engine/cache/radix_cache.rs @@ -1,5 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team // SPDX-License-Identifier: Apache-2.0 +// Modified behavioral adaptation of SGLang v0.5.20; see THIRD_PARTY_NOTICES.md. +// Source: https://github.com/sgl-project/sglang/tree/94602c9c2b7cbdb8efd5c52802dac6a1c180089e //! Radix-tree KV cache for SGLang engine simulation. //! @@ -518,6 +521,9 @@ impl RadixCache { if common_len < child_len { if common_len > 0 { + // Refresh the matched child before splitting it, so both the retained + // suffix and the new prefix node it seeds count as just accessed. + self.touch_node(child_id, now); let intermediate = self.split_node(child_id, common_len); current = intermediate; } @@ -537,7 +543,7 @@ impl RadixCache { self.touch_node(id, now); } } else if lock_match { - // A partial edge match keeps the split node's inherited timestamp. + // A split node inherits the timestamp just given to its suffix. self.inc_lock_ref(current); } @@ -735,9 +741,6 @@ impl RadixCache { prefix_len <= aligned_len, "prefix length {prefix_len} exceeds aligned key length {aligned_len}" ); - if aligned_len == prefix_len { - return start_node; - } let page_keys = self.page_hashes(&key[prefix_len..aligned_len]); self.insert_page_hash_suffix( start_node, @@ -760,11 +763,12 @@ impl RadixCache { page_keys.len(), page_ids.len() ); + // Re-inserting an already cached sequence still refreshes its path. + let now = self.now(); + self.touch_path(start_node, now); if page_keys.is_empty() { return start_node; } - let now = self.now(); - self.touch_path(start_node, now); let mut current = start_node; let mut key_offset = 0; @@ -807,11 +811,12 @@ impl RadixCache { child_key.len(), ) }; + // Refresh a partially matched child too, before it is split. + self.touch_node(child_id, now); if common_len == child_len { key_offset += common_len; current = child_id; - self.touch_node(current, now); } else { if common_len > 0 { let intermediate = self.split_node(child_id, common_len); @@ -1589,6 +1594,123 @@ mod tests { assert_eq!(cache.match_prefix(&[4, 5, 6]).0, 3); } + /// An older prompt cached one second before a two-page newer prompt, with the test clock + /// one more second later: unless something refreshes it, the older prompt is evicted first. + struct OlderAndNewer { + cache: RadixCache, + older: Vec, + older_indices: Vec, + older_node: NodeId, + newer: Vec, + newer_indices: Vec, + } + + impl OlderAndNewer { + fn new(page_size: usize, older_pages: usize) -> Self { + let page_tokens = page_size as u32; + let mut cache = RadixCache::new((older_pages + 2) * page_size, page_size); + let older = (0..older_pages as u32 * page_tokens).collect::>(); + let newer = (1_000..1_000 + 2 * page_tokens).collect::>(); + let base = Instant::now(); + cache.set_test_now(base); + let older_indices = cache.page_pool.allocate(older.len()).unwrap(); + let older_node = cache.insert(&older, &older_indices); + cache.set_test_now(base + std::time::Duration::from_secs(1)); + let newer_indices = cache.page_pool.allocate(newer.len()).unwrap(); + cache.insert(&newer, &newer_indices); + cache.set_test_now(base + std::time::Duration::from_secs(2)); + Self { + cache, + older, + older_indices, + older_node, + newer, + newer_indices, + } + } + + fn assert_newer_is_evicted_first(mut self) { + let page_size = self.cache.page_size(); + let newer_pages = self.cache.page_ids(&self.newer_indices, 2); + assert_eq!( + self.cache.evict(2 * page_size), + (2 * page_size, newer_pages) + ); + assert_eq!(self.cache.prefix_match_len(&self.newer), 0); + assert_eq!(self.cache.prefix_match_len(&self.older), self.older.len()); + } + } + + #[rstest::rstest] + fn partial_match_refreshes_the_suffix_it_splits( + #[values(1, 4, 64)] page_size: usize, + #[values(false, true)] lock: bool, + ) { + let mut fixture = OlderAndNewer::new(page_size, 4); + let prefix = &fixture.older[..2 * page_size]; + let (matched, node) = if lock { + let hashes = fixture.cache.page_hashes(prefix); + fixture.cache.match_prefix_hashes_and_lock(&hashes) + } else { + fixture.cache.match_prefix(prefix) + }; + assert_eq!(matched, prefix.len()); + if lock { + fixture.cache.dec_lock_ref(node); + } + fixture.assert_newer_is_evicted_first(); + } + + #[rstest::rstest] + fn partial_insert_refreshes_the_suffix_it_splits( + #[values(1, 4, 64)] page_size: usize, + #[values(false, true)] hashed: bool, + ) { + let mut fixture = OlderAndNewer::new(page_size, 4); + let prefix = &fixture.older[..2 * page_size]; + let prefix_indices = &fixture.older_indices[..2 * page_size]; + if hashed { + let hashes = fixture.cache.page_hashes(prefix); + let pages = fixture.cache.page_ids(prefix_indices, hashes.len()); + let root = fixture.cache.root(); + fixture + .cache + .insert_page_hashes_from_node(root, 0, &hashes, &pages); + } else { + fixture.cache.insert(prefix, prefix_indices); + } + fixture.assert_newer_is_evicted_first(); + } + + #[rstest::rstest] + fn reinserting_a_cached_sequence_without_new_pages_refreshes_its_path( + #[values(1, 4, 64)] page_size: usize, + #[values(false, true)] hashed: bool, + ) { + // A request holding the older prompt finishes without completing another page, so + // the insert that caches its sequence has no suffix to add. + let mut fixture = OlderAndNewer::new(page_size, 2); + let node = fixture.older_node; + fixture.cache.inc_lock_ref(node); + let inserted = if hashed { + let hashes = fixture.cache.page_hashes(&fixture.older); + let pages = fixture.cache.page_ids(&fixture.older_indices, hashes.len()); + fixture + .cache + .insert_page_hashes_from_node(node, fixture.older.len(), &hashes, &pages) + } else { + fixture.cache.insert_from_node( + node, + fixture.older.len(), + &fixture.older, + &fixture.older_indices, + ) + }; + assert_eq!(inserted, node); + fixture.cache.dec_lock_ref(node); + fixture.assert_newer_is_evicted_first(); + } + #[test] fn partial_eviction_and_parent_promotion_keep_index_consistent() { let mut cache = RadixCache::new(100, 1); diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 086e1cf3f..30d14900b 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -1443,3 +1443,22 @@ https://github.com/alibaba/tair-kvcache/blob/85402f1a012fb4279be5bb4b57613bb0a06 The AISimulate files are modified re-implementations, not verbatim copies. Both derived files carry the source URLs and revisions in their headers. + +## SGLang v0.5.20 partial-edge access recency and reinsertion of an already cached prefix + +The following repository-root files independently model partial-edge access recency and reinsertion of an already cached prefix: + +- `crates/core/src/engine/cache/radix_cache.rs` + +The modified Rust behavior and synthetic regression fixtures adapt the contract +observed in SGLang v0.5.20, commit `94602c9c2b7cbdb8efd5c52802dac6a1c180089e`. +Original source paths: python/sglang/srt/mem_cache/radix_cache.py (_match_prefix_helper and _insert_helper). + +Source: https://github.com/sgl-project/sglang/tree/94602c9c2b7cbdb8efd5c52802dac6a1c180089e + +Copyright 2023-2024 SGLang Team and SGLang contributors. Apache-2.0; +the upstream license is at +https://github.com/sgl-project/sglang/blob/94602c9c2b7cbdb8efd5c52802dac6a1c180089e/LICENSE +and its full text is included in this repository as `LICENSE`. +Modified for AISimulate request leases, physical page accounting and replay. +No serving implementation, production request content or upstream test fixture is vendored.