Consolidate streaming session tests (4 files → 3, fewer server starts) (#20961)

test/registered/sessions/test_session_control.py

@@ -3,13 +3,11 @@ Usage:
 python3 -m unittest test_session_control.TestSessionControl.test_session_control
 python3 -m unittest test_session_control.TestSessionControl.test_session_control_with_branching
 python3 -m unittest test_session_control.TestSessionControl.test_session_control_backtrack_with_abort
-python3 -m unittest test_session_control.TestSessionControl.test_streaming_session
 python3 -m unittest test_session_control.TestSessionControlVision.test_session_control
 """
 
 import asyncio
 import json
-import time
 import unittest
 
 import aiohttp
@@ -45,7 +43,6 @@ class TestSessionControl(unittest.TestCase):
             other_args=[
                 "--attention-backend",
                 "triton",
-                "--enable-streaming-session",
             ],
         )
 
@@ -435,162 +432,6 @@ class TestSessionControl(unittest.TestCase):
         asyncio.run(self.run_session_control_backtrack_with_abort(replace=True))
         asyncio.run(self.run_session_control_backtrack_with_abort(replace=False))
 
-    def test_streaming_session(self, gen_len=12):
-        chunks = [
-            "Let me tell you something about France.",
-            "The capital of France is",
-            "The population of the city is",
-        ]
-        tokenizer = get_tokenizer(self.model)
-        chunks_ids = [tokenizer.encode(x) for x in chunks]
-        for i in range(1, len(chunks_ids)):
-            if chunks_ids[i][0] == tokenizer.bos_token_id:
-                chunks_ids[i] = chunks_ids[i][1:]
-
-        # === Part 1: streaming session ===
-        requests.post(self.base_url + "/flush_cache")
-        session_id = requests.post(
-            self.base_url + "/open_session",
-            json={"capacity_of_str_len": 1000, "streaming": True},
-        ).json()
-        rid = None
-        outputs_from_session = []
-
-        prev_kv_len = 0
-        for turn_idx, chunk_ids in enumerate(chunks_ids):
-            response = requests.post(
-                self.base_url + "/generate",
-                json={
-                    "input_ids": chunk_ids,
-                    "session_params": {"id": session_id, "rid": rid},
-                    "sampling_params": {
-                        "temperature": 0,
-                        "max_new_tokens": gen_len,
-                        "no_stop_trim": True,
-                        "skip_special_tokens": False,
-                    },
-                },
-            ).json()
-            rid = response["meta_info"]["id"]
-            outputs_from_session.append(response["text"])
-            cached = response["meta_info"]["cached_tokens"]
-            prompt_tokens = response["meta_info"]["prompt_tokens"]
-            completion_tokens = response["meta_info"]["completion_tokens"]
-
-            if turn_idx == 0:
-                # Turn 1 should have no cache hit (cache was flushed).
-                self.assertEqual(
-                    cached, 0, "Turn 1 should have 0 cached tokens (clean start)"
-                )
-            else:
-                # Turns 2+ inherit KV from the previous turn (via inherit_kv_states,
-                # not radix tree matching). cached_tokens reflects the inherited prefix.
-                self.assertEqual(
-                    cached,
-                    prev_kv_len,
-                    f"Turn {turn_idx + 1}: should inherit {prev_kv_len} KV tokens from previous turn",
-                )
-            prev_kv_len = prompt_tokens + completion_tokens
-
-        # Close the session before checking cache/memory state.
-        ret = requests.post(
-            self.base_url + "/close_session",
-            json={"session_id": session_id},
-        )
-        self.assertEqual(ret.status_code, 200)
-
-        # === Cache verification (after close, before flush) ===
-
-        # Assertion 2: turn 1's prompt was inserted to the cache.
-        verify_resp = requests.post(
-            self.base_url + "/generate",
-            json={
-                "input_ids": chunks_ids[0],
-                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
-            },
-        ).json()
-        self.assertGreater(
-            verify_resp["meta_info"]["cached_tokens"],
-            0,
-            "Turn 1's prompt should be cached in the radix tree",
-        )
-
-        # Assertion 3 (insertion): turn 2's prompt tokens should NOT be in cache.
-        # The tree should only contain turn 1's extent (prompt + output from
-        # cache_unfinished_req during decode). Turn 2's prompt starts fresh tokens
-        # that were never inserted.
-        verify_resp2 = requests.post(
-            self.base_url + "/generate",
-            json={
-                "input_ids": chunks_ids[1],
-                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
-            },
-        ).json()
-        self.assertEqual(
-            verify_resp2["meta_info"]["cached_tokens"],
-            0,
-            "Turn 2's prompt should not be in cache (no insertion for turns 2+)",
-        )
-
-        # === Memory verification ===
-
-        # Assertion 4 & 5: KV is released properly and no memory leak.
-        # SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE is True by default;
-        # the scheduler will crash if it detects a leak during idle.
-        time.sleep(2)
-        health_resp = requests.get(self.base_url + "/health")
-        self.assertEqual(
-            health_resp.status_code,
-            200,
-            "Server should be healthy after session close (no memory leak)",
-        )
-
-        # After flush, all cache should be reclaimed.
-        requests.post(self.base_url + "/flush_cache")
-        verify_resp3 = requests.post(
-            self.base_url + "/generate",
-            json={
-                "input_ids": chunks_ids[0],
-                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
-            },
-        ).json()
-        self.assertEqual(
-            verify_resp3["meta_info"]["cached_tokens"],
-            0,
-            "After session close + flush, cache should be fully reclaimed",
-        )
-
-        # === Part 2: non-session baseline for output comparison ===
-        requests.post(self.base_url + "/flush_cache")
-
-        outputs_normal = []
-        input_ids = chunks_ids[0][:]
-        for i in range(len(chunks_ids)):
-            response = requests.post(
-                self.base_url + "/generate",
-                json={
-                    "input_ids": input_ids,
-                    "sampling_params": {
-                        "temperature": 0,
-                        "max_new_tokens": gen_len,
-                        "no_stop_trim": True,
-                        "skip_special_tokens": False,
-                    },
-                },
-            ).json()
-            outputs_normal.append(response["text"])
-            if i + 1 < len(chunks_ids):
-                out_ids = tokenizer.encode(response["text"])
-                if out_ids and out_ids[0] == tokenizer.bos_token_id:
-                    out_ids = out_ids[1:]
-                input_ids = input_ids + out_ids + chunks_ids[i + 1]
-
-        print("outputs from streaming session:")
-        print(outputs_from_session)
-        print("outputs from normal queries:")
-        print(outputs_normal)
-        self.assertEqual(outputs_from_session, outputs_normal)
-
     def run_session_control_with_branching(
         self, root_prompt, chunks_per_step, gen_len=16
     ):

test/registered/sessions/test_session_latency.py

@@ -32,11 +32,11 @@ from sglang.test.test_utils import (
 )
 
 register_cuda_ci(
-    est_time=100,
+    est_time=122,
     suite="stage-b-test-large-1-gpu",
 )
 
-NUM_TURNS = 300
+NUM_TURNS = 150
 INPUT_LEN = 16
 GEN_LEN = 8
 NUM_CONCURRENT = 8
@@ -417,8 +417,8 @@ class TestSessionLatency(CustomTestCase):
             speedup = reg_tail / stm_tail if stm_tail > 0 else float("inf")
             self.assertGreaterEqual(
                 speedup,
-                2.0,
+                1.4,
-                f"streaming should be >=2x faster on last {TAIL_TURNS} turns "
+                f"streaming should be >=1.4x faster on last {TAIL_TURNS} turns "
                 f"(regular={reg_tail:.1f}ms, streaming={stm_tail:.1f}ms, speedup={speedup:.2f}x)",
             )
 

test/registered/sessions/test_session_logprob_leak.py

@@ -1,138 +0,0 @@
-"""
-Test for token leak in streaming sessions with return_logprob enabled.
-
-When logprob_start_len=0, init_next_round_input truncates the prefix match
-key to length 0, which bypasses the session slot's committed KV and orphans
-allocated tokens. This test verifies that sessions with logprobs enabled
-do not leak tokens.
-
-Usage:
-    python3 -m pytest test_session_logprob_leak.py -xvs
-"""
-
-import time
-import unittest
-from typing import Any
-
-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_TURNS = 5
-NUM_ROUNDS = 30
-
-PROMPTS = [
-    "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.",
-]
-
-
-def _generate(base_url, input_ids, **kwargs) -> dict:
-    payload: dict[str, Any] = {
-        "input_ids": input_ids,
-        "sampling_params": {
-            "temperature": 0,
-            "max_new_tokens": kwargs.get("max_new_tokens", 8),
-            "no_stop_trim": True,
-            "skip_special_tokens": False,
-        },
-    }
-    if kwargs.get("return_logprob"):
-        payload["return_logprob"] = True
-    if kwargs.get("logprob_start_len") is not None:
-        payload["logprob_start_len"] = kwargs["logprob_start_len"]
-    if kwargs.get("session_params"):
-        payload["session_params"] = kwargs["session_params"]
-    resp = requests.post(base_url + "/generate", json=payload, timeout=120)
-    assert resp.status_code == 200, f"Generate failed: {resp.text}"
-    return resp.json()
-
-
-def _run_one_session(base_url, tokenizer, **gen_kwargs):
-    """Open session → N turns → close."""
-    resp = requests.post(
-        base_url + "/open_session",
-        json={"capacity_of_str_len": 50000, "streaming": True},
-    )
-    assert resp.status_code == 200
-    session_id = resp.json()
-
-    rid = None
-    for turn in range(NUM_TURNS):
-        turn_ids = tokenizer.encode(f"Turn {turn}: {PROMPTS[turn % len(PROMPTS)]}")
-        result = _generate(
-            base_url,
-            turn_ids,
-            session_params={"id": session_id, "rid": rid},
-            **gen_kwargs,
-        )
-        rid = result["meta_info"]["id"]
-
-    requests.post(base_url + "/close_session", json={"session_id": session_id})
-
-
-def _assert_no_leak(base_url, tokenizer, **gen_kwargs):
-    """Run many session rounds and verify server stays healthy."""
-    requests.post(base_url + "/flush_cache")
-    for _ in range(NUM_ROUNDS):
-        _run_one_session(base_url, tokenizer, **gen_kwargs)
-    time.sleep(3)
-    assert (
-        requests.get(base_url + "/health").status_code == 200
-    ), "Server unhealthy — likely a token memory leak."
-
-
-class TestSessionLogprobLeak(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"],
-        )
-
-    @classmethod
-    def tearDownClass(cls) -> None:
-        kill_process_tree(cls.process.pid)
-
-    def _tokenizer(self):
-        from sglang.srt.utils.hf_transformers_utils import get_tokenizer
-
-        return get_tokenizer(self.model)
-
-    def test_session_without_logprob(self) -> None:
-        """Streaming sessions without logprobs must not leak tokens."""
-        _assert_no_leak(self.base_url, self._tokenizer())
-
-    def test_session_with_output_logprob(self) -> None:
-        """Streaming sessions with output logprobs must not leak tokens."""
-        _assert_no_leak(self.base_url, self._tokenizer(), return_logprob=True)
-
-    def test_session_with_input_logprob(self) -> None:
-        """Streaming sessions with logprob_start_len=0 must not leak tokens."""
-        _assert_no_leak(
-            self.base_url,
-            self._tokenizer(),
-            return_logprob=True,
-            logprob_start_len=0,
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/sessions/test_streaming_session.py

@@ -0,0 +1,403 @@
+"""
+Streaming session tests: KV cache mechanics, logprob leak, chunked prefill leak.
+
+All tests share a single server (DEFAULT_SMALL_MODEL) with streaming sessions
+and chunked prefill enabled.
+
+Usage:
+    python -m pytest test_streaming_session.py -xvs
+    python -m unittest test_streaming_session.TestStreamingSession
+"""
+
+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.srt.utils.hf_transformers_utils import get_tokenizer
+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=70, suite="stage-b-test-large-1-gpu")
+
+# ---------------------------------------------------------------------------
+# Logprob leak constants
+# ---------------------------------------------------------------------------
+
+LOGPROB_NUM_TURNS = 5
+LOGPROB_NUM_ROUNDS = 30
+
+LOGPROB_PROMPTS = [
+    "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.",
+]
+
+# ---------------------------------------------------------------------------
+# Chunked prefill leak constants
+# ---------------------------------------------------------------------------
+
+LEAK_NUM_SESSIONS = 4
+LEAK_NUM_TURNS = 5
+LEAK_GEN_LEN = 16
+
+# Filler text to trigger chunked prefill (200+ tokens per turn)
+LEAK_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
+
+
+# ---------------------------------------------------------------------------
+# Logprob leak helpers
+# ---------------------------------------------------------------------------
+
+
+def _logprob_generate(base_url, input_ids, **kwargs) -> dict:
+    payload: dict[str, Any] = {
+        "input_ids": input_ids,
+        "sampling_params": {
+            "temperature": 0,
+            "max_new_tokens": kwargs.get("max_new_tokens", 8),
+            "no_stop_trim": True,
+            "skip_special_tokens": False,
+        },
+    }
+    if kwargs.get("return_logprob"):
+        payload["return_logprob"] = True
+    if kwargs.get("logprob_start_len") is not None:
+        payload["logprob_start_len"] = kwargs["logprob_start_len"]
+    if kwargs.get("session_params"):
+        payload["session_params"] = kwargs["session_params"]
+    resp = requests.post(base_url + "/generate", json=payload, timeout=120)
+    assert resp.status_code == 200, f"Generate failed: {resp.text}"
+    return resp.json()
+
+
+def _logprob_run_one_session(base_url, tokenizer, **gen_kwargs):
+    """Open session -> N turns -> close."""
+    resp = requests.post(
+        base_url + "/open_session",
+        json={"capacity_of_str_len": 50000, "streaming": True},
+    )
+    assert resp.status_code == 200
+    session_id = resp.json()
+
+    rid = None
+    for turn in range(LOGPROB_NUM_TURNS):
+        turn_ids = tokenizer.encode(
+            f"Turn {turn}: {LOGPROB_PROMPTS[turn % len(LOGPROB_PROMPTS)]}"
+        )
+        result = _logprob_generate(
+            base_url,
+            turn_ids,
+            session_params={"id": session_id, "rid": rid},
+            **gen_kwargs,
+        )
+        rid = result["meta_info"]["id"]
+
+    requests.post(base_url + "/close_session", json={"session_id": session_id})
+
+
+def _logprob_assert_no_leak(base_url, tokenizer, **gen_kwargs):
+    """Run many session rounds and verify server stays healthy."""
+    requests.post(base_url + "/flush_cache")
+    for _ in range(LOGPROB_NUM_ROUNDS):
+        _logprob_run_one_session(base_url, tokenizer, **gen_kwargs)
+    time.sleep(3)
+    assert (
+        requests.get(base_url + "/health").status_code == 200
+    ), "Server unhealthy — likely a token memory leak."
+
+
+# ---------------------------------------------------------------------------
+# Chunked prefill leak helpers
+# ---------------------------------------------------------------------------
+
+
+async def _leak_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": LEAK_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 _leak_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(LEAK_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(LEAK_NUM_TURNS):
+            tasks = []
+            # Streaming requests for all sessions
+            for s in range(LEAK_NUM_SESSIONS):
+                offset = (s * LEAK_NUM_TURNS + turn) * 200
+                text = f"Session {s} turn {turn}: {LEAK_FILLER[offset : offset + 1500]}"
+                ids = tokenizer.encode(text)
+                tasks.append(
+                    _leak_async_generate(
+                        base_url,
+                        http,
+                        ids,
+                        session_params={"id": sids[s], "rid": None},
+                    )
+                )
+
+            # Non-streaming requests interleaved
+            for ns in range(LEAK_NUM_SESSIONS // 2):
+                text = f"Non-streaming {ns} turn {turn}: {LEAK_FILLER[ns * 100 : ns * 100 + 500]}"
+                ids = tokenizer.encode(text)
+                tasks.append(_leak_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
+
+
+# ===================================================================
+# 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",
+            ],
+        )
+        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 = [
+            "Let me tell you something about France.",
+            "The capital of France is",
+            "The population of the city is",
+        ]
+        chunks_ids = [self.tokenizer.encode(x) for x in chunks]
+        for i in range(1, len(chunks_ids)):
+            if chunks_ids[i][0] == self.tokenizer.bos_token_id:
+                chunks_ids[i] = chunks_ids[i][1:]
+
+        # === Part 1: streaming session — check KV inheritance ===
+        requests.post(self.base_url + "/flush_cache")
+        session_id = requests.post(
+            self.base_url + "/open_session",
+            json={"capacity_of_str_len": 1000, "streaming": True},
+        ).json()
+        rid = None
+
+        prev_kv_len = 0
+        for turn_idx, chunk_ids in enumerate(chunks_ids):
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": chunk_ids,
+                    "session_params": {"id": session_id, "rid": rid},
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": gen_len,
+                        "no_stop_trim": True,
+                        "skip_special_tokens": False,
+                    },
+                },
+            ).json()
+            rid = response["meta_info"]["id"]
+            cached = response["meta_info"]["cached_tokens"]
+            prompt_tokens = response["meta_info"]["prompt_tokens"]
+            completion_tokens = response["meta_info"]["completion_tokens"]
+
+            if turn_idx == 0:
+                # Turn 1 should have no cache hit (cache was flushed).
+                self.assertEqual(
+                    cached, 0, "Turn 1 should have 0 cached tokens (clean start)"
+                )
+            else:
+                # Turns 2+ inherit KV from the previous turn (via inherit_kv_states,
+                # not radix tree matching). cached_tokens reflects the inherited prefix.
+                self.assertEqual(
+                    cached,
+                    prev_kv_len,
+                    f"Turn {turn_idx + 1}: should inherit {prev_kv_len} KV tokens from previous turn",
+                )
+            prev_kv_len = prompt_tokens + completion_tokens
+
+        # Close the session before checking cache/memory state.
+        ret = requests.post(
+            self.base_url + "/close_session",
+            json={"session_id": session_id},
+        )
+        self.assertEqual(ret.status_code, 200)
+
+        # === Cache verification (after close, before flush) ===
+
+        # Turn 1's prompt was inserted to the cache.
+        verify_resp = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": chunks_ids[0],
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        ).json()
+        self.assertGreater(
+            verify_resp["meta_info"]["cached_tokens"],
+            0,
+            "Turn 1's prompt should be cached in the radix tree",
+        )
+
+        # Turn 2's prompt tokens should NOT be in cache.
+        # The tree should only contain turn 1's extent (prompt + output from
+        # cache_unfinished_req during decode). Turn 2's prompt starts fresh tokens
+        # that were never inserted.
+        verify_resp2 = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": chunks_ids[1],
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        ).json()
+        self.assertEqual(
+            verify_resp2["meta_info"]["cached_tokens"],
+            0,
+            "Turn 2's prompt should not be in cache (no insertion for turns 2+)",
+        )
+
+        # === Flush reclamation ===
+
+        requests.post(self.base_url + "/flush_cache")
+        verify_resp3 = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": chunks_ids[0],
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        ).json()
+        self.assertEqual(
+            verify_resp3["meta_info"]["cached_tokens"],
+            0,
+            "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)
+
+    def test_leak_logprob_output(self) -> None:
+        """Streaming sessions with output logprobs must not leak tokens."""
+        _logprob_assert_no_leak(self.base_url, self.tokenizer, return_logprob=True)
+
+    def test_leak_logprob_input(self) -> None:
+        """Streaming sessions with logprob_start_len=0 must not leak tokens."""
+        _logprob_assert_no_leak(
+            self.base_url,
+            self.tokenizer,
+            return_logprob=True,
+            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")
+
+        asyncio.run(_leak_run_all(self.base_url, self.tokenizer))
+
+        # Run a few non-streaming requests to flush state
+        for i in range(3):
+            ids = self.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()

test/registered/sessions/test_streaming_session_leak.py

@@ -1,173 +0,0 @@
-"""
-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()