minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Fix streaming session busy check double-counting; add compat CI tests (#22213)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-12 01:48:16 -07:00
committed by GitHub
parent 3f60df8012
commit f1eb4ca90c
2 changed files with 111 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/session_aware_cache.py
+16 -4
View File
@@ -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 --
test/registered/sessions/test_streaming_session.py
+95 -27
View File
@@ -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()