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.