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

[mem_cache] Share one ReqKvInfo between a streaming session slot and its request (#37108)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-30 16:36:48 -07:00
committed by GitHub
parent 4bea51d885
commit 4bb8de34cc
4 changed files with 35 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+5
View File
@@ -1282,6 +1282,11 @@ class Req(ReqDllmMixin):
or self.mamba_host_hit_length > 0
)
def detach_kv(self) -> ReqKvInfo:
# Hand the KV record to a new holder; the req keeps a fresh empty one.
kv, self.kv = self.kv, ReqKvInfo()
return kv
def effective_kv_committed_len(self) -> int:
# Report only the prompt prefix so thinking + answer fall into the
# overallocated range and are reclaimed by release_kv_cache. #22373.
python/sglang/srt/session/streaming_session.py
+18 -28
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import copy
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, Optional
@@ -64,17 +63,16 @@ class SessionSlot:
def save_from_req(self, req: Req, is_first: bool):
"""Save KV state from a finishing request into this slot."""
kv = req.detach_kv()
if is_first:
self.last_node = req.last_node
self.swa_uuid_for_lock = req.swa_uuid_for_lock
self.skip_lock_node_ids = req.skip_lock_node_ids
# The slot takes over the request's KV record.
self.kv = kv
else:
# The protected prefix is the first request's tree lock; nothing hands
# KV to the tree after that, so later turns must not have moved it.
assert req.kv.cache_protected_len == self.kv.cache_protected_len
# Transfer the ownership of this kv row
self.kv = copy.copy(req.kv)
# Later turns run on the slot's record (see restore_to_req).
assert kv is self.kv
self.mamba_pool_idx = req.mamba_pool_idx
self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer
@@ -83,9 +81,8 @@ class SessionSlot:
self.mamba_last_track_seqlen = req.mamba_last_track_seqlen
self.mamba_branching_seqlen = req.mamba_branching_seqlen
# Ownership moved to the slot; clear the req's references so a later
# alloc/retract path cannot mistake slot-owned mamba state for its own.
req.kv = ReqKvInfo()
# The mamba state moved to the slot too; clear the req's references so a
# later alloc/retract path cannot mistake slot-owned state for its own.
req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None
req.mamba_next_track_idx = None
@@ -95,7 +92,7 @@ class SessionSlot:
def restore_to_req(self, req: Req):
"""Restore KV state from this slot into an incoming request."""
req.kv = copy.copy(self.kv)
req.kv = self.kv
req.swa_uuid_for_lock = self.swa_uuid_for_lock
req.skip_lock_node_ids = self.skip_lock_node_ids
@@ -270,14 +267,13 @@ class StreamingSession(BasePrefixCache):
if is_npu() and self.page_size > 1:
prefix_len = (prefix_len // self.page_size) * self.page_size
req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len)
slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len)
# Free orphaned tail: alloc_for_extend will overwrite
# req_to_token[prefix_len:] with new indices. The range
# [prefix_len, kv_allocated_len) has stale indices from the
# previous turn's decode (e.g. alloc-commit gap on retract,
# or speculative draft tokens).
self._free_tail(slot, req, prefix_len)
self._free_tail(req.kv, prefix_len)
device_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, :prefix_len
@@ -311,6 +307,7 @@ class StreamingSession(BasePrefixCache):
# in req_nodes (finish_req was never called -> last successful
# req). Next request re-prefills from scratch.
if isinstance(req.finished_reason, FINISH_ABORT):
kv = req.detach_kv()
if slot is None:
# First-request mid-processing abort: create ephemeral
# slot from req state so release_session handles cleanup.
@@ -320,7 +317,7 @@ class StreamingSession(BasePrefixCache):
# return the (possibly extra_buffer ping-pong) slots to
# the mamba pool; otherwise the abort orphans them.
slot = SessionSlot(
kv=copy.copy(req.kv),
kv=kv,
last_node=req.last_node,
swa_uuid_for_lock=req.swa_uuid_for_lock,
skip_lock_node_ids=req.skip_lock_node_ids,
@@ -332,11 +329,9 @@ class StreamingSession(BasePrefixCache):
# the abort fall-through doesn't double-free.
req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None
slot.kv.kv_allocated_len = max(
slot.kv.kv_allocated_len, req.kv.kv_allocated_len
)
else:
assert kv is slot.kv
self.release_session(session_id)
req.kv = ReqKvInfo()
req.session.abort_req()
return True
@@ -542,21 +537,16 @@ class StreamingSession(BasePrefixCache):
# -- Internal helpers (streaming body bits) --
def _free_tail(self, slot: SessionSlot, req: Req, prefix_len: int) -> None:
def _free_tail(self, kv: ReqKvInfo, prefix_len: int) -> None:
"""match_prefix path: free orphaned KV in [prefix_len, kv_allocated_len)
before alloc_for_extend overwrites it. The gap appears when spec
decoding pushes allocated above committed, or when retract retry's
logit-reserve pulls prefix_len below committed.
"""
self._free_kv_aligned(
slot.kv.req_pool_idx, prefix_len, slot.kv.kv_allocated_len
)
slot.kv.kv_allocated_len = prefix_len
slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len)
slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len)
req.kv.kv_allocated_len = prefix_len
req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len)
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, prefix_len)
self._free_kv_aligned(kv.req_pool_idx, prefix_len, kv.kv_allocated_len)
kv.kv_allocated_len = prefix_len
kv.kv_committed_len = min(kv.kv_committed_len, prefix_len)
kv.swa_evicted_seqlen = min(kv.swa_evicted_seqlen, prefix_len)
def _trim_overshoot(self, req: Req, finished_len: int) -> None:
"""Trim slot KV to finished_len boundary. Spec v2 may overshoot