From f1eb4ca90c87051081df2d9cd7ee1dbbffc44d3b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 12 Apr 2026 01:48:16 -0700 Subject: [PATCH] Fix streaming session busy check double-counting; add compat CI tests (#22213) --- .../srt/mem_cache/session_aware_cache.py | 20 ++- .../sessions/test_streaming_session.py | 122 ++++++++++++++---- 2 files changed, 111 insertions(+), 31 deletions(-) diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 0298a3a5a..983395ec6 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/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 -- diff --git a/test/registered/sessions/test_streaming_session.py b/test/registered/sessions/test_streaming_session.py index ffefe30be..89c1a1fd8 100644 --- a/test/registered/sessions/test_streaming_session.py +++ b/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( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--enable-streaming-session", - "--chunked-prefill-size", - "512", - ], - ) + 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", + ], + ) 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()