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

[mem_cache] Move cache_protected_len and swa_evict_floor into ReqKvInfo (#36982)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-29 19:37:53 -07:00
committed by GitHub
parent ed39568e79
commit 6be767c2d2
29 changed files with 146 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/staging_handler.py
+3 -3
View File
@@ -171,12 +171,12 @@ class DecodeStagingHandler:
# exact only when the prefix is page-aligned. Fail just this request on a
# mismatch instead of raising, which would kill the prefill scheduler.
page_size = self.kv_buffer_info["page_size"]
if decode_req.req.cache_protected_len % page_size != 0:
if decode_req.req.kv.cache_protected_len % page_size != 0:
logger.error(
"[STAGING] decode prefix length %s is not page-aligned "
"(page_size=%s); failing room=%s (staging scatter offsets "
"would be wrong).",
decode_req.req.cache_protected_len,
decode_req.req.kv.cache_protected_len,
page_size,
room,
)
@@ -414,7 +414,7 @@ class DecodeStagingHandler:
req_pool_idx = decode_req.req.req_pool_idx
# page_start is suffix-relative (pages after the decode-side cached
# prefix); req_to_token rows are absolute.
prefix_tokens = decode_req.req.cache_protected_len
prefix_tokens = decode_req.req.kv.cache_protected_len
token_start = prefix_tokens + page_start * page_size
token_end = token_start + num_pages * page_size
prefill_tp = receiver.prefill_info.attn_tp_size
python/sglang/srt/disaggregation/decode.py
+1 -1
View File
@@ -1360,7 +1360,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
count_retracted=True,
extra_reserved_reqs=len(preallocated_reqs) + 1,
)
decode_req.req.cache_protected_len = total_prefix_len
decode_req.req.kv.cache_protected_len = total_prefix_len
page_size = self.token_to_kv_pool_allocator.page_size
kv_transfer_page_size = page_size
python/sglang/srt/managers/schedule_batch.py
+18 -15
View File
@@ -817,13 +817,22 @@ class ReqLogprob:
class ReqKvInfo:
# Device KV a request holds outside the prefix cache. Always present on the Req;
# whether any KV is held is `req.req_pool_idx is not None` (Req.is_holding_kv).
# The request's own KV is [cache_protected_len, kv_allocated_len).
cache_protected_len: int = 0 # tree cache owns [0, here) (matched or inserted)
kv_allocated_len: int = 0
# The length of KV that have been removed in swa cache.
# SWA KV cache eviction behavior differs by cache type:
# - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in
# `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction.
# - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`.
swa_evicted_seqlen: int = 0
# SWA slots in [swa_dead_lo(page_size), swa_evicted_seqlen) are already freed.
swa_evict_floor: int = 0 # [0, here) never window-evicted (prefill-aware SWA)
swa_evicted_seqlen: int = 0 # SWA eviction cursor
def swa_dead_lo(self, page_size: int) -> int:
# Lowest SWA position this request may free itself: above the tree-owned
# prefix and above the eviction shield, page-aligned upward.
lo = max(self.cache_protected_len, self.swa_evict_floor)
if page_size > 1 and lo > self.cache_protected_len:
lo = ceil_align(lo, page_size)
return lo
@property
def is_released(self) -> bool:
@@ -920,10 +929,6 @@ class Req(ReqDllmMixin):
# for cross-encoder model
self.token_type_ids = token_type_ids
# Tokens in [0, swa_evict_floor) are protected from SWA window eviction.
# This is used by prefill-aware SWA models such as Unlimited-OCR to keep prompt/image KV visible during decode.
self.swa_evict_floor: int = 0
# The index of the extend / decode batch
self.extend_batch_idx = 0
self.decode_batch_idx = 0
@@ -1051,8 +1056,6 @@ class Req(ReqDllmMixin):
# per-component nodes this req skipped locking (e.g. mamba on the decode
# hold, already COW'd), so their dec releases only what it took.
self.skip_lock_node_ids: dict = {}
# The prefix length that is inserted into the tree cache
self.cache_protected_len: int = 0
# Whether or not if it is chunked. It increments whenever
# it is chunked, and decrement whenever chunked request is
@@ -1436,9 +1439,9 @@ class Req(ReqDllmMixin):
match_result.mamba_branching_seqlen,
)
if match_result.cache_protected_len is not None:
self.cache_protected_len = match_result.cache_protected_len
self.kv.cache_protected_len = match_result.cache_protected_len
else:
self.cache_protected_len = len(self.prefix_indices)
self.kv.cache_protected_len = len(self.prefix_indices)
if self.is_dllm():
self._update_block_offset_for_dllm()
@@ -1721,7 +1724,7 @@ class Req(ReqDllmMixin):
self.routed_experts = None
self.indexer_topk = None
self.last_node = None
self.cache_protected_len = 0
self.kv.cache_protected_len = 0
self.num_matched_prefix_tokens = 0
self.swa_uuid_for_lock = None
self.swa_prefix_lock_released = False
python/sglang/srt/managers/schedule_policy.py
+2 -2
View File
@@ -193,7 +193,7 @@ def match_prefix_for_req(
if match_result.mamba_branching_seqlen is not None:
req.mamba_branching_seqlen = match_result.mamba_branching_seqlen
if match_result.cache_protected_len is not None:
req.cache_protected_len = match_result.cache_protected_len
req.kv.cache_protected_len = match_result.cache_protected_len
return match_result
@@ -1327,7 +1327,7 @@ class PrefillAdder:
)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
prefix_len = len(req.prefix_indices)
req.cache_protected_len = prefix_len
req.kv.cache_protected_len = prefix_len
input_tokens = self.ceil_paged_tokens(
len(req.full_untruncated_fill_ids) - len(req.prefix_indices)
python/sglang/srt/managers/scheduler.py
+1 -1
View File
@@ -3647,7 +3647,7 @@ class Scheduler(
if self.tp_worker.model_runner.prefill_aware_swa:
for req in can_run_list:
req.swa_evict_floor = req.extend_range.end
req.kv.swa_evict_floor = req.extend_range.end
# Record prefill stats for logging after forward.
new_batch.prefill_stats = PrefillStats.from_adder(
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+3 -3
View File
@@ -258,12 +258,12 @@ class SchedulerInvariantChecker:
allocated_len = req.kv.kv_allocated_len
if self.page_size > 1:
allocated_len = ceil_align(allocated_len, self.page_size)
assert req.cache_protected_len % self.page_size == 0
assert req.kv.cache_protected_len % self.page_size == 0
full_uncached += allocated_len - req.cache_protected_len
full_uncached += allocated_len - req.kv.cache_protected_len
if self.is_hybrid_swa:
swa_uncached += allocated_len - max(
req.cache_protected_len, req.kv.swa_evicted_seqlen
req.kv.cache_protected_len, req.kv.swa_evicted_seqlen
)
if req.beam_group is not None:
python/sglang/srt/mem_cache/chunk_cache.py
+4 -4
View File
@@ -82,7 +82,7 @@ class ChunkCache(BasePrefixCache):
# For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
# The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, req.cache_protected_len : kv_len_to_handle
req.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
]
self.token_to_kv_pool_allocator.free(kv_indices)
@@ -147,7 +147,7 @@ class PureSWAChunkCache(SWAChunkCache):
explicitly skip the range already freed by ``free_swa_out_of_window_slots``
(a.k.a. _evict_swa) during decode.
``req.swa_evict_floor`` shields the prompt/image KV from window eviction
``req.kv.swa_evict_floor`` shields the prompt/image KV from window eviction
only while the request is active, so that range IS released here on
finish. Distinct from the ``cache_protected_len`` prefix, which is owned
elsewhere and never freed by this path.
@@ -161,8 +161,8 @@ class PureSWAChunkCache(SWAChunkCache):
req.req_pool_idx, :kv_committed_len
]
# The cache_protected_len prefix is not this req's to free.
protected_len = req.cache_protected_len
evict_floor = req.swa_evict_floor
protected_len = req.kv.cache_protected_len
evict_floor = req.kv.swa_evict_floor
evicted_seqlen = req.kv.swa_evicted_seqlen
if evicted_seqlen > evict_floor:
parts = []
python/sglang/srt/mem_cache/common.py
+4 -5
View File
@@ -67,12 +67,11 @@ def free_swa_out_of_window_slots(
# For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
assert (
req.cache_protected_len % page_size == 0
req.kv.cache_protected_len % page_size == 0
), "cache_protected_len must be page aligned"
evict_floor = max(req.cache_protected_len, getattr(req, "swa_evict_floor", 0))
if page_size > 1 and evict_floor > req.cache_protected_len:
evict_floor = -(-evict_floor // page_size) * page_size
req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, evict_floor)
req.kv.swa_evicted_seqlen = max(
req.kv.swa_evicted_seqlen, req.kv.swa_dead_lo(page_size)
)
if is_chunk_cache:
# Chunk cache builds no radix tree, so no tombstone-leaf concern; evict
python/sglang/srt/mem_cache/mamba_radix_cache.py
+10 -10
View File
@@ -577,7 +577,7 @@ class MambaRadixCache(BasePrefixCache):
if cache_len is None:
cache_len = 0
if cache_len != len(token_ids):
cache_end_idx = max(cache_len, req.cache_protected_len)
cache_end_idx = max(cache_len, req.kv.cache_protected_len)
self.token_to_kv_pool_allocator.free_segment(
kv_indices[cache_end_idx:], start_pos=cache_end_idx
)
@@ -639,7 +639,7 @@ class MambaRadixCache(BasePrefixCache):
),
value=page_aligned_kv_indices,
mamba_value=mamba_value,
prev_prefix_len=req.cache_protected_len,
prev_prefix_len=req.kv.cache_protected_len,
)
)
mamba_exist = result.mamba_exist
@@ -648,8 +648,8 @@ class MambaRadixCache(BasePrefixCache):
self.int8_ckpt_pool.free(mamba_value)
else:
self.token_to_kv_pool_allocator.free_segment(
kv_indices[req.cache_protected_len :],
start_pos=req.cache_protected_len,
kv_indices[req.kv.cache_protected_len :],
start_pos=req.kv.cache_protected_len,
)
mamba_exist = True
@@ -753,7 +753,7 @@ class MambaRadixCache(BasePrefixCache):
),
value=page_aligned_kv_indices,
mamba_value=mamba_value_donated,
prev_prefix_len=req.cache_protected_len,
prev_prefix_len=req.kv.cache_protected_len,
chunked=chunked,
)
)
@@ -780,15 +780,15 @@ class MambaRadixCache(BasePrefixCache):
assert torch.equal(new_last_node.mamba_value, mamba_value_donated)
assert (
req.cache_protected_len <= len(new_indices) + self.page_size - 1
), f"{req.cache_protected_len=}, {len(new_indices)=}, {len(page_aligned_token_ids)=}, {mamba_exist=}"
req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1
), f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {len(page_aligned_token_ids)=}, {mamba_exist=}"
assert new_prefix_len <= len(
new_indices
), f"{new_prefix_len=}, {len(new_indices)=}"
self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
new_indices[req.cache_protected_len :],
(req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))),
new_indices[req.kv.cache_protected_len :],
)
self.dec_lock_ref(req.last_node)
@@ -799,7 +799,7 @@ class MambaRadixCache(BasePrefixCache):
req.prefix_indices = torch.cat(
[new_indices, kv_indices_orig[len(new_indices) :]]
)
req.cache_protected_len = len(new_indices)
req.kv.cache_protected_len = len(new_indices)
req.mamba_last_track_seqlen = None
req.last_node = new_last_node
python/sglang/srt/mem_cache/pure_swa_radix_cache.py
+2 -2
View File
@@ -96,8 +96,8 @@ class PureSWARadixCache(RadixCache):
).page_aligned(self.page_size)
keys_len = len(radix_key)
old_prefix_len = req.cache_protected_len
swa_evict_floor = req.swa_evict_floor
old_prefix_len = req.kv.cache_protected_len
swa_evict_floor = req.kv.swa_evict_floor
swa_evicted_seqlen = req.kv.swa_evicted_seqlen
if self.page_size > 1 and swa_evict_floor > 0:
python/sglang/srt/mem_cache/radix_cache.py
+9 -9
View File
@@ -467,10 +467,10 @@ class RadixCache(BasePrefixCache):
if self.disable:
# The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, req.cache_protected_len : kv_len_to_handle
req.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
]
self.token_to_kv_pool_allocator.free_segment(
kv_indices, start_pos=req.cache_protected_len
kv_indices, start_pos=req.kv.cache_protected_len
)
return
@@ -502,8 +502,8 @@ class RadixCache(BasePrefixCache):
self.token_to_kv_pool_allocator.free_segments(
[
(
kv_indices[req.cache_protected_len : freed_end],
req.cache_protected_len,
kv_indices[req.kv.cache_protected_len : freed_end],
req.kv.cache_protected_len,
),
(kv_indices[key_len:], key_len),
]
@@ -543,8 +543,8 @@ class RadixCache(BasePrefixCache):
new_prefix_len = result.prefix_len
self.token_to_kv_pool_allocator.free_segment(
kv_indices[req.cache_protected_len : new_prefix_len],
start_pos=req.cache_protected_len,
kv_indices[req.kv.cache_protected_len : new_prefix_len],
start_pos=req.kv.cache_protected_len,
)
# The prefix indices could be updated, reuse it
@@ -558,15 +558,15 @@ class RadixCache(BasePrefixCache):
), f"{len(new_indices)=}, {len(radix_key)=}"
self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
new_indices[req.cache_protected_len :],
(req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))),
new_indices[req.kv.cache_protected_len :],
)
# The cache_protected_len is not always equal to len(req.prefix_indices)
# since for page_size > 1, the partial part is added to req.prefix_indices, but that part of kv indices is not added to the tree.
# It should be freed in the next cache_unfinished_req and final cache_finished_req to avoid memory leak.
# So we introduce this `cache_protected_len` field to make sure the partial part can be freed correctly.
req.cache_protected_len = len(new_indices)
req.kv.cache_protected_len = len(new_indices)
self.dec_lock_ref(req.last_node)
self.inc_lock_ref(new_last_node)
python/sglang/srt/mem_cache/swa_radix_cache.py
+3 -3
View File
@@ -483,7 +483,7 @@ class SWARadixCache(BasePrefixCache):
).page_aligned(self.page_size)
page_aligned_len = len(radix_key)
values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
old_prefix_len = req.cache_protected_len
old_prefix_len = req.kv.cache_protected_len
# Radix Cache takes one ref in memory pool
# Note: the insert function already frees the overlapped kv_indices
@@ -535,7 +535,7 @@ class SWARadixCache(BasePrefixCache):
cache_salt=req.cache_salt,
).page_aligned(self.page_size)
values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True)
old_prefix_len = req.cache_protected_len
old_prefix_len = req.kv.cache_protected_len
# Radix Cache takes one ref in memory pool
# Note: the insert function already frees the overlapped kv_indices
@@ -565,7 +565,7 @@ class SWARadixCache(BasePrefixCache):
new_indices[old_prefix_len:],
)
req.cache_protected_len = len(new_indices)
req.kv.cache_protected_len = len(new_indices)
self.dec_lock_ref(
req.last_node,
python/sglang/srt/mem_cache/unified_cache/components/README.md
+2 -2
View File
@@ -242,7 +242,7 @@ Cache an in-progress request's partial KV data (chunked prefill).
| **Purpose** | During chunked prefill, insert partial results so the next chunk can match the prefix |
| **Inputs** | `req` — the in-progress request |
| **Output** | `None` |
| **Mutation** | Inserts partial KV → re-matches prefix → updates `req.prefix_indices`, `req.cache_protected_len`, `req.last_node`; transfers lock from old node to new node |
| **Mutation** | Inserts partial KV → re-matches prefix → updates `req.prefix_indices`, `req.kv.cache_protected_len`, `req.last_node`; transfers lock from old node to new node |
| **Complexity** | **O(K + D·C)** — two tree traversals: insert O(K + D·C) + re-match O(K + D·C) + lock transfer O(D). Simplifies to **O(K)**. |
**Algorithm detail:**
@@ -252,7 +252,7 @@ Cache an in-progress request's partial KV data (chunked prefill).
4. Writes new prefix indices into `req_to_token_pool`
5. `dec_lock_ref()` on old `req.last_node`
6. `inc_lock_ref()` on new matched node
7. Updates `req.prefix_indices`, `req.cache_protected_len`, `req.last_node`
7. Updates `req.prefix_indices`, `req.kv.cache_protected_len`, `req.last_node`
8. `cleanup_after_caching_req()` per component
---
python/sglang/srt/mem_cache/unified_radix_cache.py
+10 -10
View File
@@ -838,7 +838,7 @@ class UnifiedRadixCache(BasePrefixCache):
if is_insert:
insert_params = InsertParams(
prev_prefix_len=req.cache_protected_len,
prev_prefix_len=req.kv.cache_protected_len,
priority=getattr(req, "priority", 0) or 0,
)
@@ -859,7 +859,7 @@ class UnifiedRadixCache(BasePrefixCache):
kv_indices_full = kv_indices
tail_free_start = None
if effective_cache_len < len(token_ids):
tail_free_start = max(effective_cache_len, req.cache_protected_len)
tail_free_start = max(effective_cache_len, req.kv.cache_protected_len)
token_ids = token_ids[:effective_cache_len]
kv_indices = kv_indices[:effective_cache_len]
@@ -883,8 +883,8 @@ class UnifiedRadixCache(BasePrefixCache):
self.token_to_kv_pool_allocator.free_segments(segments)
else:
self.token_to_kv_pool_allocator.free_segment(
kv_indices[req.cache_protected_len :],
start_pos=req.cache_protected_len,
kv_indices[req.kv.cache_protected_len :],
start_pos=req.kv.cache_protected_len,
)
self._dec_req_lock(req, skip_swa=req.swa_prefix_lock_released)
@@ -925,7 +925,7 @@ class UnifiedRadixCache(BasePrefixCache):
# components prepare insert data + return effective cache_len
insert_params = InsertParams(
prev_prefix_len=req.cache_protected_len,
prev_prefix_len=req.kv.cache_protected_len,
chunked=chunked,
priority=getattr(req, "priority", 0) or 0,
)
@@ -981,14 +981,14 @@ class UnifiedRadixCache(BasePrefixCache):
new_last_node = match_result.last_device_node
new_prefix_len = result.prefix_len
assert (
req.cache_protected_len <= len(new_indices) + self.page_size - 1
), f"{req.cache_protected_len=}, {len(new_indices)=}, {page_aligned_len=}"
req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1
), f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {page_aligned_len=}"
assert new_prefix_len <= len(
new_indices
), f"{new_prefix_len=}, {len(new_indices)=}"
self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
new_indices[req.cache_protected_len :],
(req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))),
new_indices[req.kv.cache_protected_len :],
)
self._dec_req_lock(req)
@@ -1014,7 +1014,7 @@ class UnifiedRadixCache(BasePrefixCache):
)
else:
req.prefix_indices = new_indices
req.cache_protected_len = len(new_indices)
req.kv.cache_protected_len = len(new_indices)
req.last_node = new_last_node
req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock
# carry the skip set so this node's dec releases only what we locked
python/sglang/srt/session/streaming_session.py
+20 -20
View File
@@ -51,7 +51,6 @@ class SessionSlot:
# First req's radix tree node (for dec_lock_ref on session close)
last_node: Any = None
cache_protected_len: int = 0
swa_uuid_for_lock: Optional[str] = None
# components the first req skipped locking on last_node, so release dec
# releases only what it took (may share the node with another req).
@@ -74,13 +73,18 @@ class SessionSlot:
"""Save KV state from a finishing request into this slot."""
self.req_pool_idx = req.req_pool_idx
self.kv_committed_len = req.kv_committed_len
self.kv = copy.copy(req.kv)
if is_first:
self.last_node = req.last_node
self.cache_protected_len = req.cache_protected_len
self.swa_uuid_for_lock = req.swa_uuid_for_lock
self.skip_lock_node_ids = req.skip_lock_node_ids
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)
self.mamba_pool_idx = req.mamba_pool_idx
self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer
@@ -89,14 +93,8 @@ class SessionSlot:
self.mamba_last_track_seqlen = req.mamba_last_track_seqlen
self.mamba_branching_seqlen = req.mamba_branching_seqlen
# Ownership has transferred to the slot. Null *all* of the req's
# references so any later alloc()/free path that inspects the req
# (e.g. the alloc-skip check on `req.mamba_ping_pong_track_buffer
# is None`, or the retract cleanup) sees no dangling pointers
# into slot-owned tensors. Without this the alloc path can decide
# the req still has a ping-pong buffer and skip alloc, causing
# the slot's tensor to be reused by a new req and leaked when
# the slot is later freed.
# 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.req_pool_idx = None
req.kv = ReqKvInfo()
req.mamba_pool_idx = None
@@ -258,7 +256,10 @@ class StreamingSession(BasePrefixCache):
aligned_prefix_len = (
expected_prefix_len // self.page_size
) * self.page_size
if aligned_prefix_len < slot.cache_protected_len or aligned_prefix_len == 0:
if (
aligned_prefix_len < slot.kv.cache_protected_len
or aligned_prefix_len == 0
):
# Release KV to avoid leak and fallback to full prefill.
# req remains unassigned, so alloc_for_extend treats it as new.
self.release_session(req.session.session_id)
@@ -273,9 +274,9 @@ class StreamingSession(BasePrefixCache):
# Streaming sessions are append-only (session_controller rollback
# ensures req_nodes always points to the last successful req).
assert prefix_len >= slot.cache_protected_len, (
assert prefix_len >= slot.kv.cache_protected_len, (
f"streaming session prefix shrank: {prefix_len=} < "
f"{slot.cache_protected_len=}"
f"{slot.kv.cache_protected_len=}"
)
# Floor-align prefix_len to page boundary (NPU workaround).
@@ -300,7 +301,7 @@ class StreamingSession(BasePrefixCache):
last_device_node=slot.virtual_node,
last_host_node=slot.virtual_node,
best_match_node=slot.virtual_node,
cache_protected_len=slot.cache_protected_len,
cache_protected_len=slot.kv.cache_protected_len,
)
def try_cache_finished_req(
@@ -326,7 +327,7 @@ class StreamingSession(BasePrefixCache):
if slot is None:
# First-request mid-processing abort: create ephemeral
# slot from req state so release_session handles cleanup.
# Include last_node/cache_protected_len from the req so
# Include last_node from the req so
# release_session calls dec_lock_ref on the tree lock.
# Also carry the mamba refs over so _free_slot_mamba can
# return the (possibly extra_buffer ping-pong) slots to
@@ -335,7 +336,6 @@ class StreamingSession(BasePrefixCache):
req_pool_idx=req.req_pool_idx,
kv=copy.copy(req.kv),
last_node=req.last_node,
cache_protected_len=req.cache_protected_len,
swa_uuid_for_lock=req.swa_uuid_for_lock,
skip_lock_node_ids=req.skip_lock_node_ids,
mamba_pool_idx=req.mamba_pool_idx,
@@ -441,7 +441,7 @@ class StreamingSession(BasePrefixCache):
slot = self.slots.pop(session_id, None)
if slot is None:
return
protected_len = slot.cache_protected_len
protected_len = slot.kv.cache_protected_len
lock_node = slot.last_node
tokens_freed = (
max(0, slot.kv.kv_allocated_len - protected_len)
@@ -490,7 +490,7 @@ class StreamingSession(BasePrefixCache):
)
if slot.is_holding_kv and not in_batch:
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - slot.cache_protected_len
total += allocated - slot.kv.cache_protected_len
return total
def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
@@ -507,7 +507,7 @@ class StreamingSession(BasePrefixCache):
if slot.is_holding_kv and not in_batch:
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - max(
slot.cache_protected_len, slot.kv.swa_evicted_seqlen
slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen
)
return total