""" 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-1-gpu-large") # --------------------------------------------------------------------------- # 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()