diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5e48302d3..18050ac9c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1420,7 +1420,15 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): ), "req_pool_indices is full! There is a bug in memory estimation." fill_len = self._pre_alloc_fill_len(req) - req.kv.kv_allocated_len = fill_len + # TODO(th4): co-locate this req.kv bookkeeping with the real KV + # allocation; the pool alloc above and the kv_allocated_len assignment + # below should become a single owned-kv allocation step. + if req.kv is None: + from sglang.srt.managers.schedule_batch import ReqKvInfo + + req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0) + else: + req.kv.kv_allocated_len = fill_len req.kv_committed_len = fill_len if prefix_len > 0: diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index f0e86361d..4ab5e9ff3 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -286,6 +286,7 @@ class DecodeKVCacheOffloadManager: self.token_to_kv_pool_allocator.free(overalloc_indices) self.req_to_token_pool.free(req) + req.kv = None self.tree_cache.protected_size_ -= len(req.prefix_indices) if req.rid in self.offloaded_state: del self.offloaded_state[req.rid] diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 071d9d50c..7177b5787 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -748,7 +748,7 @@ class Req(ReqDllmMixin): # For req-level memory management self.kv_committed_len = 0 - self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) + self.kv: Optional[ReqKvInfo] = None # for cross-encoder model self.token_type_ids = token_type_ids @@ -1501,9 +1501,8 @@ class Req(ReqDllmMixin): self.mamba_cow_src_index = None self.mamba_needs_clear = False self.already_computed = 0 - self.kv.kv_allocated_len = 0 + assert self.kv is None, "expect it is already released" self.kv_committed_len = 0 - self.kv.swa_evicted_seqlen = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 @@ -2170,8 +2169,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): req.extend_batch_idx += 1 # update req-level memory management fields + # TODO(th4): co-locate this req.kv bookkeeping with the real KV + # allocation in alloc_for_extend above; they are currently a few + # steps apart and should become one owned-kv allocation step. req.kv_committed_len = seq_len - req.kv.kv_allocated_len = seq_len + if req.kv is None: + req.kv = ReqKvInfo(kv_allocated_len=seq_len, swa_evicted_seqlen=0) + else: + req.kv.kv_allocated_len = seq_len # If input_embeds are available, store them if req.input_embeds is not None: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 0b533a84d..9722699bf 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -315,6 +315,8 @@ class SchedulerInvariantChecker: batch = self.get_last_batch() if batch is not None: for req in batch.reqs: + if req.kv is None: + continue _add_owner( req, f"req {req.rid}", diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index ff440de77..479b2322f 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -721,6 +721,7 @@ class SchedulerPPMixin: ] self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free(req) + req.kv = None logger.info( f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: " diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 9e4f0924b..28cf9520b 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -70,6 +70,9 @@ def free_swa_out_of_window_slots( token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, is_chunk_cache: bool = False, ) -> None: + if req.kv is None: + return + # For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window assert ( req.cache_protected_len % page_size == 0 @@ -683,6 +686,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr tree_cache.req_to_token_pool.free_mamba_cache(req) # DSV4-NPU's free() also releases c4/c128 state pages; no-op for others. tree_cache.req_to_token_pool.free(req) + req.kv = None def available_and_evictable_str(tree_cache: BasePrefixCache) -> str: diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index bc819cbca..946aa438e 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -37,12 +37,6 @@ class _VirtualNode: pass -def _new_kv() -> ReqKvInfo: - from sglang.srt.managers.schedule_batch import ReqKvInfo - - return ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) - - @dataclass class SessionSlot: """Holds KV state between streaming session turns.""" @@ -52,7 +46,7 @@ class SessionSlot: # KV pool state (None means no KV is currently held by this slot) req_pool_idx: Optional[int] = None kv_committed_len: int = 0 - kv: ReqKvInfo = field(default_factory=_new_kv) + kv: Optional[ReqKvInfo] = None # First req's radix tree node (for dec_lock_ref on session close) last_node: Any = None @@ -97,6 +91,7 @@ class SessionSlot: # the slot's tensor to be reused by a new req and leaked when # the slot is later freed. req.req_pool_idx = None + req.kv = None req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None @@ -324,6 +319,7 @@ class StreamingSession(BasePrefixCache): ) self.release_session(session_id) req.req_pool_idx = None + req.kv = None req.session.abort_req() return True diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 1e6a614ff..bd5e40328 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -20,7 +20,7 @@ from sglang.srt.disaggregation.kv_events import ( StorageMedium, ) from sglang.srt.environ import envs -from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.schedule_batch import Req, ReqKvInfo from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( @@ -649,6 +649,7 @@ class UnifiedRadixCacheSuite: ) self._rid += 1 req_to_token_pool.alloc([req]) + req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) return req def _apply_match_to_req(self, req, match): @@ -883,7 +884,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len - req.kv.kv_allocated_len = kv_len + req.kv = ReqKvInfo(kv_allocated_len=kv_len, swa_evicted_seqlen=0) req.last_node = cache.root_node req.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -1769,7 +1770,7 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.kv.swa_evicted_seqlen = 0 + req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) swa_avail_before = allocator.swa_attn_allocator.available_size() @@ -1821,7 +1822,7 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.kv.swa_evicted_seqlen = 0 + req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): cache.cache_unfinished_req(req)