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]