Fix streaming session busy check double-counting; add compat CI tests (#22213)

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -58,6 +58,11 @@ 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."""
@@ -65,6 +70,7 @@ 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
@@ -98,6 +104,8 @@ 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.
@@ -269,10 +277,14 @@ class SessionAwareCache(BasePrefixCache):
             self.req_to_token_pool.free_slots.append(slot.req_pool_idx)
 
     def session_held_tokens(self) -> int:
-        """Total KV tokens held by session slots, not tracked by the tree."""
+        """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.
+        """
         total = 0
         for slot in self.slots.values():
-            if slot.is_holding_kv:
+            if slot.is_holding_kv and not slot.is_active:
                 allocated = ceil_align(slot.kv_allocated_len, self.page_size)
                 total += allocated - slot.cache_protected_len
         return total
@@ -285,7 +297,7 @@ class SessionAwareCache(BasePrefixCache):
         """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:
+            if slot.is_holding_kv and not slot.is_active:
                 allocated = ceil_align(slot.kv_allocated_len, self.page_size)
                 total += allocated - max(
                     slot.cache_protected_len, slot.swa_evicted_seqlen
@@ -294,7 +306,7 @@ class SessionAwareCache(BasePrefixCache):
 
     def session_held_req_count(self) -> int:
         """Number of req pool slots held by session slots."""
-        return sum(s.is_holding_kv for s in self.slots.values())
+        return sum(s.is_holding_kv and not s.is_active for s in self.slots.values())
 
     # -- Pass-through methods --
 

test/registered/sessions/test_streaming_session.py

@@ -17,6 +17,7 @@ from typing import Any, Optional
 import aiohttp
 import requests
 
+from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
 from sglang.srt.utils.hf_transformers_utils import get_tokenizer
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -205,36 +206,28 @@ async def _leak_run_all(base_url: str, tokenizer: Any) -> None:
                 assert resp.status == 200
 
 
-# ===================================================================
-# Test class
-# ===================================================================
-
-
 class TestStreamingSession(CustomTestCase):
     @classmethod
     def setUpClass(cls):
         cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
         cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
+        with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
-            cls.model,
+            cls.process = popen_launch_server(
-            cls.base_url,
+                cls.model,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                cls.base_url,
-            other_args=[
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-                "--enable-streaming-session",
+                other_args=[
-                "--chunked-prefill-size",
+                    "--enable-streaming-session",
-                "512",
+                    "--chunked-prefill-size",
-            ],
+                    "512",
-        )
+                ],
+            )
         cls.tokenizer = get_tokenizer(cls.model)
 
     @classmethod
     def tearDownClass(cls):
         kill_process_tree(cls.process.pid)
 
-    # ------------------------------------------------------------------
-    # KV cache mechanics
-    # ------------------------------------------------------------------
-
     def test_kv_cache_inheritance(self, gen_len=12):
         """Verify KV inheritance, radix cache insertion, and flush reclamation."""
         chunks = [
@@ -346,10 +339,6 @@ class TestStreamingSession(CustomTestCase):
             "After session close + flush, cache should be fully reclaimed",
         )
 
-    # ------------------------------------------------------------------
-    # Logprob leak tests
-    # ------------------------------------------------------------------
-
     def test_leak_logprob_none(self) -> None:
         """Streaming sessions without logprobs must not leak tokens."""
         _logprob_assert_no_leak(self.base_url, self.tokenizer)
@@ -367,10 +356,6 @@ class TestStreamingSession(CustomTestCase):
             logprob_start_len=0,
         )
 
-    # ------------------------------------------------------------------
-    # Chunked prefill leak test
-    # ------------------------------------------------------------------
-
     def test_leak_chunked_prefill(self) -> None:
         """Concurrent multi-turn streaming sessions then idle health check."""
         requests.post(self.base_url + "/flush_cache")
@@ -399,5 +384,88 @@ class TestStreamingSession(CustomTestCase):
         )
 
 
+class TestStreamingSessionMixedChunk(TestStreamingSession):
+    """Streaming session with --enable-mixed-chunk."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--enable-streaming-session",
+                    "--chunked-prefill-size",
+                    "512",
+                    "--enable-mixed-chunk",
+                ],
+            )
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
+@unittest.skip("streaming session + retract has a token leak — tracked separately")
+class TestStreamingSessionRetract(TestStreamingSession):
+    """Streaming session under retract decode pressure."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        with envs.SGLANG_TEST_RETRACT.override(
+            True
+        ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--enable-streaming-session",
+                    "--chunked-prefill-size",
+                    "128",
+                ],
+            )
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
+@unittest.skip("streaming session + retract has a token leak — tracked separately")
+class TestStreamingSessionRetractMixedChunk(TestStreamingSession):
+    """Streaming session under retract decode with --enable-mixed-chunk."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        with envs.SGLANG_TEST_RETRACT.override(
+            True
+        ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--enable-streaming-session",
+                    "--chunked-prefill-size",
+                    "128",
+                    "--enable-mixed-chunk",
+                ],
+            )
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
 if __name__ == "__main__":
     unittest.main()