diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 086e1cf3f..05dda5569 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -1443,3 +1443,24 @@ 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 stable length-policy retraction order and post-release cache eviction toward the survivor reserve + +The following repository-root files independently model stable length-policy retraction order and post-release cache eviction toward the survivor reserve: + +- `crates/core/src/engine/scheduler/sglang/core.rs` +- `crates/core/src/engine/scheduler/sglang/decode.rs` +- `crates/core/src/engine/scheduler/sglang/tests.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/managers/schedule_batch.py (_get_decode_retraction_order, retract_decode and release_req). + +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/scheduler/sglang/core.rs b/crates/core/src/engine/scheduler/sglang/core.rs index 0e57fda97..33deaa64c 100644 --- a/crates/core/src/engine/scheduler/sglang/core.rs +++ b/crates/core/src/engine/scheduler/sglang/core.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 //! Behavioral model of the SGLang scheduler loop (`Scheduler.get_next_batch_to_run`, //! `update_running_batch`, `ScheduleBatch.retract_decode`) as of sgl-project/sglang v0.5.6.post2 @@ -7,8 +10,9 @@ //! Rust from the observed semantics; no SGLang source is copied. The overlap-iteration wiring //! hosted here (the `HostLoop` drain, `inflight_source`, the prefix-cache commits deferred until //! the forward's result is observed, and the pricing of vision-cache misses) was verified against -//! v0.5.19 (`0bcd822`); batch selection, retraction and the prefill/decode interval keep the -//! `5c8bd8b5` pin. +//! v0.5.19 (`0bcd822`); batch selection and the prefill/decode interval keep the `5c8bd8b5` pin. +//! The decode retraction order and its post-release eviction (`retract_decode`, `release_req`) +//! were verified against v0.5.20 (`94602c9c`). //! //! The optional prefill/decode interval follows the scheduling contract at //! https://github.com/sgl-project/sglang/blob/20621aa14bda7726a8a968f326198eac61717fef/python/sglang/srt/managers/scheduler.py @@ -1236,7 +1240,8 @@ impl SglangCore { let observed = observed.unwrap_or_default(); // `maybe_cache_unfinished_req`: the observed forward's prefixes enter the // radix cache now, after this iteration's batch was selected. The commits - // are in batch order, which `running` keeps, so search on from the last match. + // are in batch order, which `running` keeps unless this pass retracted and + // reordered it, so search on from the last match, wrapping around. // A retracted or cancelled request's commits were discarded with its outputs. let mut cursor = 0; let len = self.running.len(); diff --git a/crates/core/src/engine/scheduler/sglang/decode.rs b/crates/core/src/engine/scheduler/sglang/decode.rs index 770766753..d646f75a8 100644 --- a/crates/core/src/engine/scheduler/sglang/decode.rs +++ b/crates/core/src/engine/scheduler/sglang/decode.rs @@ -1,6 +1,10 @@ // 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 +use std::cmp::Reverse; use std::time::Duration; use uuid::Uuid; @@ -18,7 +22,8 @@ use super::config::{SglangConfig, floor_to_block}; use super::request::SglangRequest; /// SGLang `SGLANG_RETRACT_DECODE_STEPS`: decode steps worth of KV reserved per surviving request -/// when `retract_decode` re-estimates `new_token_ratio`. +/// when `retract_decode` re-estimates `new_token_ratio`, and the per-request token reserve that +/// `release_req` evicts cached pages toward after each retraction. pub(super) const RETRACT_DECODE_STEPS: usize = 20; #[derive(Default)] @@ -235,12 +240,13 @@ fn check_decode_mem_for_burst( break; } - let Some((idx, _)) = running - .iter() - .enumerate() - .filter(|(_, req)| !req.pending_terminal) - .min_by_key(|(_, req)| req.output_len()) - else { + // `retract_decode` stably sorts the batch once, longest output first and then shortest + // prompt first, retracts from the end, and keeps the survivors in that order. Ghost rows + // are sorted with the batch but never retracted. + if retracted.is_empty() && running.iter().any(|req| !req.pending_terminal) { + running.sort_by_key(|req| (Reverse(req.output_len()), req.prompt_len())); + } + let Some(idx) = running.iter().rposition(|req| !req.pending_terminal) else { break; }; @@ -264,6 +270,13 @@ fn check_decode_mem_for_burst( logical_available_blocks_before: Some(logical_available_blocks_before), required_blocks_before: Some(required_blocks_before), }); + // `release_req` evicts cached pages toward a reserve for the remaining requests. Only the + // next step's demand decides whether to retract another request. + let reserve = running.len() * RETRACT_DECODE_STEPS; + let available = kv_manager.cache().available_tokens(); + if available < reserve { + kv_manager.evict(reserve - available); + } retracted.push(req); } diff --git a/crates/core/src/engine/scheduler/sglang/tests.rs b/crates/core/src/engine/scheduler/sglang/tests.rs index 1fca55bde..ced136f24 100644 --- a/crates/core/src/engine/scheduler/sglang/tests.rs +++ b/crates/core/src/engine/scheduler/sglang/tests.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 use super::request::WaitingQueue; @@ -674,12 +677,155 @@ fn fresh_prefill_tracks_cache_owned_prefix_pages_and_pressure_event() { assert_eq!(pressure.logical_available_blocks_before, Some(0)); assert_eq!(pressure.required_blocks_before, Some(1)); assert!(pressure.state_after.active_blocks <= pressure.state_before.active_blocks); - // The retracted request frees the page it owned beyond the cached prefix (one stored block - // removed); the cached first page survives, the second cached page had already been evicted - // for the blocker. - assert_eq!(removed_event_count(&buffer.drain()), 1); - assert_eq!(kv_manager.cache().page_pool.available(), 4); - assert_eq!(kv_manager.cache_mut().match_prefix(&prompt).0, 4); + // The retraction itself frees only the page the request owned beyond the cached prefix; the + // second cached page had already been evicted for the blocker. + assert_eq!(pressure.state_before.active_blocks, 3); + assert_eq!(pressure.state_after.active_blocks, 2); + // That leaves the cached first page unprotected, and `release_req` evicts it toward the + // blocker's 20-token reserve: two stored blocks are removed in total. + assert_eq!(removed_event_count(&buffer.drain()), 2); + assert_eq!(kv_manager.cache().page_pool.available(), 8); + assert_eq!(kv_manager.cache_mut().match_prefix(&prompt).0, 0); +} + +mod decode_retraction { + use super::*; + + /// A page-size-1 decode row that owns KV for its whole sequence. + fn row( + kv_manager: &mut SglangKvManager, + id: u128, + prompt: &[u32], + output: &[u32], + ) -> SglangRequest { + let sequence_tokens = [prompt, output].concat(); + let sequence_len = sequence_tokens.len(); + let allocation = kv_manager.allocate_for_request(&sequence_tokens).unwrap(); + SglangRequest { + is_decode_handoff: false, + uuid: Uuid::from_u128(id), + sequence_tokens, + prompt_len: prompt.len(), + max_output_tokens: 8, + planned_output_ids: None, + kv_lease: allocation.lease, + materialized_tokens: sequence_len, + allocated_tokens: sequence_len, + images: Vec::new(), + pending_terminal: false, + } + } + + /// A page-size-1 host-loop ghost row that finished on its first output token; the next + /// decode step still allocates that token's slot. + fn ghost_row( + kv_manager: &mut SglangKvManager, + id: u128, + prompt: &[u32], + output: u32, + ) -> SglangRequest { + let allocation = kv_manager.allocate_for_request(prompt).unwrap(); + let mut sequence_tokens = prompt.to_vec(); + sequence_tokens.push(output); + SglangRequest { + is_decode_handoff: false, + uuid: Uuid::from_u128(id), + sequence_tokens, + prompt_len: prompt.len(), + max_output_tokens: 1, + planned_output_ids: None, + kv_lease: allocation.lease, + materialized_tokens: prompt.len(), + allocated_tokens: prompt.len(), + images: Vec::new(), + pending_terminal: true, + } + } + + fn ids(requests: &[SglangRequest]) -> Vec { + requests + .iter() + .map(|request| request.uuid.as_u128()) + .collect() + } + + #[test] + fn victim_has_the_shortest_output_then_the_longest_prompt_and_is_the_latest_tie() { + // No free slot and each row needs one for its next token: retracting any three-token + // row lets the other three decode. + let config = SglangConfig::from_args(&test_args(11, 1, 16)); + let mut kv_manager = SglangKvManager::new(11, 1, KvEventPublishers::default(), 0); + let mut running = vec![ + row(&mut kv_manager, 1, &[10, 11], &[12]), + row(&mut kv_manager, 2, &[20, 21, 22], &[]), + row(&mut kv_manager, 3, &[30, 31, 32], &[]), + row(&mut kv_manager, 4, &[40, 41], &[]), + ]; + assert_eq!(kv_manager.cache().available_tokens(), 0); + + let retracted = decode::check_decode_mem(&mut running, &mut kv_manager, &config); + + // Rows 2 to 4 have no output and row 4 has the shortest prompt. Rows 2 and 3 tie + // exactly, and the later one is retracted first. + assert_eq!(ids(&retracted), vec![3]); + // The survivors keep the retraction order: longest output, then shortest prompt. + assert_eq!(ids(&running), vec![1, 4, 2]); + } + + #[test] + fn ghost_rows_are_sorted_with_the_batch_but_never_retracted() { + let config = SglangConfig::from_args(&test_args(10, 1, 16)); + let mut kv_manager = SglangKvManager::new(10, 1, KvEventPublishers::default(), 0); + // The ghost has the shortest output and the longest prompt, so it sorts last. + let mut running = vec![ + ghost_row(&mut kv_manager, 1, &[10, 11, 12], 13), + row(&mut kv_manager, 2, &[20, 21], &[22, 23]), + row(&mut kv_manager, 3, &[30, 31], &[32]), + ]; + assert_eq!(kv_manager.cache().available_tokens(), 0); + + let retracted = decode::check_decode_mem(&mut running, &mut kv_manager, &config); + + assert_eq!(ids(&retracted), vec![3]); + assert_eq!(ids(&running), vec![2, 1]); + assert!(running[1].pending_terminal); + + // With only ghost rows to choose from, nothing is retracted or reordered. + let mut kv_manager = SglangKvManager::new(5, 1, KvEventPublishers::default(), 0); + let mut running = vec![ + ghost_row(&mut kv_manager, 4, &[40, 41, 42], 43), + ghost_row(&mut kv_manager, 5, &[50, 51], 52), + ]; + assert!(decode::check_decode_mem(&mut running, &mut kv_manager, &config).is_empty()); + assert_eq!(ids(&running), vec![4, 5]); + } + + #[test] + fn retraction_evicts_toward_the_reserve_without_retracting_for_it() { + let config = SglangConfig::from_args(&test_args(11, 1, 16)); + let mut kv_manager = SglangKvManager::new(11, 1, KvEventPublishers::default(), 0); + let prefix = [1, 2, 3, 4]; + let cached = kv_manager.allocate_for_request(&prefix).unwrap(); + kv_manager.finish(&prefix, cached.lease); + // Only the row without output protects the cached prefix. + let mut running = vec![ + row(&mut kv_manager, 1, &[10, 11], &[12]), + row(&mut kv_manager, 2, &[20, 21], &[22]), + row(&mut kv_manager, 3, &[1, 2, 3, 4, 5], &[]), + ]; + assert_eq!(running[2].cached_tokens(), prefix.len()); + assert_eq!(kv_manager.cache().available_tokens(), 0); + assert_eq!(kv_manager.cache().evictable_size, 0); + + let retracted = decode::check_decode_mem(&mut running, &mut kv_manager, &config); + + // Its own slot covers the survivors' next step, so nothing else is retracted although + // even evicting the unprotected prefix stays short of 20 tokens for each survivor. + assert_eq!(ids(&retracted), vec![3]); + assert_eq!(ids(&running), vec![1, 2]); + assert_eq!(kv_manager.cache().prefix_match_len(&prefix), 0); + assert_eq!(kv_manager.cache().available_tokens(), 5); + } } mod source_holds { diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 086e1cf3f..05dda5569 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -1443,3 +1443,24 @@ 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 stable length-policy retraction order and post-release cache eviction toward the survivor reserve + +The following repository-root files independently model stable length-policy retraction order and post-release cache eviction toward the survivor reserve: + +- `crates/core/src/engine/scheduler/sglang/core.rs` +- `crates/core/src/engine/scheduler/sglang/decode.rs` +- `crates/core/src/engine/scheduler/sglang/tests.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/managers/schedule_batch.py (_get_decode_retraction_order, retract_decode and release_req). + +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.