Fix streaming session with paged KV cache (SWA/MLA) (#20070)

Co-authored-by: Yilong Zhao <74357408+happierpig@users.noreply.github.com>
Co-authored-by: Aurick Qiao <6137920+aurickq@users.noreply.github.com>

python/sglang/srt/managers/schedule_batch.py

@@ -902,7 +902,10 @@ class Req(ReqDllmMixin):
                 match_result.host_hit_length,
                 match_result.mamba_branching_seqlen,
             )
-            self.cache_protected_len = len(self.prefix_indices)
+            if match_result.cache_protected_len is not None:
+                self.cache_protected_len = match_result.cache_protected_len
+            else:
+                self.cache_protected_len = len(self.prefix_indices)
 
             if self.is_dllm():
                 self._update_block_offset_for_dllm()

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -26,6 +26,16 @@ class SchedulerRuntimeCheckerMixin:
             return self.tree_cache.session_held_tokens()
         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 0
+
+    def _session_held_swa_tokens(self: Scheduler) -> int:
+        if isinstance(self.tree_cache, SessionAwareCache):
+            return self.tree_cache.session_held_swa_tokens()
+        return 0
+
     def _session_held_req_count(self: Scheduler) -> int:
         if isinstance(self.tree_cache, SessionAwareCache):
             return self.tree_cache.session_held_req_count()
@@ -104,11 +114,20 @@ class SchedulerRuntimeCheckerMixin:
             swa_available_size,
             swa_evictable_size,
         ) = self._get_swa_token_info()
-        session_held = self._session_held_tokens()
-        memory_leak = (full_num_used - session_held) != 0 or swa_num_used != 0
+        session_held_full = self._session_held_full_tokens()
+        session_held_swa = self._session_held_swa_tokens()
+
+        # Streaming sessions hold tree locks during idle, so tree-protected
+        # tokens must be accounted for alongside session-held tokens.
+        full_protected = self.tree_cache.full_protected_size()
+        swa_protected = self.tree_cache.swa_protected_size()
+        full_leaked = full_num_used - full_protected - session_held_full
+        swa_leaked = swa_num_used - swa_protected - session_held_swa
+        memory_leak = full_leaked != 0 or swa_leaked != 0
         token_msg = (
-            f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {self.tree_cache.full_protected_size()=}\n"
-            f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {self.tree_cache.swa_protected_size()=}, {session_held=}\n"
+            f"{full_leaked=}, {swa_leaked=}\n"
+            f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {full_protected=}, {session_held_full=}\n"
+            f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {swa_protected=}, {session_held_swa=}\n"
         )
         return memory_leak, token_msg
 

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -111,6 +111,7 @@ class MatchResult(NamedTuple):
     last_host_backup_node: Any = None
     host_hit_length: int = 0
     mamba_branching_seqlen: Optional[int] = None
+    cache_protected_len: Optional[int] = None
 
 
 class BasePrefixCache(ABC, PrefixCacheTrait):

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -12,6 +12,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
     MatchResult,
 )
+from sglang.srt.utils.common import ceil_align
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -178,6 +179,7 @@ class SessionAwareCache(BasePrefixCache):
             device_indices=device_indices,
             last_device_node=slot.virtual_node,
             last_host_node=slot.virtual_node,
+            cache_protected_len=slot.cache_protected_len,
         )
 
     def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
@@ -242,7 +244,23 @@ class SessionAwareCache(BasePrefixCache):
         total = 0
         for slot in self.slots.values():
             if slot.req_pool_idx is not None:
-                total += slot.kv_allocated_len - slot.cache_protected_len
+                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:
+        """An alias to align the naming style of SWA"""
+        return self.session_held_tokens()
+
+    def session_held_swa_tokens(self) -> int:
+        """Total SWA tokens held by session slots, not tracked by the tree."""
+        total = 0
+        for slot in self.slots.values():
+            if slot.req_pool_idx is not None:
+                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:
@@ -305,7 +323,14 @@ class SessionAwareCache(BasePrefixCache):
     def init_metrics_collector(self):
         return self.inner.init_metrics_collector()
 
-    # Forward attribute access for cache-specific methods (e.g. sanity_check,
+    def sanity_check(self):
+        # Skip inner sanity check when sessions hold tree locks, because
+        # the check asserts all nodes are unlocked during idle.
+        if any(s.req_pool_idx is not None for s in self.slots.values()):
+            return
+        self.inner.sanity_check()
+
+    # Forward attribute access for cache-specific methods (e.g.
     # sliding_window_size, all_values_flatten, etc.)
     def __getattr__(self, name):
         return getattr(self.inner, name)