chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fix chunked prefill and KV cache leaks for streaming sessions (#20476)

Browse Source 
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Leon Gao
2026-03-13 02:36:55 -07:00
committed by GitHub
parent 287dc12b05
commit b1246c50f8
3 changed files with 198 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -156,10 +156,12 @@ class ReqToTokenPool:
# Indices of reqs that already have a req_pool_idx and will reuse
# their existing slot (e.g. chunked prefill continuing across chunks).
reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
if not any(r.is_dllm() for r in reqs):
assert (
sum(1 for i in reusing if reqs[i].is_chunked > 0) <= 1
), "only one chunked request may reuse req_pool_idx in a batch"
# NOTE: this check is relaxed temporarily
# https://github.com/sgl-project/sglang/pull/20476
# if not any(r.is_dllm() for r in reqs):
# assert (
# sum(1 for i in reusing if reqs[i].is_chunked > 0) <= 1
# ), "only one chunked request may reuse req_pool_idx in a batch"
assert all(
reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in reusing
), "reusing request must be chunked or have committed KV"

23
python/sglang/srt/mem_cache/session_aware_cache.py
View File

@@ -94,8 +94,11 @@ class SessionSlot:
req.mamba_last_track_seqlen = self.mamba_last_track_seqlen
req.mamba_branching_seqlen = self.mamba_branching_seqlen
self.req_pool_idx = None
self.mamba_pool_idx = None
# 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.
# Each retry calls match_prefix -> restore_to_req again, so the slot
# must remain intact for idempotent restoration.
def _is_streaming(req: Optional[Req]) -> bool:
@@ -201,8 +204,20 @@ class SessionAwareCache(BasePrefixCache):
slot.save_from_req(req, is_first=is_first)
def cache_unfinished_req(self, req: Req, **kwargs):
if _is_streaming(req) and req.session.session_id in self.slots:
return
if _is_streaming(req):
# in chunked_prefill for streaming, we skip the stash path which triggers radix.
# only the last chunk in first turn trigger a full prompt radix insert.
if kwargs.get("chunked", False):
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(req.fill_ids)
]
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
return
if req.session.session_id in self.slots:
# Subsequent turns: slot exists, skip inner entirely.
return
# First turn (no slot): fall through to inner for lock management,
# tree insertion, and cache_protected_len updates between chunks.
self.inner.cache_unfinished_req(req, **kwargs)
def evict(self, params: EvictParams) -> EvictResult: