Fix chunked prefill and KV cache leaks for streaming sessions (#20476)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

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):
+        # NOTE: this check is relaxed temporarily
-            assert (
+        # https://github.com/sgl-project/sglang/pull/20476
-                sum(1 for i in reusing if reqs[i].is_chunked > 0) <= 1
+        # if not any(r.is_dllm() for r in reqs):
-            ), "only one chunked request may reuse req_pool_idx in a batch"
+        #     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"

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
+        # NOTE: req_pool_idx and mamba_pool_idx are intentionally NOT cleared
-        self.mamba_pool_idx = None
+        # 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:
+        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:

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