From b1246c50f8873bc1bd5fd4ad57dcabd707692795 Mon Sep 17 00:00:00 2001 From: Leon Gao Date: Fri, 13 Mar 2026 02:36:55 -0700 Subject: [PATCH] Fix chunked prefill and KV cache leaks for streaming sessions (#20476) Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: hnyls2002 --- python/sglang/srt/mem_cache/memory_pool.py | 10 +- .../srt/mem_cache/session_aware_cache.py | 23 ++- .../sessions/test_streaming_session_leak.py | 173 ++++++++++++++++++ 3 files changed, 198 insertions(+), 8 deletions(-) create mode 100644 test/registered/sessions/test_streaming_session_leak.py diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 476b3248d..40fe210df 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -156,10 +156,12 @@ class ReqToTokenPool: # Indices of reqs that already have a req_pool_idx and will reuse # their existing slot (e.g. chunked prefill continuing across chunks). reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] - if not any(r.is_dllm() for r in reqs): - assert ( - sum(1 for i in reusing if reqs[i].is_chunked > 0) <= 1 - ), "only one chunked request may reuse req_pool_idx in a batch" + # NOTE: this check is relaxed temporarily + # https://github.com/sgl-project/sglang/pull/20476 + # if not any(r.is_dllm() for r in reqs): + # assert ( + # sum(1 for i in reusing if reqs[i].is_chunked > 0) <= 1 + # ), "only one chunked request may reuse req_pool_idx in a batch" assert all( reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in reusing ), "reusing request must be chunked or have committed KV" diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index c1eca00be..873acf28d 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -94,8 +94,11 @@ class SessionSlot: req.mamba_last_track_seqlen = self.mamba_last_track_seqlen req.mamba_branching_seqlen = self.mamba_branching_seqlen - self.req_pool_idx = None - self.mamba_pool_idx = None + # 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. + # Each retry calls match_prefix -> restore_to_req again, so the slot + # must remain intact for idempotent restoration. def _is_streaming(req: Optional[Req]) -> bool: @@ -201,8 +204,20 @@ class SessionAwareCache(BasePrefixCache): slot.save_from_req(req, is_first=is_first) def cache_unfinished_req(self, req: Req, **kwargs): - if _is_streaming(req) and req.session.session_id in self.slots: - return + if _is_streaming(req): + # in chunked_prefill for streaming, we skip the stash path which triggers radix. + # only the last chunk in first turn trigger a full prompt radix insert. + if kwargs.get("chunked", False): + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : len(req.fill_ids) + ] + req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) + return + if req.session.session_id in self.slots: + # Subsequent turns: slot exists, skip inner entirely. + return + # First turn (no slot): fall through to inner for lock management, + # tree insertion, and cache_protected_len updates between chunks. self.inner.cache_unfinished_req(req, **kwargs) def evict(self, params: EvictParams) -> EvictResult: diff --git a/test/registered/sessions/test_streaming_session_leak.py b/test/registered/sessions/test_streaming_session_leak.py new file mode 100644 index 000000000..65c4517a9 --- /dev/null +++ b/test/registered/sessions/test_streaming_session_leak.py @@ -0,0 +1,173 @@ +""" +Test for token leak in streaming sessions with chunked prefill. + +Runs concurrent multi-turn streaming sessions interleaved with non-streaming +requests (to create mixed batches), closes all sessions, waits for idle, +and checks server health. + +Usage: + python3 -m pytest test_streaming_session_leak.py -xvs +""" + +import asyncio +import time +import unittest +from typing import Any, Optional + +import aiohttp +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=120, suite="stage-b-test-large-1-gpu") + +NUM_SESSIONS = 4 +NUM_TURNS = 5 +GEN_LEN = 16 + +# Filler text to trigger chunked prefill (200+ tokens per turn) +FILLER = ( + "The quick brown fox jumps over the lazy dog. " + "Pack my box with five dozen liquor jugs. " + "How vexingly quick daft zebras jump. " + "Sphinx of black quartz, judge my vow. " + "The five boxing wizards jump quickly. " + "Jackdaws love my big sphinx of quartz. " + "A wizard's job is to vex chumps quickly in fog. " + "We promptly judged antique ivory buckles for the next prize. " +) * 20 + + +async def _async_generate( + base_url: str, + session: aiohttp.ClientSession, + input_ids: list[int], + session_params: Optional[dict[str, Any]] = None, +) -> Any: + payload: dict[str, Any] = { + "input_ids": input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": GEN_LEN, + "no_stop_trim": True, + "skip_special_tokens": False, + }, + } + if session_params: + payload["session_params"] = session_params + timeout = aiohttp.ClientTimeout(total=300) + async with session.post( + base_url + "/generate", json=payload, timeout=timeout + ) as resp: + assert resp.status == 200, f"Generate failed: {await resp.text()}" + return await resp.json() + + +async def _run_all(base_url: str, tokenizer: Any) -> None: + """Fire all requests per turn simultaneously to create mixed batches.""" + timeout = aiohttp.ClientTimeout(total=300) + async with aiohttp.ClientSession(timeout=timeout) as http: + # Open all sessions + sids = [] + for s in range(NUM_SESSIONS): + async with http.post( + base_url + "/open_session", + json={"capacity_of_str_len": 50000, "streaming": True}, + ) as resp: + sids.append(await resp.json()) + + # For each turn, fire ALL streaming + non-streaming requests at once + for turn in range(NUM_TURNS): + tasks = [] + # Streaming requests for all sessions + for s in range(NUM_SESSIONS): + offset = (s * NUM_TURNS + turn) * 200 + text = f"Session {s} turn {turn}: {FILLER[offset : offset + 1500]}" + ids = tokenizer.encode(text) + tasks.append( + _async_generate( + base_url, + http, + ids, + session_params={"id": sids[s], "rid": None}, + ) + ) + + # Non-streaming requests interleaved + for ns in range(NUM_SESSIONS // 2): + text = f"Non-streaming {ns} turn {turn}: {FILLER[ns * 100 : ns * 100 + 500]}" + ids = tokenizer.encode(text) + tasks.append(_async_generate(base_url, http, ids)) + + # Fire all at once — creates mixed batch of streaming + non-streaming + await asyncio.gather(*tasks) + + # Close all sessions + for sid in sids: + async with http.post( + base_url + "/close_session", json={"session_id": sid} + ) as resp: + assert resp.status == 200 + + +class TestStreamingSessionLeak(CustomTestCase): + @classmethod + def setUpClass(cls) -> None: + 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", + ], + ) + + @classmethod + def tearDownClass(cls) -> None: + kill_process_tree(cls.process.pid) + + def test_streaming_session_no_leak(self) -> None: + """Concurrent multi-turn streaming sessions then idle health check.""" + from sglang.srt.utils.hf_transformers_utils import get_tokenizer + + tokenizer = get_tokenizer(self.model) + requests.post(self.base_url + "/flush_cache") + + asyncio.run(_run_all(self.base_url, tokenizer)) + + # Run a few non-streaming requests to flush state + for i in range(3): + ids = tokenizer.encode(f"Flush request {i}: final cleanup.") + requests.post( + self.base_url + "/generate", + json={ + "input_ids": ids, + "sampling_params": {"temperature": 0, "max_new_tokens": 4}, + }, + ) + + # Wait for server to go idle and run memory check + time.sleep(5) + health = requests.get(self.base_url + "/health") + self.assertEqual( + health.status_code, + 200, + "Server unhealthy after streaming session close — " + "likely a token memory leak from streaming session lifecycle.", + ) + + +if __name__ == "__main__": + unittest.main()