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

[Minor] Add SessionSlot.is_holding_kv property for readability (#20120)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2026-03-08 03:25:13 -07:00
committed by GitHub
co-authored by gemini-code-assist[bot]
parent 0f0c8b2f18
commit 29f3a5396e
1 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/session_aware_cache.py
+10 -5
View File
@@ -54,6 +54,11 @@ class SessionSlot:
mamba_last_track_seqlen: Any = None mamba_last_track_seqlen: Any = None
mamba_branching_seqlen: Any = None mamba_branching_seqlen: Any = None
@property
def is_holding_kv(self) -> bool:
"""Whether this slot currently holds KV pool resources."""
return self.req_pool_idx is not None
def save_from_req(self, req: Req, is_first: bool): def save_from_req(self, req: Req, is_first: bool):
"""Save KV state from a finishing request into this slot.""" """Save KV state from a finishing request into this slot."""
self.req_pool_idx = req.req_pool_idx self.req_pool_idx = req.req_pool_idx
@@ -229,7 +234,7 @@ class SessionAwareCache(BasePrefixCache):
else: else:
self.inner.dec_lock_ref(slot.last_node) self.inner.dec_lock_ref(slot.last_node)
if slot.req_pool_idx is not None: if slot.is_holding_kv:
start = slot.cache_protected_len start = slot.cache_protected_len
end = slot.kv_allocated_len end = slot.kv_allocated_len
if start < end: if start < end:
@@ -243,7 +248,7 @@ class SessionAwareCache(BasePrefixCache):
"""Total KV tokens held by session slots, not tracked by the tree.""" """Total KV tokens held by session slots, not tracked by the tree."""
total = 0 total = 0
for slot in self.slots.values(): for slot in self.slots.values():
if slot.req_pool_idx is not None: if slot.is_holding_kv:
allocated = ceil_align(slot.kv_allocated_len, self.page_size) allocated = ceil_align(slot.kv_allocated_len, self.page_size)
total += allocated - slot.cache_protected_len total += allocated - slot.cache_protected_len
return total return total
@@ -256,7 +261,7 @@ class SessionAwareCache(BasePrefixCache):
"""Total SWA tokens held by session slots, not tracked by the tree.""" """Total SWA tokens held by session slots, not tracked by the tree."""
total = 0 total = 0
for slot in self.slots.values(): for slot in self.slots.values():
if slot.req_pool_idx is not None: if slot.is_holding_kv:
allocated = ceil_align(slot.kv_allocated_len, self.page_size) allocated = ceil_align(slot.kv_allocated_len, self.page_size)
total += allocated - max( total += allocated - max(
slot.cache_protected_len, slot.swa_evicted_seqlen slot.cache_protected_len, slot.swa_evicted_seqlen
@@ -265,7 +270,7 @@ class SessionAwareCache(BasePrefixCache):
def session_held_req_count(self) -> int: def session_held_req_count(self) -> int:
"""Number of req pool slots held by session slots.""" """Number of req pool slots held by session slots."""
return sum(1 for s in self.slots.values() if s.req_pool_idx is not None) return sum(s.is_holding_kv for s in self.slots.values())
# -- Pass-through methods -- # -- Pass-through methods --
@@ -326,7 +331,7 @@ class SessionAwareCache(BasePrefixCache):
def sanity_check(self): def sanity_check(self):
# Skip inner sanity check when sessions hold tree locks, because # Skip inner sanity check when sessions hold tree locks, because
# the check asserts all nodes are unlocked during idle. # the check asserts all nodes are unlocked during idle.
if any(s.req_pool_idx is not None for s in self.slots.values()): if any(s.is_holding_kv for s in self.slots.values()):
return return
self.inner.sanity_check() self.inner.sanity_check()