From 0cb7295698f01e948ab757dbdf73ece49ffbf548 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 14 Apr 2026 13:11:06 -0700 Subject: [PATCH] Fix streaming session busy-check double-counting via active_pool_idxs (#22753) --- .../scheduler_runtime_checker_mixin.py | 21 ++++++++-- .../srt/mem_cache/session_aware_cache.py | 42 +++++++++++-------- 2 files changed, 42 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index d35efdb53..14edf1e67 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -141,19 +141,34 @@ class SchedulerRuntimeCheckerMixin: if session.streaming ) + def _active_pool_idxs(self: Scheduler) -> set: + """Pool idxs currently owned by reqs in last_batch / running_batch. + + Used to decide which session slots' KV is owned by batch reqs + (and thus counted via uncached_size, not session_held). + """ + idxs = set() + for batch in [self.last_batch, self.running_batch]: + if batch is None or batch.is_empty(): + continue + for req in batch.reqs: + if req.req_pool_idx is not None: + idxs.add(req.req_pool_idx) + return idxs + def _session_held_tokens(self: Scheduler) -> int: if isinstance(self.tree_cache, SessionAwareCache): - return self.tree_cache.session_held_tokens() + return self.tree_cache.session_held_tokens(self._active_pool_idxs()) return 0 def _session_held_full_tokens(self: Scheduler) -> int: if isinstance(self.tree_cache, SessionAwareCache): - return self.tree_cache.session_held_full_tokens() + return self.tree_cache.session_held_full_tokens(self._active_pool_idxs()) return 0 def _session_held_swa_tokens(self: Scheduler) -> int: if isinstance(self.tree_cache, SessionAwareCache): - return self.tree_cache.session_held_swa_tokens() + return self.tree_cache.session_held_swa_tokens(self._active_pool_idxs()) return 0 def _session_held_req_count(self: Scheduler) -> int: diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index c453d0ebc..134bf99d1 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -62,11 +62,6 @@ class SessionSlot: mamba_last_track_seqlen: Any = None mamba_branching_seqlen: Any = None - # True while the slot's KV has been restored to an active request. - # Prevents double-counting in token accounting (the request's tokens - # are already tracked via uncached_size in the busy mem check). - is_active: bool = False - @property def is_holding_kv(self) -> bool: """Whether this slot currently holds KV pool resources.""" @@ -74,7 +69,6 @@ class SessionSlot: def save_from_req(self, req: Req, is_first: bool): """Save KV state from a finishing request into this slot.""" - self.is_active = False self.req_pool_idx = req.req_pool_idx self.kv_committed_len = req.kv_committed_len self.kv_allocated_len = req.kv_allocated_len @@ -108,8 +102,6 @@ class SessionSlot: req.mamba_last_track_seqlen = self.mamba_last_track_seqlen req.mamba_branching_seqlen = self.mamba_branching_seqlen - self.is_active = True - # NOTE: req_pool_idx and mamba_pool_idx are intentionally NOT cleared # from the slot. During chunked prefill, a request may be rejected by # the scheduler (e.g. budget exhausted) and retried in the next cycle. @@ -348,37 +340,51 @@ class SessionAwareCache(BasePrefixCache): self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free_slots.append(slot.req_pool_idx) - def session_held_tokens(self) -> int: + def session_held_tokens(self, active_pool_idxs: Optional[set] = None) -> int: """Total KV tokens held by session slots, not tracked by the tree. - Excludes active slots whose tokens are already counted as part of - the running request's uncached_size in the busy mem check. + Excludes slots whose KV is currently owned by an owning request — + those tokens are counted via uncached_size in the busy mem check. + A slot's pool_idx being in active_pool_idxs indicates a req owns it. """ total = 0 for slot in self.slots.values(): - if slot.is_holding_kv and not slot.is_active: + in_batch = ( + active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs + ) + if slot.is_holding_kv and not in_batch: allocated = ceil_align(slot.kv_allocated_len, self.page_size) total += allocated - slot.cache_protected_len return total - def session_held_full_tokens(self) -> int: + def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int: """An alias to align the naming style of SWA""" - return self.session_held_tokens() + return self.session_held_tokens(active_pool_idxs) - def session_held_swa_tokens(self) -> int: + def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int: """Total SWA tokens held by session slots, not tracked by the tree.""" total = 0 for slot in self.slots.values(): - if slot.is_holding_kv and not slot.is_active: + in_batch = ( + active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs + ) + if slot.is_holding_kv and not in_batch: allocated = ceil_align(slot.kv_allocated_len, self.page_size) total += allocated - max( slot.cache_protected_len, slot.swa_evicted_seqlen ) return total - def session_held_req_count(self) -> int: + def session_held_req_count(self, active_pool_idxs: Optional[set] = None) -> int: """Number of req pool slots held by session slots.""" - return sum(s.is_holding_kv and not s.is_active for s in self.slots.values()) + + def _owned(s): + in_batch = ( + active_pool_idxs is not None and s.req_pool_idx in active_pool_idxs + ) + return s.is_holding_kv and not in_batch + + return sum(_owned(s) for s in self.slots.values()) # -- Pass-through methods --