diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e6634ff8c..5597d8bae 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -902,7 +902,10 @@ class Req(ReqDllmMixin): match_result.host_hit_length, match_result.mamba_branching_seqlen, ) - self.cache_protected_len = len(self.prefix_indices) + if match_result.cache_protected_len is not None: + self.cache_protected_len = match_result.cache_protected_len + else: + self.cache_protected_len = len(self.prefix_indices) if self.is_dllm(): self._update_block_offset_for_dllm() diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 78eaef316..dc8b7597e 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -26,6 +26,16 @@ class SchedulerRuntimeCheckerMixin: return self.tree_cache.session_held_tokens() return 0 + def _session_held_full_tokens(self: Scheduler) -> int: + if isinstance(self.tree_cache, SessionAwareCache): + return self.tree_cache.session_held_full_tokens() + return 0 + + def _session_held_swa_tokens(self: Scheduler) -> int: + if isinstance(self.tree_cache, SessionAwareCache): + return self.tree_cache.session_held_swa_tokens() + return 0 + def _session_held_req_count(self: Scheduler) -> int: if isinstance(self.tree_cache, SessionAwareCache): return self.tree_cache.session_held_req_count() @@ -104,11 +114,20 @@ class SchedulerRuntimeCheckerMixin: swa_available_size, swa_evictable_size, ) = self._get_swa_token_info() - session_held = self._session_held_tokens() - memory_leak = (full_num_used - session_held) != 0 or swa_num_used != 0 + session_held_full = self._session_held_full_tokens() + session_held_swa = self._session_held_swa_tokens() + + # Streaming sessions hold tree locks during idle, so tree-protected + # tokens must be accounted for alongside session-held tokens. + full_protected = self.tree_cache.full_protected_size() + swa_protected = self.tree_cache.swa_protected_size() + full_leaked = full_num_used - full_protected - session_held_full + swa_leaked = swa_num_used - swa_protected - session_held_swa + memory_leak = full_leaked != 0 or swa_leaked != 0 token_msg = ( - f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {self.tree_cache.full_protected_size()=}\n" - f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {self.tree_cache.swa_protected_size()=}, {session_held=}\n" + f"{full_leaked=}, {swa_leaked=}\n" + f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {full_protected=}, {session_held_full=}\n" + f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {swa_protected=}, {session_held_swa=}\n" ) return memory_leak, token_msg diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 0d6d5313c..e5c251a83 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -111,6 +111,7 @@ class MatchResult(NamedTuple): last_host_backup_node: Any = None host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None + cache_protected_len: Optional[int] = None class BasePrefixCache(ABC, PrefixCacheTrait): diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 512e1347b..df6fd0e5b 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -12,6 +12,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -178,6 +179,7 @@ class SessionAwareCache(BasePrefixCache): device_indices=device_indices, last_device_node=slot.virtual_node, last_host_node=slot.virtual_node, + cache_protected_len=slot.cache_protected_len, ) def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs): @@ -242,7 +244,23 @@ class SessionAwareCache(BasePrefixCache): total = 0 for slot in self.slots.values(): if slot.req_pool_idx is not None: - total += slot.kv_allocated_len - slot.cache_protected_len + allocated = ceil_align(slot.kv_allocated_len, self.page_size) + total += allocated - slot.cache_protected_len + return total + + def session_held_full_tokens(self) -> int: + """An alias to align the naming style of SWA""" + return self.session_held_tokens() + + def session_held_swa_tokens(self) -> int: + """Total SWA tokens held by session slots, not tracked by the tree.""" + total = 0 + for slot in self.slots.values(): + if slot.req_pool_idx is not None: + allocated = ceil_align(slot.kv_allocated_len, self.page_size) + total += allocated - max( + slot.cache_protected_len, slot.swa_evicted_seqlen + ) return total def session_held_req_count(self) -> int: @@ -305,7 +323,14 @@ class SessionAwareCache(BasePrefixCache): def init_metrics_collector(self): return self.inner.init_metrics_collector() - # Forward attribute access for cache-specific methods (e.g. sanity_check, + def sanity_check(self): + # Skip inner sanity check when sessions hold tree locks, because + # the check asserts all nodes are unlocked during idle. + if any(s.req_pool_idx is not None for s in self.slots.values()): + return + self.inner.sanity_check() + + # Forward attribute access for cache-specific methods (e.g. # sliding_window_size, all_values_flatten, etc.) def __getattr__(self, name): return getattr(self.inner, name) diff --git a/test/registered/sessions/test_session_latency.py b/test/registered/sessions/test_session_latency.py index 021ac5e1c..774ae8a85 100644 --- a/test/registered/sessions/test_session_latency.py +++ b/test/registered/sessions/test_session_latency.py @@ -25,7 +25,6 @@ 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_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, @@ -318,19 +317,25 @@ def _print_summary(all_results: Dict[str, List[ModeResult]]): class TestSessionLatency(CustomTestCase): - @classmethod def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST + cls.model = "openai/gpt-oss-20b" cls.base_url = DEFAULT_URL_FOR_TEST + # NOTE: Overlap scheduling commits KV cache one step ahead, + # so the last decode token is cached (unlike non-overlap). + # Disable overlap to keep session cache behavior consistent. cls.process = popen_launch_server( cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=[ - "--attention-backend", - "flashinfer", + "--disable-overlap-schedule", "--enable-streaming-session", + "--mem-fraction-static", + "0.70", + "--disable-piecewise-cuda-graph", + "--page-size", + "4", ], ) cls.tokenizer = get_tokenizer(cls.model)