diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c3e5a562b..d864e55bb 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1282,6 +1282,11 @@ class Req(ReqDllmMixin): or self.mamba_host_hit_length > 0 ) + def detach_kv(self) -> ReqKvInfo: + # Hand the KV record to a new holder; the req keeps a fresh empty one. + kv, self.kv = self.kv, ReqKvInfo() + return kv + def effective_kv_committed_len(self) -> int: # Report only the prompt prefix so thinking + answer fall into the # overallocated range and are reclaimed by release_kv_cache. #22373. diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index f46a62d36..e4dc8c2c4 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -1,6 +1,5 @@ from __future__ import annotations -import copy import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional @@ -64,17 +63,16 @@ class SessionSlot: def save_from_req(self, req: Req, is_first: bool): """Save KV state from a finishing request into this slot.""" + kv = req.detach_kv() if is_first: self.last_node = req.last_node self.swa_uuid_for_lock = req.swa_uuid_for_lock self.skip_lock_node_ids = req.skip_lock_node_ids + # The slot takes over the request's KV record. + self.kv = kv else: - # The protected prefix is the first request's tree lock; nothing hands - # KV to the tree after that, so later turns must not have moved it. - assert req.kv.cache_protected_len == self.kv.cache_protected_len - - # Transfer the ownership of this kv row - self.kv = copy.copy(req.kv) + # Later turns run on the slot's record (see restore_to_req). + assert kv is self.kv self.mamba_pool_idx = req.mamba_pool_idx self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer @@ -83,9 +81,8 @@ class SessionSlot: self.mamba_last_track_seqlen = req.mamba_last_track_seqlen self.mamba_branching_seqlen = req.mamba_branching_seqlen - # Ownership moved to the slot; clear the req's references so a later - # alloc/retract path cannot mistake slot-owned mamba state for its own. - req.kv = ReqKvInfo() + # The mamba state moved to the slot too; clear the req's references so a + # later alloc/retract path cannot mistake slot-owned state for its own. req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None @@ -95,7 +92,7 @@ class SessionSlot: def restore_to_req(self, req: Req): """Restore KV state from this slot into an incoming request.""" - req.kv = copy.copy(self.kv) + req.kv = self.kv req.swa_uuid_for_lock = self.swa_uuid_for_lock req.skip_lock_node_ids = self.skip_lock_node_ids @@ -270,14 +267,13 @@ class StreamingSession(BasePrefixCache): if is_npu() and self.page_size > 1: prefix_len = (prefix_len // self.page_size) * self.page_size req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len) - slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len) # Free orphaned tail: alloc_for_extend will overwrite # req_to_token[prefix_len:] with new indices. The range # [prefix_len, kv_allocated_len) has stale indices from the # previous turn's decode (e.g. alloc-commit gap on retract, # or speculative draft tokens). - self._free_tail(slot, req, prefix_len) + self._free_tail(req.kv, prefix_len) device_indices = self.req_to_token_pool.req_to_token[ req.kv.req_pool_idx, :prefix_len @@ -311,6 +307,7 @@ class StreamingSession(BasePrefixCache): # in req_nodes (finish_req was never called -> last successful # req). Next request re-prefills from scratch. if isinstance(req.finished_reason, FINISH_ABORT): + kv = req.detach_kv() if slot is None: # First-request mid-processing abort: create ephemeral # slot from req state so release_session handles cleanup. @@ -320,7 +317,7 @@ class StreamingSession(BasePrefixCache): # return the (possibly extra_buffer ping-pong) slots to # the mamba pool; otherwise the abort orphans them. slot = SessionSlot( - kv=copy.copy(req.kv), + kv=kv, last_node=req.last_node, swa_uuid_for_lock=req.swa_uuid_for_lock, skip_lock_node_ids=req.skip_lock_node_ids, @@ -332,11 +329,9 @@ class StreamingSession(BasePrefixCache): # the abort fall-through doesn't double-free. req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None - slot.kv.kv_allocated_len = max( - slot.kv.kv_allocated_len, req.kv.kv_allocated_len - ) + else: + assert kv is slot.kv self.release_session(session_id) - req.kv = ReqKvInfo() req.session.abort_req() return True @@ -542,21 +537,16 @@ class StreamingSession(BasePrefixCache): # -- Internal helpers (streaming body bits) -- - def _free_tail(self, slot: SessionSlot, req: Req, prefix_len: int) -> None: + def _free_tail(self, kv: ReqKvInfo, prefix_len: int) -> None: """match_prefix path: free orphaned KV in [prefix_len, kv_allocated_len) before alloc_for_extend overwrites it. The gap appears when spec decoding pushes allocated above committed, or when retract retry's logit-reserve pulls prefix_len below committed. """ - self._free_kv_aligned( - slot.kv.req_pool_idx, prefix_len, slot.kv.kv_allocated_len - ) - slot.kv.kv_allocated_len = prefix_len - slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len) - slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len) - req.kv.kv_allocated_len = prefix_len - req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len) - req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, prefix_len) + self._free_kv_aligned(kv.req_pool_idx, prefix_len, kv.kv_allocated_len) + kv.kv_allocated_len = prefix_len + kv.kv_committed_len = min(kv.kv_committed_len, prefix_len) + kv.swa_evicted_seqlen = min(kv.swa_evicted_seqlen, prefix_len) def _trim_overshoot(self, req: Req, finished_len: int) -> None: """Trim slot KV to finished_len boundary. Spec v2 may overshoot diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 34657afc6..ab1d63b01 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -83,6 +83,10 @@ class _FakeReq: self.skip_lock_node_ids = {} self.mamba_pool_idx = None self.mamba_ping_pong_track_buffer = None + + def detach_kv(self): + kv, self.kv = self.kv, ReqKvInfo() + return kv self.mamba_next_track_idx = None self.mamba_last_track_seqlen = None self.mamba_branching_seqlen = None @@ -177,27 +181,17 @@ def test_nth_mid_abort_nukes_session_slot(): inner = _FakeInnerCache(req_to_token_pool, allocator, page_size) tree_cache = StreamingSession(inner) - # Session already has a slot from a previous turn. - tree_cache.slots["session-a"] = SessionSlot( - kv=ReqKvInfo( - req_pool_idx=0, - kv_committed_len=50, - kv_allocated_len=50, - swa_evicted_seqlen=0, - cache_protected_len=0, - ), - last_node=None, - ) - - # Mid-processing abort: req has the SESSION slot's pool_idx (restore_to_req ran). + # Mid-processing abort: restore_to_req ran, so the req runs on the slot's + # record, which this turn has grown to committed=60 / allocated=65. req = _FakeReq("session-a", req_pool_idx=0, committed=60, allocated=65) req.finished_reason = FINISH_ABORT("client disconnected") + tree_cache.slots["session-a"] = SessionSlot(kv=req.kv, last_node=None) tree_cache.cache_finished_req(req) # Slot wiped — deleted from slots dict. assert "session-a" not in tree_cache.slots - # All KV freed: [0, 65) from release_session (slot extended to req's allocated). + # All KV freed: [0, 65) from release_session. assert len(allocator.freed) == 1 assert allocator.freed[0].tolist() == list(range(65)) # Pool slot returned. diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 4c6b7f576..58840ffd4 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -105,15 +105,14 @@ _OWNER_SITES = { "kv_allocated_len", ): 1, # streaming session tail trimming - (_SS, "StreamingSession._free_tail", "kv_committed_len"): 2, - (_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2, + (_SS, "StreamingSession._free_tail", "kv_committed_len"): 1, + (_SS, "StreamingSession._free_tail", "kv_allocated_len"): 1, (_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1, (_SS, "StreamingSession._trim_overshoot", "kv_allocated_len"): 1, - (_SS, "StreamingSession.try_cache_finished_req", "kv_allocated_len"): 1, # Inherit the authoritative finished length (not the lagging req clock). (_SS, "StreamingSession.try_cache_finished_req", "kv_committed_len"): 1, - # NPU page-boundary clamp on req and slot clocks. - (_SS, "StreamingSession.try_match_prefix", "kv_committed_len"): 2, + # NPU page-boundary clamp on the shared req/slot clock. + (_SS, "StreamingSession.try_match_prefix", "kv_committed_len"): 1, }