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

Delete dead rematch path in SessionAwareCache.release_session (#22735)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-13 17:02:40 -07:00
committed by GitHub
parent 9fb00ede15
commit 33a3ba256f
3 changed files with 10 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/session_controller.py
+1 -3
View File
@@ -329,9 +329,7 @@ class SessionController:
node.req.multimodal_inputs = None node.req.multimodal_inputs = None
if isinstance(self.tree_cache, SessionAwareCache): if isinstance(self.tree_cache, SessionAwareCache):
self.tree_cache.release_session( self.tree_cache.release_session(session_id)
session_id, req if session.streaming else None
)
del self.sessions[session_id] del self.sessions[session_id]
log_info_on_rank0( log_info_on_rank0(
logger, f"Session closed: {session_id} (active={len(self.sessions)})" logger, f"Session closed: {session_id} (active={len(self.sessions)})"
python/sglang/srt/mem_cache/session_aware_cache.py
+9 -70
View File
@@ -308,81 +308,20 @@ class SessionAwareCache(BasePrefixCache):
# -- Session lifecycle -- # -- Session lifecycle --
def _resolve_release_state( def release_session(self, session_id: str):
self, slot: SessionSlot, req: Optional[Req] """Release all KV resources held by a streaming session.
) -> tuple[int, Any]:
"""Resolve the currently tree-owned prefix for a session slot.
A long-lived session can outlive radix-tree splits caused by unrelated `slot.last_node` + `slot.cache_protected_len` are trusted directly: radix
traffic. In that case, the saved `last_node` may no longer represent the tree splits mutate TreeNode objects in place (see RadixCache._split_node),
full protected prefix even though the slot's req_to_token row still so a saved TreeNode reference remains valid and the locked prefix length
contains tree-owned indices at the front. Re-match the current request is unchanged. No rematch needed -- and `match_prefix` here would cause
text, then intersect the returned tree indices with the slot's row so further splits that disturb accounting.
release uses the prefix that is still actually backed by the tree.
""" """
protected_len = slot.cache_protected_len
lock_node = slot.last_node
# TODO: re-match logic disabled — match_prefix has side effects
# (splits) that disturb tree accounting. Directly using
# slot.last_node + cache_protected_len is safe after split analysis.
return protected_len, lock_node
if (
req is None
or not slot.is_holding_kv
or slot.req_pool_idx is None
or protected_len <= 0
):
return protected_len, lock_node
from sglang.srt.mem_cache.radix_cache import RadixKey
token_ids = (req.origin_input_ids + req.output_ids)[: slot.kv_committed_len]
if not token_ids:
return 0, None
match = self.inner.match_prefix(
MatchPrefixParams(
key=RadixKey(token_ids=token_ids, extra_key=req.extra_key),
req=None,
)
)
if len(match.device_indices) == 0:
return 0, None
max_protected_len = min(len(match.device_indices), protected_len)
row_indices = self.req_to_token_pool.req_to_token[
slot.req_pool_idx, :max_protected_len
].to(dtype=torch.int64)
match_indices = match.device_indices[:max_protected_len]
mismatches = (match_indices != row_indices).nonzero(as_tuple=False)
if mismatches.numel() == 0 and max_protected_len == len(match.device_indices):
common_len = max_protected_len
return common_len, match.last_device_node
common_len = (
int(mismatches[0].item()) if mismatches.numel() > 0 else max_protected_len
)
if self.page_size > 1:
common_len = (common_len // self.page_size) * self.page_size
if common_len <= 0:
return 0, None
rematch = self.inner.match_prefix(
MatchPrefixParams(
key=RadixKey(token_ids=token_ids[:common_len], extra_key=req.extra_key),
req=None,
)
)
return len(rematch.device_indices), rematch.last_device_node
def release_session(self, session_id: str, req: Optional[Req] = None):
"""Release all KV resources held by a streaming session."""
slot = self.slots.pop(session_id, None) slot = self.slots.pop(session_id, None)
if slot is None: if slot is None:
return return
protected_len, lock_node = self._resolve_release_state(slot, req) protected_len = slot.cache_protected_len
lock_node = slot.last_node
tokens_freed = ( tokens_freed = (
max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0 max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0
) )
test/registered/unit/mem_cache/test_streaming_session_unit.py
-82
View File
@@ -111,88 +111,6 @@ def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch):
assert allocator.freed[0].tolist() == list(range(32, 40)) assert allocator.freed[0].tolist() == list(range(32, 40))
def test_release_session_recomputes_current_tree_owned_prefix():
page_size = 16
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
full_match = MatchResult(
device_indices=torch.tensor(list(range(16)) + list(range(64, 96))),
last_device_node="stale-expanded",
last_host_node="stale-expanded",
)
protected_match = MatchResult(
device_indices=torch.tensor(list(range(16))),
last_device_node="current-protected",
last_host_node="current-protected",
)
inner = _FakeInnerCache(
req_to_token_pool,
allocator,
page_size,
match_results=[full_match, protected_match],
)
tree_cache = SessionAwareCache(inner)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
last_node="outdated-node",
cache_protected_len=32,
)
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
tree_cache.release_session("session-a", req)
assert inner.dec_lock_ref_calls == ["current-protected"]
assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(16, 48))
def test_release_session_never_grows_tree_owned_prefix():
page_size = 16
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
overmatched = MatchResult(
device_indices=torch.tensor(list(range(48))),
last_device_node="overmatched-node",
last_host_node="overmatched-node",
)
capped_match = MatchResult(
device_indices=torch.tensor(list(range(16))),
last_device_node="original-lock-node",
last_host_node="original-lock-node",
)
inner = _FakeInnerCache(
req_to_token_pool,
allocator,
page_size,
match_results=[overmatched, capped_match],
)
tree_cache = SessionAwareCache(inner)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
last_node="outdated-node",
cache_protected_len=16,
)
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
tree_cache.release_session("session-a", req)
assert inner.dec_lock_ref_calls == ["original-lock-node"]
assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(16, 48))
def test_match_prefix_abort_does_not_restore_live_session_slot(): def test_match_prefix_abort_does_not_restore_live_session_slot():
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128) req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])