diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 43324d495..900b3eb42 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ef200f82d..4fac93960 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f3b9e0fd1..cc832ec1d 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 2d4a497fc..460bea0be 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e83d6e4ac..d66e61077 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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( diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index fb0af2db2..767304389 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -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: diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 51b365b3d..854d35ff4 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -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 = [] diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 46c7a41cc..ee2432e3b 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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 diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index a5afb89d4..843228045 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index 938e8b5a6..256300554 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -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: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 60ec95137..6552854e0 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -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) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index acadbc352..361c2517a 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -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, diff --git a/python/sglang/srt/mem_cache/unified_cache/components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md index 5429f1a90..e317a0ce8 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -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 --- diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 8c26a8b24..225dffedf 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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 diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index c80d49a97..8077cfc23 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -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 diff --git a/test/manual/chunked_prefill/test_scripted_radix.py b/test/manual/chunked_prefill/test_scripted_radix.py index 015a4d090..c868dccf5 100644 --- a/test/manual/chunked_prefill/test_scripted_radix.py +++ b/test/manual/chunked_prefill/test_scripted_radix.py @@ -304,7 +304,7 @@ class TestRadixNoTailChunked(ScriptedTestCase): if req is not None and req.rid == r.rid: observed_mid_chunk = True prefix_len: int = len(req.prefix_indices) - protected_len: int = req.cache_protected_len + protected_len: int = req.kv.cache_protected_len assert prefix_len == protected_len, ( f"page_size=1 must take the no-tail else branch: " f"len(prefix_indices)={prefix_len} != " @@ -460,7 +460,7 @@ class TestRadixPartialPage(ScriptedTestCase): req = s.chunked_req if req is not None and req.rid == r.rid: prefix_len: int = len(req.prefix_indices) - protected_len: int = req.cache_protected_len + protected_len: int = req.kv.cache_protected_len assert prefix_len >= protected_len, ( f"len(prefix_indices)={prefix_len} dropped below " f"cache_protected_len={protected_len}: tail was freed " diff --git a/test/registered/unit/layers/test_minicpm_sparse_cache.py b/test/registered/unit/layers/test_minicpm_sparse_cache.py index ea471ab75..0cf4aeb80 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_cache.py +++ b/test/registered/unit/layers/test_minicpm_sparse_cache.py @@ -226,7 +226,7 @@ def test_streaming_session_release_frees_compressed_slots(): ) session.slots["session-a"] = SessionSlot( req_pool_idx=req_pool_idx, - kv=SimpleNamespace(kv_allocated_len=16), + kv=SimpleNamespace(kv_allocated_len=16, cache_protected_len=0), ) session.release_session("session-a") diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index f5ce8520f..2a208d037 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -97,7 +97,7 @@ class TestDecodePreallocQueuePriority(unittest.TestCase): finished_reason=FINISH_ABORT("failed") if failed else None, return_logprob=False, sampling_params=SimpleNamespace(max_new_tokens=8), - cache_protected_len=0, + kv=SimpleNamespace(cache_protected_len=0), time_stats=MagicMock(), ) return SimpleNamespace( diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 4d15cdc94..dc7c45e5c 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -12,7 +12,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() -from sglang.srt.managers.schedule_batch import NextBatchPlan, Req +from sglang.srt.managers.schedule_batch import NextBatchPlan, Req, ReqKvInfo from sglang.srt.managers.scheduler import Scheduler from sglang.srt.mem_cache.chunk_cache import ChunkCache from sglang.srt.utils.common import Range @@ -38,7 +38,7 @@ def _make_req( req.extend_range = Range(fill_len - extend_input_len, fill_len) req.inflight_middle_chunks = 0 req.host_hit_length = 0 - req.cache_protected_len = 0 + req.kv = ReqKvInfo() req.skip_radix_cache_insert = False req.last_node = None req.swa_uuid_for_lock = None diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 4216f1998..d0b792b23 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -76,14 +76,15 @@ class MockReq: ) self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else []) self.req_pool_idx = req_pool_idx - self.cache_protected_len = cache_protected_len self.last_node = last_node self.extra_key = None self.cache_salt = None self.prefix_indices = torch.empty(0, dtype=torch.int64) self.priority = 0 self.kv_committed_len = len(fill_ids) - self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids)) + self.kv = SimpleNamespace( + kv_allocated_len=len(fill_ids), cache_protected_len=cache_protected_len + ) def get_fill_ids(self): return self.full_untruncated_fill_ids[: self.extend_range.end] @@ -368,7 +369,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): req.output_ids = [99] req.last_node = object() req.finished_reason = None - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = 123 req.swa_prefix_lock_released = False req.pd_rebootstrap_in_progress = False diff --git a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py index dfe8cedb2..e2a358efd 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py @@ -22,9 +22,11 @@ class _FakeAllocator: class _FakeReq: req_pool_idx = 0 - swa_evict_floor = 3 - cache_protected_len = 0 - kv = SimpleNamespace(swa_evicted_seqlen=6) + + def __init__(self): + self.kv = SimpleNamespace( + swa_evicted_seqlen=6, swa_evict_floor=3, cache_protected_len=0 + ) def pop_committed_kv_cache(self): return 8 @@ -51,7 +53,7 @@ class TestPureSWAChunkCache(CustomTestCase): def test_finished_req_skips_protected_prefix(self): cache = self._make_cache() req = _FakeReq() - req.cache_protected_len = 2 + req.kv.cache_protected_len = 2 cache.cache_finished_req(req, kv_len_to_handle=8) diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 680b86149..0998f3827 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -27,6 +27,7 @@ import random import unittest import unittest.mock from array import array +from types import SimpleNamespace import torch @@ -518,7 +519,7 @@ class TestRadixCache(unittest.TestCase): cache.req_to_token_pool = ReqToTokenPool(request_indices.clone()) req = unittest.mock.Mock( req_pool_idx=0, - cache_protected_len=0, + kv=SimpleNamespace(cache_protected_len=0), extra_key=None, cache_salt=None, priority=0, diff --git a/test/registered/unit/mem_cache/test_radix_force_miss.py b/test/registered/unit/mem_cache/test_radix_force_miss.py index 4e1337428..c72e3c30d 100644 --- a/test/registered/unit/mem_cache/test_radix_force_miss.py +++ b/test/registered/unit/mem_cache/test_radix_force_miss.py @@ -12,6 +12,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") import unittest import unittest.mock from array import array +from types import SimpleNamespace import torch @@ -39,7 +40,7 @@ class _StubReq: self.host_hit_length = None self.num_matched_prefix_tokens = 0 self.mamba_branching_seqlen = None - self.cache_protected_len = None + self.kv = SimpleNamespace(cache_protected_len=None) def _compute_max_prefix_len(self, input_len): return max(input_len - 1, 0) diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 0697762e9..e002cac63 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -72,13 +72,13 @@ class _FakeReq: self.kv = SimpleNamespace( kv_allocated_len=allocated, swa_evicted_seqlen=0, + cache_protected_len=0, ) self.origin_input_ids = list(range(committed)) self.output_ids = [] self.extra_key = None self.cache_salt = None self.last_node = None - self.cache_protected_len = 0 self.swa_uuid_for_lock = None self.skip_lock_node_ids = {} self.mamba_pool_idx = None @@ -114,8 +114,9 @@ def test_preabort_detaches_session_and_preserves_slot(): tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, kv_committed_len=48, - kv=SimpleNamespace(kv_allocated_len=48, swa_evicted_seqlen=0), - cache_protected_len=16, + kv=SimpleNamespace( + kv_allocated_len=48, swa_evicted_seqlen=0, cache_protected_len=16 + ), ) req = _FakeReq("session-a", req_pool_idx=1, committed=1, allocated=1) @@ -178,9 +179,10 @@ def test_nth_mid_abort_nukes_session_slot(): tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, kv_committed_len=50, - kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0), + kv=SimpleNamespace( + kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0 + ), last_node=None, - cache_protected_len=0, ) # Mid-processing abort: req has the SESSION slot's pool_idx (restore_to_req ran). @@ -216,9 +218,10 @@ def test_release_session_threads_mamba_skip_ids(): tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, kv_committed_len=50, - kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0), + kv=SimpleNamespace( + kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0 + ), last_node=lock_node, - cache_protected_len=0, skip_lock_node_ids={ComponentType.MAMBA: {42}}, ) diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index bcc4df1c8..1ca7f063b 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -18,7 +18,7 @@ from types import SimpleNamespace import torch -from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.managers.schedule_batch import ReqKvInfo, ScheduleBatch from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import free_swa_out_of_window_slots @@ -107,10 +107,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): is_holding_kv=True, origin_input_ids=token_ids, output_ids=[], - cache_protected_len=cache_protected_len, - kv=SimpleNamespace( - swa_evicted_seqlen=0, - ), + kv=ReqKvInfo(cache_protected_len=cache_protected_len), extra_key=None, cache_salt=None, last_node=tree.root_node, diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 128c9a61b..9595926e3 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -36,7 +36,7 @@ class _DummyReq: def __init__(self): self._kv_committed_len = 0 self.swa_prefix_lock_released = False - self.kv = SimpleNamespace(swa_evicted_seqlen=0) + self.kv = SimpleNamespace(swa_evicted_seqlen=0, cache_protected_len=0) def _build_swa_tree( @@ -682,7 +682,7 @@ class TestSWA(unittest.TestCase): req.last_node = tree.root_node req.swa_uuid_for_lock = None req.kv.swa_evicted_seqlen = 0 - req.cache_protected_len = 1 + req.kv.cache_protected_len = 1 # Intentionally mismatch to ensure code does not use len(prefix_indices). req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device) @@ -700,7 +700,7 @@ class TestSWA(unittest.TestCase): req, is_insert=True, kv_len_to_handle=req._kv_committed_len ) - self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len) + self.assertEqual(captured["prev_prefix_len"], req.kv.cache_protected_len) self.assertTrue(captured["is_bigram"]) self.assertEqual(captured["key_len"], len(req.origin_input_ids) - 1) @@ -720,7 +720,7 @@ class TestSWA(unittest.TestCase): req2.last_node = tree.root_node req2.swa_uuid_for_lock = None req2.kv.swa_evicted_seqlen = 0 - req2.cache_protected_len = 1 + req2.kv.cache_protected_len = 1 req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device) freed_lens = [] @@ -917,7 +917,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase): req.get_fill_ids = lambda: token_ids req.extra_key = None req.cache_salt = None - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.last_node = tree.root_node req.swa_uuid_for_lock = None req.prefix_indices = torch.empty(0, dtype=torch.int64, device=tree.device) @@ -928,7 +928,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase): # The insert itself frees nothing. self.assertEqual(allocator.swa_available_size(), swa_before) # The live leaf holds a full window, so the whole key stays matchable. - self.assertEqual(req.cache_protected_len, num_tokens) + self.assertEqual(req.kv.cache_protected_len, num_tokens) # [0, evicted) is a tombstone; only [evicted, num_tokens) counts as SWA. (first,) = tree.root_node.children.values() self.assertTrue(first.swa_tombstone) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 78f5f6092..314f5652a 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -642,7 +642,7 @@ def bench_cache_finished( len(req.prefix_indices), len(req.full_untruncated_fill_ids) ) req.last_node = node - req.cache_protected_len = matched_len + req.kv.cache_protected_len = matched_len req.kv_committed_len = len(seq) if hasattr(lr, "swa_uuid_for_lock"): req.swa_uuid_for_lock = lr.swa_uuid_for_lock diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 6be425392..db3530223 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1228,7 +1228,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None req.full_untruncated_fill_ids = array("q", input_ids + output_ids) @@ -1269,7 +1269,7 @@ class UnifiedRadixCacheSuite: req.kv_committed_len = kv_len req.kv.kv_allocated_len = kv_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None if self.cfg.has_mamba: @@ -1313,7 +1313,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.swa_prefix_lock_released = True req.extra_key = None @@ -1348,7 +1348,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None if self.cfg.has_mamba: @@ -1357,7 +1357,7 @@ class UnifiedRadixCacheSuite: cache.cache_unfinished_req(req) self.assertGreater(len(req.prefix_indices), 0) - self.assertEqual(req.cache_protected_len, len(req.prefix_indices)) + self.assertEqual(req.kv.cache_protected_len, len(req.prefix_indices)) self.assertIsNotNone(req.last_node) self.assertFalse(req.swa_prefix_lock_released) @@ -1385,7 +1385,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) req.kv_committed_len = len(tokens) req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None req.kv.swa_evicted_seqlen = evicted_len @@ -1487,7 +1487,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None req.full_untruncated_fill_ids = array("q", input_ids) @@ -1612,7 +1612,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value) req.kv_committed_len = kv_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None req.kv.swa_evicted_seqlen = 0 @@ -1639,7 +1639,7 @@ class UnifiedRadixCacheSuite: swa_value, ) ) - self.assertEqual(req.cache_protected_len, len(tokens)) + self.assertEqual(req.kv.cache_protected_len, len(tokens)) cache.dec_lock_ref( req.last_node, @@ -2199,7 +2199,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req.kv_committed_len = pre_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -2288,7 +2288,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req.kv_committed_len = pre_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -6770,7 +6770,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): req.output_ids = array("q") req.kv_committed_len = len(tokens) req.kv.kv_allocated_len = len(tokens) - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None req.mamba_last_track_seqlen = len(tokens) @@ -8097,7 +8097,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase): req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices) req.kv_committed_len = seq_len req.last_node = cache.root_node_handle() - req.cache_protected_len = 0 + req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -8113,7 +8113,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase): f"{self.cfg.sliding_window_size} window", ) self.assertEqual( - req.cache_protected_len, + req.kv.cache_protected_len, boundary, "the match after the insert must reach the leaf the insert created", ) diff --git a/test/registered/xpu/test_lmcache_radix_cache.py b/test/registered/xpu/test_lmcache_radix_cache.py index 46b0090af..354767e2a 100644 --- a/test/registered/xpu/test_lmcache_radix_cache.py +++ b/test/registered/xpu/test_lmcache_radix_cache.py @@ -72,7 +72,7 @@ def _make_req(rid, req_pool_idx, token_ids, tree): extra_key=None, cache_salt=None, last_node=tree.root_node, - cache_protected_len=0, + kv=SimpleNamespace(cache_protected_len=0), priority=0, kv_committed_freed=False, kv_committed_len=len(token_ids),