Fix streaming session busy-check double-counting via active_pool_idxs (#22753)

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:

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 --