From ac9111d8e80cffde0b7f09501e0eb5d155a2896e Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 6 Oct 2026 23:28:12 -0700 Subject: [PATCH] fix: emit SGLang-mode KV events at SGLang's insert points and cadence The SGLang scheduler published a Stored event when it allocated prompt pages, and it inserted each decode page into the radix tree as soon as the page filled. At page size 1 a request therefore produced about one Stored event per output token, plus one Removed event for every decode step that evicted. SGLang v0.5.21 behaves differently: - it inserts a request's prompt into its radix tree after the prefill forward (`cache_unfinished_req`, once per chunk for chunked prefill). It does so whether or not the request's first decode slot then fits, and for a disaggregated prefill while its KV transfers. It inserts the rest when it releases the finished request (`release_kv_cache`), including a running request that an abort finished. Decoded tokens are neither matchable nor published before then, and retraction releases them without inserting them (`release_kv_cache(is_insert=False)`); - it queues events in the tree cache. A removal merges into a removal at the tail of the queue, and a store whose parent is the last block of the tail store extends that store; - it publishes the queue after every prefill batch, after every `--decode-log-interval` (default 40) decode batches, and when idle. The SGLang KV manager and scheduler now do the same. Allocation no longer publishes, and decode no longer inserts per token. The post-prefill insert covers the prompt only. It runs even when the first decode slots cannot be reserved, and also for requests that finish with their first token, which include disaggregated prefill sources. Cancelling a running request finishes it, so the KV it decoded is inserted. Events go through a queue that merges like SGLang's and assigns event IDs at publish time, and the queue is flushed at SGLang's publish points. Commands flush the queue when no request is running. Behavior changes in SGLang mode: - KV events are fewer and larger, and are published at the next publish point (up to 40 decode passes later) instead of immediately. Stored and removed block totals are unchanged. Pages that never enter the radix tree are no longer announced as stored or removed. - KV-aware routing sees a request's decoded blocks only when the request finishes or is cancelled, as with SGLang. - A disaggregated prefill source's prompt is inserted and published after its prefill pass. Before, it was published when its pages were allocated, and it entered the radix tree only when the source was released. - Decoded tokens of in-flight requests are no longer prefix-matchable by the simulator's own radix cache. - A retracted request discards its decoded KV and recomputes it when it is readmitted; only its prompt stays cached. - Simulated TTFT and throughput can therefore change under KV pressure. vLLM and TRT-LLM modes are unchanged. We recorded a multi-turn coding-agent workload on SGLang 0.5.21 at page size 1 (16 closed-loop lanes, 6,262 requests). Real SGLang emitted 2.00 Stored events per request at 370 blocks per event (median 98), and 0.49 Removed events per request at 1,375 blocks per event. Replaying the same schedule, the simulator went from 112.2 to 2.00 Stored events per request (6.6 to 371 blocks per event, median 98), and from 11.4 to 1.98 Removed events per request (59 to 341 blocks per event). A 4-lane run with a quarter of the KV capacity agrees on stores (real and simulated: 2.00 per request at 383 and 384 blocks per event). At page size 16, Stored events per request fell from 7.7 to 1.9. The removal gap that remains comes from eviction granularity. SGLang evicts whole radix leaves, while this cache evicts only the pages it needs. Evicting whole leaves in a local experiment brought the simulator to 0.47 Removed events per request at 1,437 blocks per event. That change alters cache retention, so it is left for a follow-up. Signed-off-by: PeaBrane --- .../src/engine/kv_manager/sglang_backend.rs | 138 ++++++++-- .../core/src/engine/scheduler/sglang/core.rs | 56 ++++- .../src/engine/scheduler/sglang/decode.rs | 20 +- .../core/src/engine/scheduler/sglang/tests.rs | 237 +++++++++++++++++- 4 files changed, 404 insertions(+), 47 deletions(-) diff --git a/crates/core/src/engine/kv_manager/sglang_backend.rs b/crates/core/src/engine/kv_manager/sglang_backend.rs index 8d53160c2..42f38551e 100644 --- a/crates/core/src/engine/kv_manager/sglang_backend.rs +++ b/crates/core/src/engine/kv_manager/sglang_backend.rs @@ -171,6 +171,10 @@ pub struct SglangKvManager { /// transient slot ownership. block_hash_refcounts: FxHashMap, pending_admission: Option, + /// Committed events not yet handed to the publishers. Like SGLang's radix + /// cache event queue, a new event merges into a compatible tail, and the + /// scheduler drains the queue only where SGLang publishes a batch. + queued_events: Vec, } #[must_use = "admission must be committed or rolled back"] @@ -287,6 +291,7 @@ impl SglangKvManager { page_to_block_hash, block_hash_refcounts: FxHashMap::default(), pending_admission: None, + queued_events: Vec::new(), } } @@ -383,8 +388,33 @@ impl SglangKvManager { 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}"); + } else { + self.enqueue_event(event); + } + } + + fn enqueue_event(&mut self, event: KvEvent) { + let unmerged = match self.queued_events.last_mut() { + Some(tail) => merge_into_tail(tail, event), + None => Some(event), + }; + if let Some(event) = unmerged { + self.queued_events.push(event); + } + } + + /// Hand every queued event to the publishers, numbering them in publish order. + pub(crate) fn flush_events(&mut self) { + debug_assert!( + self.pending_admission.is_none(), + "flush during KV admission" + ); + for mut event in std::mem::take(&mut self.queued_events) { + event.event_id = self.next_event_id; + self.next_event_id += 1; + if let Err(error) = self.kv_event_publishers.publish(event, None) { + tracing::warn!("Failed to publish SGLang KV event: {error}"); + } } } @@ -457,10 +487,8 @@ impl SglangKvManager { }; pages.append(&mut new_pages); let allocated_tokens = available_before - self.cache.available_tokens(); - - // Observer-visible KV events are complete-block only. - self.publish_stored_hashes(materialized_hashes, &pages, prefix_len, token_ids); - + // New pages become visible to routers only when their tokens are inserted into the + // radix tree, as in SGLang (`cache_unfinished_req` / `release_kv_cache`). self.log_trace("allocation", allocated_tokens); lease.pages = pages; @@ -523,13 +551,6 @@ impl SglangKvManager { lease.materialized_tokens = token_ids.len(); lease.ensure_page_hashes(token_ids, page_size); let allocated_tokens = available_before - self.cache.available_tokens(); - - self.publish_stored_hashes( - lease.page_hashes_through(token_ids.len(), page_size), - &lease.pages, - prefix_len, - token_ids, - ); self.log_trace("allocation", allocated_tokens); Some(prefix_len) } @@ -1141,8 +1162,9 @@ impl SglangKvManager { return hashed_blocks; } + // `flush_events` assigns the published event ID. let event = KvEvent { - event_id: self.next_event_id, + event_id: 0, data: KvEventData::Stored(StoredBlocks { parent_hash, start_position: None, @@ -1151,8 +1173,6 @@ impl SglangKvManager { dp_rank: self.dp_rank, tier: KvEventTier::Device, }; - self.next_event_id += 1; - self.publish_event(event); hashed_blocks @@ -1190,17 +1210,41 @@ impl SglangKvManager { } let event = KvEvent { - event_id: self.next_event_id, + event_id: 0, data: KvEventData::Removed { block_hashes }, dp_rank: self.dp_rank, tier: KvEventTier::Device, }; - self.next_event_id += 1; - self.publish_event(event); } } +/// Merge `event` into `tail` where SGLang's radix cache event queue would, and +/// return it unchanged otherwise. +fn merge_into_tail(tail: &mut KvEvent, event: KvEvent) -> Option { + if tail.dp_rank != event.dp_rank || tail.tier != event.tier { + return Some(event); + } + match (&mut tail.data, event.data) { + ( + KvEventData::Removed { + block_hashes: tail_hashes, + }, + KvEventData::Removed { block_hashes }, + ) => { + tail_hashes.extend(block_hashes); + None + } + (KvEventData::Stored(tail_stored), KvEventData::Stored(stored)) + if tail_stored.blocks.last().map(|block| block.block_hash) == stored.parent_hash => + { + tail_stored.blocks.extend(stored.blocks); + None + } + (_, data) => Some(KvEvent { data, ..event }), + } +} + #[cfg(test)] mod tests { use super::*; @@ -1217,6 +1261,18 @@ mod tests { lease } + /// Publish a lease's complete pages without inserting them, as a publication + /// inside an admission would. + fn publish_lease(mgr: &mut SglangKvManager, lease: &RadixRequestLease, tokens: &[u32]) { + let page_size = mgr.cache().page_size(); + mgr.publish_stored_hashes( + lease.page_hashes_through(tokens.len(), page_size), + lease.pages(), + 0, + tokens, + ); + } + struct MockSink { events: Mutex>, } @@ -1459,6 +1515,7 @@ mod tests { mgr.activate_destination_lease(reservation, &tokens[..4], &mut lease), 0 ); + mgr.flush_events(); let events = sink.clone_events(); let [ @@ -1631,8 +1688,10 @@ mod tests { let tokens = [1, 2, 3, 4, 5, 6, 7, 8, 9]; let mut alloc = mgr.allocate_for_request(&tokens[..4]).unwrap(); mgr.extend_cached_prefix(&tokens[..4], &mut alloc.lease); + mgr.flush_events(); mgr.extend_allocation(&tokens, &mut alloc.lease).unwrap(); mgr.finish(&tokens, alloc.lease); + mgr.flush_events(); let events = sink.clone_events(); let stored = events .iter() @@ -1654,10 +1713,12 @@ mod tests { let cached = mgr.allocate_for_request(&tokens[..8]).unwrap(); assert_eq!(cached.prefix_len, 8); mgr.finish(&tokens[..8], cached.lease); + mgr.flush_events(); assert_eq!(sink.event_count(), event_count); mgr.evict(8); let reused = mgr.allocate_for_request(&tokens[..4]).unwrap(); mgr.finish(&tokens[..4], reused.lease); + mgr.flush_events(); let events = sink.clone_events(); let KvEventData::Stored(stored) = &events.last().unwrap().data else { panic!("expected a store after eviction"); @@ -1692,13 +1753,17 @@ mod tests { let mut mgr = SglangKvManager::new(100, 1, KvEventPublishers::new(Some(sink.clone())), 0); let r = mgr.allocate_for_request(&[1, 2, 3]).unwrap(); - assert_eq!(sink.event_count(), 1); // BlockStored for 3 new pages + mgr.flush_events(); + assert_eq!(sink.event_count(), 0); // Allocated pages are not in the radix tree yet. mgr.finish(&[1, 2, 3], r.lease); + mgr.flush_events(); + assert_eq!(sink.event_count(), 1); // BlockStored for 3 new pages // Second request with full cache hit → no new events let r2 = mgr.allocate_for_request(&[1, 2, 3]).unwrap(); assert_eq!(r2.prefix_len, 3); + mgr.flush_events(); assert_eq!(sink.event_count(), 1); // no new event } @@ -1709,10 +1774,12 @@ mod tests { let first = mgr.allocate_for_request(&[1]).unwrap(); let page = first.lease.pages()[0]; - assert!(mgr.abort(first.lease)); + mgr.finish(&[1], first.lease); let second = mgr.allocate_for_request(&[2]).unwrap(); - assert_eq!(second.lease.pages(), &[page]); + assert_eq!(second.lease.pages(), &[page]); // Eviction reuses the page. + mgr.finish(&[2], second.lease); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 3); @@ -1743,19 +1810,23 @@ mod tests { let first = mgr.allocate_for_request(&[1]).unwrap(); let page = first.lease.pages()[0]; mgr.finish(&[1], first.lease); + mgr.flush_events(); let pages_before = mgr.page_to_block_hash.clone(); let counts_before = mgr.block_hash_refcounts.clone(); let admission = mgr.begin_admission(); let second = mgr.allocate_for_request(&[2]).unwrap(); assert_eq!(second.lease.pages(), &[page]); // Eviction reuses the page. + publish_lease(&mut mgr, &second.lease, &[2]); assert!(mgr.abort(second.lease)); let third = mgr.allocate_for_request(&[3]).unwrap(); assert_eq!(third.lease.pages(), &[page]); // Mutate the same slot again. + publish_lease(&mut mgr, &third.lease, &[3]); assert_eq!(sink.event_count(), 1); // Nothing speculative is published. if commit { mgr.commit_admission(admission); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 5); for (index, event) in events.iter().enumerate() { @@ -1781,6 +1852,7 @@ mod tests { assert_eq!(sink.event_count(), 1); let retry = mgr.allocate_for_request(&[2]).unwrap(); assert_eq!(retry.lease.pages(), &[page]); + mgr.flush_events(); assert_eq!(sink.clone_events()[1].event_id, 1); mgr.abort(retry.lease); } @@ -1794,7 +1866,9 @@ mod tests { let pages_before = mgr.page_to_block_hash.clone(); let admission = mgr.begin_admission(); let first = mgr.allocate_for_request(&[1]).unwrap(); + publish_lease(&mut mgr, &first.lease, &[1]); let second = mgr.allocate_for_request(&[1]).unwrap(); + publish_lease(&mut mgr, &second.lease, &[1]); let hash = mgr.page_to_block_hash[first.lease.pages()[0].index()].unwrap(); assert_eq!(mgr.block_hash_refcounts[&hash], 2); mgr.abort(first.lease); @@ -1809,9 +1883,12 @@ mod tests { assert_eq!(sink.event_count(), 0); let first = mgr.allocate_for_request(&[1]).unwrap(); + publish_lease(&mut mgr, &first.lease, &[1]); + mgr.flush_events(); let pages_before = mgr.page_to_block_hash.clone(); let admission = mgr.begin_admission(); let second = mgr.allocate_for_request(&[1]).unwrap(); + publish_lease(&mut mgr, &second.lease, &[1]); assert_eq!(mgr.block_hash_refcounts[&hash], 2); mgr.rollback_admission(admission); assert_eq!(mgr.page_to_block_hash, pages_before); @@ -1819,6 +1896,7 @@ mod tests { assert_eq!(sink.event_count(), 1); drop(second); // The speculative lease was invalidated by rollback. mgr.abort(first.lease); + mgr.flush_events(); assert_eq!(sink.event_count(), 2); } @@ -1829,6 +1907,7 @@ mod tests { let r = mgr.allocate_for_request(&[1, 2, 3, 4, 5, 6]).unwrap(); mgr.finish(&[1, 2, 3, 4, 5, 6], r.lease); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 1); @@ -1855,6 +1934,7 @@ mod tests { mgr.publish_stored_hashes(lease.page_hashes_through(4, 4), &pages[..1], 0, &tokens), 1 ); + mgr.flush_events(); assert_eq!( mgr.publish_stored_hashes(lease.page_hashes_through(8, 4), &pages, 0, &tokens), 1 @@ -1863,6 +1943,7 @@ mod tests { mgr.publish_stored_hashes(lease.page_hashes_through(8, 4), &pages, 0, &tokens), 0 ); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 2); @@ -1885,6 +1966,7 @@ mod tests { let mut alloc = mgr.allocate_for_request(&tokens[..2]).unwrap(); mgr.extend_cached_prefix(&tokens[..2], &mut alloc.lease); + mgr.flush_events(); assert_eq!(sink.event_count(), 1); assert!( @@ -1892,6 +1974,7 @@ mod tests { .is_some() ); mgr.extend_cached_prefix(&tokens[..4], &mut alloc.lease); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 2); let KvEventData::Stored(first_store) = &events[1].data else { @@ -1905,6 +1988,7 @@ mod tests { assert!(mgr.extend_allocation(&tokens, &mut alloc.lease).is_some()); mgr.finish(&tokens, alloc.lease); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 3); let KvEventData::Stored(final_store) = &events[2].data else { @@ -1924,6 +2008,9 @@ mod tests { let req1 = mgr.allocate_for_request(&[1, 2, 3]).unwrap(); let req2 = mgr.allocate_for_request(&[1, 2, 3]).unwrap(); + publish_lease(&mut mgr, &req1.lease, &[1, 2, 3]); + publish_lease(&mut mgr, &req2.lease, &[1, 2, 3]); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 1); @@ -1933,9 +2020,11 @@ mod tests { assert_eq!(store.blocks.len(), 3); mgr.free_pages(&req1.lease.pages); + mgr.flush_events(); assert_eq!(sink.event_count(), 1); mgr.free_pages(&req2.lease.pages); + mgr.flush_events(); let events = sink.clone_events(); assert_eq!(events.len(), 2); let KvEventData::Removed { block_hashes } = &events[1].data else { @@ -2062,9 +2151,12 @@ mod tests { let mut alloc1 = mgr.allocate_for_request(&tokens[..chunk1_len]).unwrap(); mgr.extend_cached_prefix(&tokens[..chunk1_len], &mut alloc1.lease); + mgr.flush_events(); - let alloc2 = mgr.allocate_for_request(&tokens[..chunk2_len]).unwrap(); + let mut alloc2 = mgr.allocate_for_request(&tokens[..chunk2_len]).unwrap(); + mgr.extend_cached_prefix(&tokens[..chunk2_len], &mut alloc2.lease); assert!(mgr.abort(alloc1.lease)); + mgr.flush_events(); let events = sink.events.lock().unwrap(); assert_eq!(events.len(), 2, "expected two stored events"); diff --git a/crates/core/src/engine/scheduler/sglang/core.rs b/crates/core/src/engine/scheduler/sglang/core.rs index d7b292568..9c6147b08 100644 --- a/crates/core/src/engine/scheduler/sglang/core.rs +++ b/crates/core/src/engine/scheduler/sglang/core.rs @@ -48,6 +48,10 @@ use crate::engine::scheduler::{ capture_kv_event_sink, }; +/// SGLang `--decode-log-interval` default. SGLang publishes its queued KV events after every +/// prefill batch, after every this many decode batches, and when the scheduler goes idle. +const DECODE_LOG_INTERVAL: usize = 40; + pub(crate) struct SglangCore { pub(super) config: SglangConfig, dp_rank: u32, @@ -73,6 +77,8 @@ pub(crate) struct SglangCore { prefill_in_pass: bool, model_work_in_pass: bool, interval_idle_in_pass: bool, + /// SGLang `forward_ct_decode`: decode forwards run so far, which paces KV event publishing. + decode_forwards: usize, } struct HeldSglangPrefill { @@ -223,6 +229,7 @@ impl SglangCore { prefill_in_pass: false, model_work_in_pass: false, interval_idle_in_pass: false, + decode_forwards: 0, #[cfg(test)] destination_reservation_attempts: 0, lifecycle_events: Vec::new(), @@ -257,6 +264,23 @@ impl SglangCore { &mut self, command: SchedulerCommand, allow_destination_admission: bool, + ) -> anyhow::Result { + let effects = self.dispatch_command(command, allow_destination_admission); + self.publish_kv_events_if_idle(); + effects + } + + /// An idle SGLang scheduler publishes queued KV events on its next loop iteration. + fn publish_kv_events_if_idle(&mut self) { + if self.running.is_empty() { + self.kv_manager.flush_events(); + } + } + + fn dispatch_command( + &mut self, + command: SchedulerCommand, + allow_destination_admission: bool, ) -> anyhow::Result { match command { SchedulerCommand::Submit(mut request) => { @@ -440,6 +464,8 @@ impl SglangCore { let reservation = self .kv_manager .reserve_destination_lease(request.kv_lease.page_hashes(), request.prompt_len()); + // Reservation can evict outside a pass, including at a rank's pass-completion boundary. + self.publish_kv_events_if_idle(); self.pending_destinations.mark_front_attempted(generation); let Some(kv) = reservation else { return Vec::new(); @@ -571,26 +597,26 @@ impl SglangCore { } fn cancel_active_request(&mut self, request_id: Uuid) -> bool { - let request = if let Some(index) = self + let (request, was_running) = if let Some(index) = self .waiting .iter() .position(|request| request.uuid == request_id) { - self.waiting.remove(index) + (self.waiting.remove(index), false) } else if let Some(index) = self .prebuilt_ready .iter() .position(|request| request.uuid == request_id) { - self.prebuilt_ready.remove(index) + (self.prebuilt_ready.remove(index), false) } else if let Some(index) = self .running .iter() .position(|request| request.uuid == request_id) { - Some(self.running.remove(index)) + (Some(self.running.remove(index)), true) } else { - None + (None, false) }; let Some(mut request) = request else { return false; @@ -599,7 +625,15 @@ impl SglangCore { if let Some(oracle) = &self.belady { oracle.retire_requests([request_id]); } - let capacity_improved = self.kv_manager.abort(std::mem::take(&mut request.kv_lease)); + // SGLang aborts a running request by finishing it, and the finish (`release_kv_cache`) + // inserts the KV it computed. Other requests release their KV without inserting it. + let capacity_improved = if was_running { + let releases_kv = request.kv_lease.is_active(); + cleanup_completed_request(&mut request, &mut self.kv_manager, self.config.block_size); + releases_kv + } else { + self.kv_manager.abort(std::mem::take(&mut request.kv_lease)) + }; self.source_holds.remove_request(request_id); self.active_destination_handoffs.remove_request(request_id); if capacity_improved { @@ -1118,6 +1152,16 @@ impl SglangCore { .max(self.config.min_new_token_ratio); } + let decode_forward = !prefill_pass && self.model_work_in_pass; + if decode_forward { + self.decode_forwards += 1; + } + if !decode_forward || self.decode_forwards.is_multiple_of(DECODE_LOG_INTERVAL) { + self.kv_manager.flush_events(); + } else { + self.publish_kv_events_if_idle(); + } + // Build FPM snapshot now that all state has settled. // Radix-cache reuse: admission matches over prompt tokens processed this pass. A chunked // request's own earlier chunks are KV context for the forward but not cache hits. diff --git a/crates/core/src/engine/scheduler/sglang/decode.rs b/crates/core/src/engine/scheduler/sglang/decode.rs index da63da613..2722df713 100644 --- a/crates/core/src/engine/scheduler/sglang/decode.rs +++ b/crates/core/src/engine/scheduler/sglang/decode.rs @@ -286,14 +286,26 @@ fn prefill_first_tokens( newly_completed.push(running.remove(idx)); } newly_completed.reverse(); - completed_requests.extend(newly_completed); // The rest need a slot for the first output token. Make room by evicting cached pages only. + // SGLang does this in the next decode pass, where the removal merges with that pass's + // queued removals; evicting before the stores below keeps it from becoming an extra event. let needed = decode_page_growth_needed(running, config.block_size, 1); let available = kv_manager.cache().available_tokens(); if available < needed { kv_manager.evict(needed - available); } + + // SGLang inserts every prefilled prompt right after the prefill forward. A continuing request + // goes through `cache_unfinished_req` whether or not its first slot fits, so a request + // retracted at the next decode step leaves its prompt cached. A finished one goes through + // `release_kv_cache`, and a disaggregated prefill source through `cache_unfinished_req` while + // its KV transfers. The first output token's KV is not computed yet and is not inserted. + for req in running.iter_mut().chain(newly_completed.iter_mut()) { + cache_materialized_prefix(req, kv_manager, config); + } + completed_requests.extend(newly_completed); + let reserved_pages = needed / config.block_size; let Some(mut reservation) = kv_manager.reserve_decode_pages(reserved_pages) else { return DecodeResult { @@ -311,7 +323,6 @@ fn prefill_first_tokens( } let token_id = req.next_output_token(); req.append_output_token(token_id, config.block_size); - cache_materialized_prefix(req, kv_manager, config); req.debug_assert_invariants(config.block_size); output_signals.push(OutputSignal { uuid: req.uuid, @@ -577,9 +588,8 @@ fn simulate_step( completed_indices.push(idx); break; } - - cache_materialized_prefix(req, kv_manager, config); - req.debug_assert_invariants(config.block_size); + // Decoded tokens stay out of the radix tree until the request finishes: + // SGLang inserts them only when it releases a finished request (`release_kv_cache`). } } diff --git a/crates/core/src/engine/scheduler/sglang/tests.rs b/crates/core/src/engine/scheduler/sglang/tests.rs index d3c8fa5dd..db956fcc2 100644 --- a/crates/core/src/engine/scheduler/sglang/tests.rs +++ b/crates/core/src/engine/scheduler/sglang/tests.rs @@ -568,7 +568,10 @@ fn fresh_prefill_tracks_cache_owned_prefix_pages_and_pressure_event() { materialized_tokens: 3, allocated_tokens: 4, }; - buffer.drain(); + // The request recomputes the second page into a fresh, still private page, so the evicted + // cached copy was the only visible one: routers see its removal right away. + kv_manager.flush_events(); + assert_eq!(removed_event_count(&buffer.drain()), 1); let mut running = vec![req, blocker]; let (retracted, pressure_events) = decode::check_decode_mem_with_pressure_events( @@ -594,10 +597,10 @@ 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); + // The retracted request frees the page it owned beyond the cached prefix. That page never + // entered the radix tree, so no removal is published; the cached first page survives. + kv_manager.flush_events(); + assert_eq!(removed_event_count(&buffer.drain()), 0); assert_eq!(kv_manager.cache().page_pool.available(), 4); assert_eq!(kv_manager.cache_mut().match_prefix(&prompt).0, 4); } @@ -679,6 +682,34 @@ mod source_holds { assert_eq!(occupied_tokens(&core), released_tokens); } + /// SGLang caches a disaggregated prefill's prompt right after the prefill forward + /// (`cache_unfinished_req` in `process_batch_result_disagg_prefill`) and publishes it with + /// that batch, so routers and later prefills see it while its KV transfers. + #[test] + fn held_source_prompt_is_cached_and_published_after_prefill() { + let mut core = SglangCore::new_with_kv_capture(args(), 0); + let handoff_id = HandoffId::from(Uuid::from_u128(408)); + let mut prefill = request(Uuid::from_u128(308)); + prefill.max_output_tokens = 1; + let prompt = prefill.tokens.clone(); + core.apply_command(SchedulerCommand::SubmitHandoffPrefill { + handoff_id, + request: prefill, + }) + .unwrap(); + + let pass = execute(&mut core, 0.0); + assert!(pass.output_signals[0].completed); + assert!(core.source_is_held(handoff_id)); + assert_eq!(stored_hashes(&pass.kv_events).len(), 2); + assert_eq!(core.kv_manager.cache().prefix_match_len(&prompt), 8); + + core.apply_command(SchedulerCommand::ReleaseSource { handoff_id }) + .unwrap(); + assert!(core.drain_kv_events().is_empty()); + assert_eq!(core.kv_manager.cache().prefix_match_len(&prompt), 8); + } + #[test] fn cancel_and_early_release_cleanup_exactly_once() { let mut core = SglangCore::new(args()); @@ -1010,8 +1041,8 @@ mod destination_lifecycle { assert_eq!(ready.kv_pages(), reserved_pages); assert_eq!(destination.running.len(), 1); assert!(destination.kv_manager.cache().protected_size >= protected_before_activation); - let activation_stores = stored_hashes(&destination.drain_kv_events()); - assert!(!activation_stores.is_empty()); + // The decode worker is busy, so the activation insert waits for SGLang's next publish. + assert!(destination.drain_kv_events().is_empty()); assert_eq!( source @@ -1033,9 +1064,12 @@ mod destination_lifecycle { .prebuilt_request(logical_uuid) .expect("full running batch must keep request ready"); assert_eq!(ready.kv_pages(), reserved_pages); - assert_no_republished_stores(&activation_stores, &stored_hashes(&blocked.kv_events)); + assert!(stored_hashes(&blocked.kv_events).is_empty()); let blocker_terminal = execute(&mut destination, blocked.end_ms); + // The blocker's completion idles the running batch, which publishes the activation. + let activation_stores = stored_hashes(&blocker_terminal.kv_events); + assert!(!activation_stores.is_empty()); assert!( blocker_terminal .admissions @@ -1362,6 +1396,91 @@ mod core_behavior { } } + /// Run `core` until it drains and return each pass's KV events as (kind, blocks) pairs, + /// keyed by pass index and skipping passes that published nothing. + fn kv_event_shapes_by_pass(core: &mut SglangCore) -> Vec<(usize, Vec<(&'static str, usize)>)> { + let mut collector = crate::engine::trace::TraceCollector::default(); + let mut now_ms = 0.0; + let mut next_event_id = None; + let mut shapes = Vec::new(); + for pass_index in 0.. { + if core.is_empty() { + break; + } + let pass = core.execute_pass(&mut collector, now_ms); + now_ms = pass.end_ms; + for event in &pass.kv_events { + let expected_id = *next_event_id.get_or_insert(event.event_id); + assert_eq!( + event.event_id, expected_id, + "published event IDs must be dense" + ); + next_event_id = Some(expected_id + 1); + } + if pass.kv_events.is_empty() { + continue; + } + let shape = pass + .kv_events + .iter() + .map(|event| match &event.data { + KvEventData::Stored(stored) => ("stored", stored.blocks.len()), + KvEventData::Removed { block_hashes } => ("removed", block_hashes.len()), + }) + .collect(); + shapes.push((pass_index, shape)); + } + shapes + } + + /// SGLang inserts a request into its radix tree after the prefill forward + /// (`cache_unfinished_req`) and again when the request finishes (`release_kv_cache`), so + /// routers see two stores per request and no decoded block before the request finishes. + #[rstest::rstest] + fn kv_events_store_prompt_after_prefill_and_outputs_at_finish( + #[values(1, 4)] page_size: usize, + ) { + let mut core = SglangCore::new_with_kv_capture(test_args(64, page_size, 64), 0); + core.receive(direct_request((1..=8).collect(), 50)); + + // One prefill pass emits the first token; 49 decode passes emit the rest. The prefill + // insert covers the prompt only: the first output token's KV is not computed yet. + let prefill_blocks = 8 / page_size; + let finished_blocks = 58 / page_size; + assert_eq!( + kv_event_shapes_by_pass(&mut core), + vec![ + (0, vec![("stored", prefill_blocks)]), + (49, vec![("stored", finished_blocks - prefill_blocks)]), + ] + ); + } + + /// SGLang merges consecutive removals in its event queue and publishes the queue after each + /// prefill batch, every `--decode-log-interval` (40) decode batches, and when it goes idle. + #[test] + fn kv_event_removals_are_batched_like_sglang() { + // Sixty cached one-token prompts fill the pool except for room for an eight-token prompt, + // so every output token of the next request evicts one cached page. + let mut core = SglangCore::new_with_kv_capture(test_args(68, 1, 64), 0); + for token in 1_000..1_060 { + let cached = core.kv_manager.allocate_for_request(&[token]).unwrap(); + core.kv_manager.finish(&[token], cached.lease); + } + core.kv_manager.flush_events(); + core.drain_kv_events(); + core.receive(direct_request((1..=8).collect(), 50)); + + assert_eq!( + kv_event_shapes_by_pass(&mut core), + vec![ + (0, vec![("removed", 1), ("stored", 8)]), + (40, vec![("removed", 40)]), + (49, vec![("removed", 9), ("stored", 50)]), + ] + ); + } + #[test] fn test_chunked_prefill_admits_next_chunk_when_full_prompt_does_not_fit() { let config = SglangConfig { @@ -1548,7 +1667,7 @@ mod core_behavior { } #[test] - fn test_unfinished_decode_request_is_cached_after_output() { + fn test_decoded_tokens_are_cached_only_when_the_request_finishes() { let config = SglangConfig::from_args( &MockEngineArgs::builder() .engine_type(EngineType::Sglang) @@ -1558,22 +1677,114 @@ mod core_behavior { .unwrap(), ); let mut kv_manager = SglangKvManager::new(64, 4, KvEventPublishers::default(), 0); - let alloc = kv_manager.allocate_for_request(&[1, 2, 3, 4]).unwrap(); + let mut alloc = kv_manager.allocate_for_request(&[1, 2, 3, 4]).unwrap(); + // SGLang `cache_unfinished_req` after the prefill forward. + kv_manager.extend_cached_prefix(&[1, 2, 3, 4], &mut alloc.lease); let mut running = vec![SglangRequest { is_decode_handoff: false, uuid: Uuid::new_v4(), sequence_tokens: vec![1, 2, 3, 4], prompt_len: 4, - max_output_tokens: 4, + max_output_tokens: 5, planned_output_ids: None, kv_lease: alloc.lease, materialized_tokens: 4, allocated_tokens: 4, }]; + for _ in 0..4 { + simulate_decode_step(&mut running, &mut kv_manager, &config, 0.0, false); + } + assert_eq!(running.len(), 1); + let sequence = running[0].sequence_prefix(8).to_vec(); + assert_eq!(kv_manager.cache().prefix_match_len(&sequence), 4); + simulate_decode_step(&mut running, &mut kv_manager, &config, 0.0, false); - let prefix = running[0].sequence_prefix(4); - assert_eq!(kv_manager.cache().prefix_match_len(prefix), 4); + assert!(running.is_empty()); + assert_eq!(kv_manager.cache().prefix_match_len(&sequence), 8); + } + + /// SGLang inserts the prompt after the prefill forward even when the first output slots do + /// not fit, and retraction (`release_kv_cache(is_insert=False)`) keeps that cached prompt. + #[test] + fn prompt_is_cached_after_prefill_when_first_output_slots_do_not_fit() { + let config = SglangConfig::from_args(&test_args(4, 4, 16)); + let (buffer, sink) = capture_kv_event_sink(); + // Four pages: the two prompts take three, and their first output tokens need two more. + let mut kv_manager = SglangKvManager::new(16, 4, KvEventPublishers::new(Some(sink)), 0); + let mut prefilled = |uuid: u128, prompt: Vec| { + let alloc = kv_manager.allocate_for_request(&prompt).unwrap(); + SglangRequest { + is_decode_handoff: false, + uuid: Uuid::from_u128(uuid), + prompt_len: prompt.len(), + materialized_tokens: prompt.len(), + allocated_tokens: prompt.len(), + sequence_tokens: prompt, + max_output_tokens: 4, + planned_output_ids: None, + kv_lease: alloc.lease, + } + }; + let short_prompt: Vec = (1..=4).collect(); + let long_prompt: Vec = (11..=18).collect(); + let mut running = vec![ + prefilled(90_101, short_prompt.clone()), + prefilled(90_102, long_prompt.clone()), + ]; + + let first = + decode::simulate_prefill_first_tokens(&mut running, &mut kv_manager, &config, 0.0) + .unwrap(); + assert!(first.output_signals.is_empty()); + assert!(running.iter().all(|req| req.output_len() == 0)); + assert_eq!(kv_manager.cache().prefix_match_len(&short_prompt), 4); + assert_eq!(kv_manager.cache().prefix_match_len(&long_prompt), 8); + kv_manager.flush_events(); + let stored_blocks: Vec<_> = buffer + .drain() + .iter() + .map(|event| match &event.data { + KvEventData::Stored(stored) => stored.blocks.len(), + KvEventData::Removed { .. } => panic!("nothing is evicted"), + }) + .collect(); + assert_eq!(stored_blocks, vec![1, 2]); + + let decode = simulate_decode_step(&mut running, &mut kv_manager, &config, 0.0, false); + let [retracted] = decode.requests.as_slice() else { + panic!("expected one retraction"); + }; + assert_eq!(retracted.uuid, Uuid::from_u128(90_101)); + assert_eq!(running[0].output_len(), 1); + assert_eq!(kv_manager.cache().prefix_match_len(&short_prompt), 4); + } + + /// SGLang aborts a running request by finishing it (`to_finish`), and the finish + /// (`release_kv_cache` with insertion) caches the tokens it decoded. + #[test] + fn cancelled_running_request_caches_its_decoded_tokens() { + let mut core = SglangCore::new_with_kv_capture(test_args(64, 4, 64), 0); + let uuid = Uuid::from_u128(90_201); + let mut request = direct_request((1..=8).collect(), 50); + request.uuid = Some(uuid); + core.receive(request); + + // One prefill pass and seven decode passes: eight output tokens, two full decoded pages. + let mut collector = crate::engine::trace::TraceCollector::default(); + let mut now_ms = 0.0; + for _ in 0..8 { + now_ms = core.execute_pass(&mut collector, now_ms).end_ms; + } + let sequence = core.running[0].sequence_prefix(16).to_vec(); + assert_eq!(core.kv_manager.cache().prefix_match_len(&sequence), 8); + core.drain_kv_events(); + + core.apply_command(SchedulerCommand::CancelRequest { request_id: uuid }) + .unwrap(); + assert!(core.is_empty()); + assert_eq!(core.kv_manager.cache().prefix_match_len(&sequence), 16); + assert_eq!(stored_hashes(&core.drain_kv_events()).len(), 2); } #[test]