From 0438b16154d3d52ef617b3297b013676bfbf161e Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 29 Aug 2026 22:42:25 -0700 Subject: [PATCH] [mem_cache] Move `kv_committed_len` into `ReqKvInfo` (#37078) --- python/sglang/srt/beam_search/coordinator.py | 2 +- python/sglang/srt/beam_search/fork.py | 2 +- python/sglang/srt/disaggregation/decode.py | 12 ++++++---- .../npu/dsv4/dsv4_allocator.py | 2 +- python/sglang/srt/managers/schedule_batch.py | 20 +++++++++------- .../batch_result_processor.py | 12 +++++----- .../scheduler_components/invariant_checker.py | 4 ++-- .../sglang/srt/mem_cache/allocation_sizing.py | 2 +- python/sglang/srt/mem_cache/common.py | 2 +- python/sglang/srt/mem_cache/memory_pool.py | 2 +- .../storage/flexkv/flexkv_radix_cache.py | 2 +- .../storage/lmcache/lmc_radix_cache.py | 2 +- .../sglang/srt/session/streaming_session.py | 19 +++++++-------- .../sglang/srt/speculative/dflash_info_v2.py | 2 +- .../test/scripted_runtime/context/queries.py | 2 +- .../test_scripted_hybrid_swa.py | 4 ++-- .../test_scripted_regression.py | 2 +- .../test_scripted_special_case.py | 4 ++-- test/registered/unit/beam_search/test_fork.py | 9 +++---- .../test_specv2_kvcache_offloading.py | 7 +++--- .../mlx/test_attention_patching.py | 1 - .../mlx/test_max_running_requests.py | 1 - .../unit/layers/test_minicpm_sparse_cache.py | 1 - ...t_batch_result_processor_mamba_boundary.py | 2 +- ...est_batch_result_processor_spec_grammar.py | 6 ++--- .../unit/managers/test_hisparse_unit.py | 9 ++++--- .../unit/managers/test_kv_page_invariants.py | 10 ++++---- .../test_schedule_batch_prepare_for_decode.py | 3 +-- .../mem_cache/test_decode_radix_lock_ref.py | 15 ++++++------ .../unit/mem_cache/test_dllm_fdfo_kv_reuse.py | 6 +++-- .../unit/mem_cache/test_hisparse_allocator.py | 2 +- .../test_session_unified_radix_cache.py | 1 - .../mem_cache/test_streaming_session_unit.py | 24 ++++++++++++------- .../test_unified_radix_cache_bench.py | 4 ++-- .../test_unified_radix_cache_unittest.py | 22 ++++++++--------- .../spec/test_decode_bookkeeping_ownership.py | 4 +--- .../xpu/test_lmcache_radix_cache.py | 9 ++----- 37 files changed, 117 insertions(+), 116 deletions(-) diff --git a/python/sglang/srt/beam_search/coordinator.py b/python/sglang/srt/beam_search/coordinator.py index ed5280872..8b224c970 100644 --- a/python/sglang/srt/beam_search/coordinator.py +++ b/python/sglang/srt/beam_search/coordinator.py @@ -442,7 +442,7 @@ class BeamCoordinator(msgspec.Struct, kw_only=True): prefix_len=group.prompt_len, # All rows are synchronized; the leader's committed length # covers the KV computed through this step. - seq_len=group.leader.kv_committed_len, + seq_len=group.leader.kv.kv_committed_len, ) group.pending_orphans.append(StagedOrphans(tick, old_map, new_map)) # Length placeholder only; the DAG owns history and member rows diff --git a/python/sglang/srt/beam_search/fork.py b/python/sglang/srt/beam_search/fork.py index e54270705..779959b1c 100644 --- a/python/sglang/srt/beam_search/fork.py +++ b/python/sglang/srt/beam_search/fork.py @@ -93,7 +93,7 @@ def free_member_rows(group, req_to_token_pool, token_to_kv_pool_allocator) -> No # release frees this decode region a second time. slots = req_to_token_pool.req_to_token[group.all_rows, start:end] token_to_kv_pool_allocator.free(slots.flatten().unique()) - leader.kv_committed_len = start + leader.kv.kv_committed_len = start leader.kv.kv_allocated_len = start req_to_token_pool.free_rows(group.member_rows_cpu.tolist()) group.member_rows = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 4fac93960..033726dda 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -193,7 +193,7 @@ class DecodeReqToTokenPool: len(reusing) <= 1 ), "only one chunked request may reuse req_pool_idx in a batch" assert all( - reqs[i].inflight_middle_chunks > 0 or reqs[i].kv_committed_len > 0 + reqs[i].inflight_middle_chunks > 0 or reqs[i].kv.kv_committed_len > 0 for i in reusing ), "reusing request must be chunked or have committed KV" @@ -1783,7 +1783,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): ), "req_pool_indices is full! There is a bug in memory estimation." fill_len = self._pre_alloc_fill_len(req) - req.kv_committed_len = fill_len + req.kv.kv_committed_len = fill_len if prefix_len > 0: self.req_to_token_pool.write( @@ -1889,7 +1889,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): req.prefix_indices = ( prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64) ) - req.set_extend_range(total_prefix_len, req.kv_committed_len) + req.set_extend_range(total_prefix_len, req.kv.kv_committed_len) # Return the transfer destination indices: if self.scheduler.enable_hisparse: @@ -2634,8 +2634,10 @@ class SchedulerDisaggregationDecodeMixin: # Truncate fill_len to kv_committed_len so cache_unfinished_req # only sees committed KV (full array includes one uncommitted # token because init_next_round_input rebuilt it as full). - if req.kv_committed_len is not None: - req.set_extend_range(len(req.prefix_indices), req.kv_committed_len) + if req.kv.kv_committed_len is not None: + req.set_extend_range( + len(req.prefix_indices), req.kv.kv_committed_len + ) else: waiting_queue.append(req) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py index bcfb5baba..0737508de 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -465,7 +465,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if req is None or req_to_token_pool is None: return - kv_len = max(req.kv_committed_len, req.kv.kv_allocated_len) + kv_len = max(req.kv.kv_committed_len, req.kv.kv_allocated_len) req_pool_idx = req.req_pool_idx if kv_len <= 0 or req_pool_idx is None: return diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index cc832ec1d..656093296 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -820,6 +820,7 @@ class ReqKvInfo: # 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_committed_len: int = 0 # KV content committed up to here, <= kv_allocated_len kv_allocated_len: int = 0 # SWA slots in [swa_dead_lo(page_size), swa_evicted_seqlen) are already freed. @@ -922,7 +923,6 @@ class Req(ReqDllmMixin): self.multi_item_delimiter_indices = multi_item_delimiter_indices # For req-level memory management - self.kv_committed_len = 0 self.kv = ReqKvInfo() self.retraction_backup: Optional[RetractionBackup] = None @@ -1286,8 +1286,8 @@ class Req(ReqDllmMixin): # Report only the prompt prefix so thinking + answer fall into the # overallocated range and are reclaimed by release_kv_cache. #22373. if get_serving().strip_thinking_cache and self.reasoning_tokens > 0: - return min(self.kv_committed_len, len(self.origin_input_ids)) - return self.kv_committed_len + return min(self.kv.kv_committed_len, len(self.origin_input_ids)) + return self.kv.kv_committed_len def update_spec_correct_drafts_histogram(self, num_correct_drafts: int): """Record one step accepted draft count (excludes bonus token) into the histogram.""" @@ -1749,7 +1749,7 @@ class Req(ReqDllmMixin): self.mamba_needs_clear = False self.already_computed = 0 assert not self.is_holding_kv, "expect it is already released" - self.kv_committed_len = 0 + self.kv.kv_committed_len = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 @@ -1937,7 +1937,7 @@ def mamba_lazy_spec_in_window( kv_committed_len lags device seq_lens by up to one verify under overlap; the 2x window absorbs it. """ - seq_len = req.kv_committed_len + seq_len = req.kv.kv_committed_len window = 2 * max_draft_tokens return seq_len // mamba_track_interval != (seq_len + window) // mamba_track_interval @@ -2513,7 +2513,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # TODO(th4): co-locate this req.kv bookkeeping with the real KV # allocation in alloc_for_extend above; they are currently a few # steps apart and should become one owned-kv allocation step. - req.kv_committed_len = seq_len + req.kv.kv_committed_len = seq_len # If input_embeds are available, store them if req.input_embeds is not None: @@ -2902,7 +2902,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ) if self.spec_algorithm.is_none(): - new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0) + new_pages = sum( + 1 for r in requests if r.kv.kv_committed_len % page_size == 0 + ) return new_pages * page_size + num_beam_member_rows(requests) return self._new_tokens_required_next_decode_spec_v2(requests, page_size) @@ -2912,7 +2914,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): reserve = get_alloc_reserve_per_decode() total = 0 for r in requests: - x = max(0, r.kv_committed_len + reserve - r.kv.kv_allocated_len) + x = max(0, r.kv.kv_committed_len + reserve - r.kv.kv_allocated_len) cur = r.kv.kv_allocated_len nxt = cur + x total += ceil_align(nxt, page_size) - ceil_align(cur, page_size) @@ -3222,7 +3224,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Update req-level memory management fields for req in self.reqs: req.decode_batch_idx += 1 - req.kv_committed_len += 1 + req.kv.kv_committed_len += 1 # New-tensor avoids racing model_worker_batch refs queued for # overlap forward. diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 9ba9160c6..832c222fe 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -748,7 +748,7 @@ class SchedulerBatchResultProcessor: # Commit the full accepted run (drafts + bonus). num_accept_tokens = len(accept_tokens) - req.kv_committed_len += num_accept_tokens + req.kv.kv_committed_len += num_accept_tokens req.spec_verify_ct += 1 num_correct_drafts = result.num_correct_drafts_per_req_cpu[i] @@ -1126,7 +1126,7 @@ class SchedulerBatchResultProcessor: if completed_mamba_boundary and not lazy: req.mamba_last_track_idx = batch.mamba_track_buffer_indices[i] - req.mamba_last_track_seqlen = req.kv_committed_len - lookahead + req.mamba_last_track_seqlen = req.kv.kv_committed_len - lookahead elif ( req.finished() and lazy @@ -1222,7 +1222,7 @@ class SchedulerBatchResultProcessor: lazy = mamba_extra_buffer_lazy_enabled() if known_boundary: self._mamba_assert_committed_len_lookahead(req) - track_seqlen = req.kv_committed_len + track_seqlen = req.kv.kv_committed_len assert track_seqlen % mamba_track_grid(self.tree_cache.page_size) == 0 at_boundary = True else: @@ -1317,8 +1317,8 @@ class SchedulerBatchResultProcessor: f"(req {req.rid}); output_ids is empty" ) token_seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1 - assert (req.kv_committed_len - token_seq_len) in (0, 1), ( - f"mamba track boundary: kv_committed_len={req.kv_committed_len} " + assert (req.kv.kv_committed_len - token_seq_len) in (0, 1), ( + f"mamba track boundary: kv_committed_len={req.kv.kv_committed_len} " f"leads seq_len={token_seq_len} by more than one (req {req.rid}); " "overlap lookahead wider than assumed" ) @@ -1340,7 +1340,7 @@ class SchedulerBatchResultProcessor: if batch.spec_algorithm.is_none(): lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i] - committed_len = req.kv_committed_len - lookahead + committed_len = req.kv.kv_committed_len - lookahead if committed_len % interval == 0: return True, committed_len elif result.num_correct_drafts_per_req_cpu is not None: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 767304389..a3557302c 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -330,7 +330,7 @@ class SchedulerInvariantChecker: req, f"req {req.rid}", req.req_pool_idx, - req.kv_committed_len, + req.kv.kv_committed_len, req.kv.kv_allocated_len, ) sess = getattr(self.tree_cache, "slots", None) @@ -341,7 +341,7 @@ class SchedulerInvariantChecker: slot, f"slot {sid[:8]}", slot.req_pool_idx, - slot.kv_committed_len, + slot.kv.kv_committed_len, slot.kv.kv_allocated_len, ) diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py index 92f114113..bcb476f92 100644 --- a/python/sglang/srt/mem_cache/allocation_sizing.py +++ b/python/sglang/srt/mem_cache/allocation_sizing.py @@ -69,7 +69,7 @@ def page_aligned_decode_alloc_lens( cur = r.kv.kv_allocated_len nxt = max( cur, - (r.kv_committed_len + reserve + page_size - 1) // page_size * page_size, + (r.kv.kv_committed_len + reserve + page_size - 1) // page_size * page_size, ) cur_kv_lens[i] = cur nxt_kv_lens[i] = nxt diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index ee2432e3b..93d814f38 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -255,7 +255,7 @@ def _release_overallocated_kv_indices( if spec_algo is None and not get_serving().strip_thinking_cache: assert ( start_p == end_p - ), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv.kv_allocated_len=}" + ), f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" if page_size > 1: start_p = ceil_align(start_p, page_size) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index d22929db4..5191304aa 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -300,7 +300,7 @@ class ReqToTokenPool: # sum(1 for i in reusing if reqs[i].inflight_middle_chunks > 0) <= 1 # ), "only one chunked request may reuse req_pool_idx in a batch" assert all( - reqs[i].inflight_middle_chunks > 0 or reqs[i].kv_committed_len > 0 + reqs[i].inflight_middle_chunks > 0 or reqs[i].kv.kv_committed_len > 0 for i in reusing ), "reusing request must be chunked or have committed KV" diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py index a7a090cd9..e1c582eb3 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -399,7 +399,7 @@ class FlexKVRadixCache(RadixCache): topk = get_spec().speculative_eagle_topk enable_kv_committed_len = topk is None or topk == 1 if enable_kv_committed_len: - kv_committed_len = req.kv_committed_len + kv_committed_len = req.kv.kv_committed_len else: kv_committed_len = len(req.origin_input_ids) + max( len(req.output_ids) - 1, 0 diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 5a24c5d19..54f70e618 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -455,7 +455,7 @@ class LMCRadixCache(RadixCache): topk = get_spec().speculative_eagle_topk enable_kv_committed_len = topk is None or topk == 1 if enable_kv_committed_len: - kv_committed_len = req.kv_committed_len + kv_committed_len = req.kv.kv_committed_len else: kv_committed_len = len(req.origin_input_ids) + max( len(req.output_ids) - 1, 0 diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 8077cfc23..631256934 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -46,7 +46,6 @@ class SessionSlot: # KV pool state req_pool_idx: Optional[int] = None - kv_committed_len: int = 0 kv: ReqKvInfo = field(default_factory=ReqKvInfo) # First req's radix tree node (for dec_lock_ref on session close) @@ -72,7 +71,6 @@ class SessionSlot: def save_from_req(self, req: Req, is_first: bool): """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 if is_first: self.last_node = req.last_node @@ -107,7 +105,6 @@ class SessionSlot: def restore_to_req(self, req: Req): """Restore KV state from this slot into an incoming request.""" req.req_pool_idx = self.req_pool_idx - req.kv_committed_len = self.kv_committed_len req.kv = copy.copy(self.kv) req.swa_uuid_for_lock = self.swa_uuid_for_lock req.skip_lock_node_ids = self.skip_lock_node_ids @@ -252,7 +249,7 @@ class StreamingSession(BasePrefixCache): # fall back to radix cache (full prefill). Once context >= page_size, # streaming session kicks in with page-aligned KV reuse. if is_npu() and self.page_size > 1: - expected_prefix_len = min(slot.kv_committed_len, len(params.key)) + expected_prefix_len = min(slot.kv.kv_committed_len, len(params.key)) aligned_prefix_len = ( expected_prefix_len // self.page_size ) * self.page_size @@ -270,7 +267,7 @@ class StreamingSession(BasePrefixCache): # token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve # already applied). min handles retract retry where committed_len # can exceed len(token_ids) by 1. - prefix_len = min(req.kv_committed_len, len(params.key)) + prefix_len = min(req.kv.kv_committed_len, len(params.key)) # Streaming sessions are append-only (session_controller rollback # ensures req_nodes always points to the last successful req). @@ -282,8 +279,8 @@ class StreamingSession(BasePrefixCache): # Floor-align prefix_len to page boundary (NPU workaround). if is_npu() and self.page_size > 1: prefix_len = (prefix_len // self.page_size) * self.page_size - req.kv_committed_len = min(req.kv_committed_len, prefix_len) - slot.kv_committed_len = min(slot.kv_committed_len, prefix_len) + 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 @@ -370,7 +367,7 @@ class StreamingSession(BasePrefixCache): # req clock (under overlap + honest committed the clock lags the in-flight # verify by ~1, which would short-change inheritance). Clamp to allocated # to keep committed <= allocated for prepare_for_decode. - slot.kv_committed_len = min(target, slot.kv.kv_allocated_len) + slot.kv.kv_committed_len = min(target, slot.kv.kv_allocated_len) # Update req_nodes to this successfully finished request. req.session.finish_req(req) @@ -564,10 +561,10 @@ class StreamingSession(BasePrefixCache): """ self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len) slot.kv.kv_allocated_len = prefix_len - slot.kv_committed_len = min(slot.kv_committed_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_committed_len = min(req.kv_committed_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) def _trim_overshoot(self, req: Req, finished_len: int) -> None: @@ -579,7 +576,7 @@ class StreamingSession(BasePrefixCache): target = len(req.origin_input_ids) + finished_len self._free_kv_aligned(req.req_pool_idx, target, req.kv.kv_allocated_len) req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target) - req.kv_committed_len = min(req.kv_committed_len, target) + req.kv.kv_committed_len = min(req.kv.kv_committed_len, target) req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target) req.output_ids = req.output_ids[:finished_len] diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 7dba1d74b..0a22093d6 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -159,7 +159,7 @@ class DFlashDraftInputV2(SpecInput): nxt_kv_lens_sum = 0 committed_seq_lens_sum = 0 for i, (req, cur, nxt) in enumerate(zip(batch.reqs, cur_kv_lens, nxt_kv_lens)): - committed_len = int(req.kv_committed_len) + committed_len = int(req.kv.kv_committed_len) committed_seq_lens_sum += committed_len top_k = int(req.sampling_params.top_k) diff --git a/python/sglang/test/scripted_runtime/context/queries.py b/python/sglang/test/scripted_runtime/context/queries.py index 7cfab1f40..d26fcffa8 100644 --- a/python/sglang/test/scripted_runtime/context/queries.py +++ b/python/sglang/test/scripted_runtime/context/queries.py @@ -117,7 +117,7 @@ def remaining_prompt_tokens(ctx: ScriptedContext, rid: str) -> int: req = find_req_by_rid(ctx, rid) if req is None: return 0 - return max(0, len(req.origin_input_ids) - req.kv_committed_len) + return max(0, len(req.origin_input_ids) - req.kv.kv_committed_len) def chunks_done(ctx: ScriptedContext, rid: str) -> int: diff --git a/test/manual/chunked_prefill/test_scripted_hybrid_swa.py b/test/manual/chunked_prefill/test_scripted_hybrid_swa.py index 3dc808ae6..b4af4cb90 100644 --- a/test/manual/chunked_prefill/test_scripted_hybrid_swa.py +++ b/test/manual/chunked_prefill/test_scripted_hybrid_swa.py @@ -96,10 +96,10 @@ class TestSWABasic(ScriptedTestCase): r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) for _ in range(400): if r.is_chunking: - assert len(r.req.prefix_indices) <= r.req.kv_committed_len, ( + assert len(r.req.prefix_indices) <= r.req.kv.kv_committed_len, ( f"prefix_indices must be bounded by kv_committed_len, " f"got prefix_indices_len={len(r.req.prefix_indices)}, " - f"kv_committed_len={r.req.kv_committed_len}" + f"kv_committed_len={r.req.kv.kv_committed_len}" ) if r.finished: break diff --git a/test/manual/chunked_prefill/test_scripted_regression.py b/test/manual/chunked_prefill/test_scripted_regression.py index be9bc25aa..2c9e11780 100644 --- a/test/manual/chunked_prefill/test_scripted_regression.py +++ b/test/manual/chunked_prefill/test_scripted_regression.py @@ -509,7 +509,7 @@ class TestRegressionGptOss(ScriptedTestCase): r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) - committed = r.req.kv_committed_len + committed = r.req.kv.kv_committed_len assert committed > 0 assert len(r.req.prefix_indices) <= committed, ( diff --git a/test/manual/chunked_prefill/test_scripted_special_case.py b/test/manual/chunked_prefill/test_scripted_special_case.py index b8a66c750..3549168c3 100644 --- a/test/manual/chunked_prefill/test_scripted_special_case.py +++ b/test/manual/chunked_prefill/test_scripted_special_case.py @@ -225,10 +225,10 @@ class TestSpecialCaseBasic(ScriptedTestCase): ) for _ in range(DEFAULT_MAX_STEPS): if r.is_chunking: - assert len(r.req.prefix_indices) <= r.req.kv_committed_len, ( + assert len(r.req.prefix_indices) <= r.req.kv.kv_committed_len, ( f"streaming-session chunked stash must stay bounded by " f"kv_committed_len; prefix_indices_len={len(r.req.prefix_indices)}, " - f"kv_committed_len={r.req.kv_committed_len}" + f"kv_committed_len={r.req.kv.kv_committed_len}" ) if r.finished: break diff --git a/test/registered/unit/beam_search/test_fork.py b/test/registered/unit/beam_search/test_fork.py index 23cfff930..1c6806f99 100644 --- a/test/registered/unit/beam_search/test_fork.py +++ b/test/registered/unit/beam_search/test_fork.py @@ -119,8 +119,9 @@ class _FakeAllocator: class TestFreeMemberRows(CustomTestCase): def _make_group(self, req_to_token, allocated_len): leader = SimpleNamespace( - kv=SimpleNamespace(kv_allocated_len=allocated_len), - kv_committed_len=allocated_len, + kv=SimpleNamespace( + kv_allocated_len=allocated_len, kv_committed_len=allocated_len + ), ) return SimpleNamespace( leader=leader, @@ -146,7 +147,7 @@ class TestFreeMemberRows(CustomTestCase): # Leader rewound to the prompt: its own release must not free the # decode region a second time. self.assertEqual(leader.kv.kv_allocated_len, 5) - self.assertEqual(leader.kv_committed_len, 5) + self.assertEqual(leader.kv.kv_committed_len, 5) self.assertEqual(sorted(pool.freed), [1, 2]) self.assertIsNone(group.member_rows) self.assertIsNone(group.member_rows_cpu) @@ -207,7 +208,7 @@ class TestRetireReclaimsStagedOrphans(CustomTestCase): allocator = _FakeAllocator() group = SimpleNamespace( leader=SimpleNamespace( - kv=SimpleNamespace(kv_allocated_len=8), kv_committed_len=8 + kv=SimpleNamespace(kv_allocated_len=8, kv_committed_len=8), ), prompt_len=5, member_rows=torch.tensor([1, 2], dtype=torch.int64), diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index c4b40bd4b..f2595158b 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py @@ -34,10 +34,11 @@ def _make_mock_req( req = MagicMock() req.rid = rid req.req_pool_idx = req_pool_idx - req.kv_committed_len = kv_committed_len - req.kv = ReqKvInfo(kv_allocated_len=kv_allocated_len) + req.kv = ReqKvInfo( + kv_committed_len=kv_committed_len, kv_allocated_len=kv_allocated_len + ) req.prefix_indices = list(range(prefix_indices_len)) - req.effective_kv_committed_len = lambda: req.kv_committed_len + req.effective_kv_committed_len = lambda: req.kv.kv_committed_len return req diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index ecc2b48e4..11c816d1f 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1521,7 +1521,6 @@ if _HAS_MLX: self.req_pool_idx = None self.mamba_pool_idx = None self.inflight_middle_chunks = 0 - self.kv_committed_len = 0 class FakeTpWorker: def __init__(self, next_token_ids): diff --git a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py index 322e63c6c..a75eccdea 100644 --- a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py +++ b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py @@ -121,7 +121,6 @@ def _fake_req(): return SimpleNamespace( req_pool_idx=None, inflight_middle_chunks=0, - kv_committed_len=0, mamba_pool_idx=None, mamba_ping_pong_track_buffer=None, ) diff --git a/test/registered/unit/layers/test_minicpm_sparse_cache.py b/test/registered/unit/layers/test_minicpm_sparse_cache.py index 0cf4aeb80..9d33d8f25 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_cache.py +++ b/test/registered/unit/layers/test_minicpm_sparse_cache.py @@ -68,7 +68,6 @@ def make_pool_and_req(capacity: int = 64): req = SimpleNamespace( req_pool_idx=None, inflight_middle_chunks=0, - kv_committed_len=0, ) req_pool_idx = pool.alloc([req])[0] return pool, req, req_pool_idx, allocator diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py index 9ba65bc57..7ec18c12a 100644 --- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py +++ b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py @@ -33,7 +33,7 @@ def _make_batch() -> tuple[Req, ScheduleBatch]: vocab_size=128, ) req.output_ids.append(3) - req.kv_committed_len = 2 + req.kv.kv_committed_len = 2 batch = ScheduleBatch(reqs=[req]) batch.tree_cache = SimpleNamespace(page_size=TRACK_INTERVAL) diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index 58f059b8d..a61365397 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -92,7 +92,7 @@ def _make_req(terminate_after: int) -> Req: sampling_params=sp, ) req.grammar = _FakeGrammar(terminate_after=terminate_after) - req.kv_committed_len = 0 + req.kv.kv_committed_len = 0 return req @@ -121,7 +121,7 @@ class TestSpecV2GrammarTruncation(CustomTestCase): self.assertEqual(predict_tokens, [[101, 102]]) # No pre-claim: commit the full retained run (no -1 refund). - self.assertEqual(req.kv_committed_len, 2) + self.assertEqual(req.kv.kv_committed_len, 2) def test_resolve_keeps_all_when_grammar_not_terminated(self): req = _make_req(terminate_after=99) @@ -131,7 +131,7 @@ class TestSpecV2GrammarTruncation(CustomTestCase): predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req])) self.assertEqual(predict_tokens, [[201, 202, 203]]) - self.assertEqual(req.kv_committed_len, 3) + self.assertEqual(req.kv.kv_committed_len, 3) class TestReasoningTokenAccounting(CustomTestCase): diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index 7fbafbae6..be5586464 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -50,8 +50,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None): fill_ids=origin_input_ids + output_ids, seqlen=len(origin_input_ids) + len(output_ids), req_pool_idx=None, - kv=SimpleNamespace(kv_allocated_len=0), - kv_committed_len=0, + kv=SimpleNamespace(kv_allocated_len=0, kv_committed_len=0), finished_reason=None, hisparse_staging=False, staging=False, @@ -219,7 +218,7 @@ class TestHiSparseUnit(unittest.TestCase): self.assertIsNotNone(kv_loc, "KV alloc failed") self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc) req.kv.kv_allocated_len = fill_len - req.kv_committed_len = fill_len + req.kv.kv_committed_len = fill_len req.full_untruncated_fill_ids = array("q", range(fill_len)) req.extend_range = Range(0, fill_len) return kv_loc @@ -579,7 +578,7 @@ class TestHiSparseUnit(unittest.TestCase): seq_len = fill_len + 1 self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc) req.kv.kv_allocated_len = seq_len - req.kv_committed_len = seq_len + req.kv.kv_committed_len = seq_len self.coordinator.map_last_loc_to_buffer( seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device), @@ -770,7 +769,7 @@ class TestHiSparseUnit(unittest.TestCase): ) ) self.assertEqual(req.kv.kv_allocated_len, fill_len) - self.assertEqual(req.kv_committed_len, fill_len) + self.assertEqual(req.kv.kv_committed_len, fill_len) self.assertEqual(req.extend_range.length, fill_len) rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size diff --git a/test/registered/unit/managers/test_kv_page_invariants.py b/test/registered/unit/managers/test_kv_page_invariants.py index 73ece119f..b9629b6f4 100644 --- a/test/registered/unit/managers/test_kv_page_invariants.py +++ b/test/registered/unit/managers/test_kv_page_invariants.py @@ -47,16 +47,18 @@ class _FakeReq: def __init__(self, rid, rpi, committed, allocated): self.rid = rid self.req_pool_idx = rpi - self.kv_committed_len = committed - self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0) + self.kv = SimpleNamespace( + kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0 + ) self.is_holding_kv = True class _FakeSlot: def __init__(self, rpi, committed, allocated): self.req_pool_idx = rpi - self.kv_committed_len = committed - self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0) + self.kv = SimpleNamespace( + kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0 + ) self.is_holding_kv = True diff --git a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py index eb05ac078..79d5e75d7 100644 --- a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py +++ b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py @@ -18,8 +18,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") def _make_req(): return types.SimpleNamespace( decode_batch_idx=0, - kv_committed_len=3, - kv_allocated_len=3, + kv=types.SimpleNamespace(kv_committed_len=3, kv_allocated_len=3), beam_group=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 d0b792b23..f574baa5a 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 @@ -81,9 +81,10 @@ class MockReq: 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), cache_protected_len=cache_protected_len + kv_committed_len=len(fill_ids), + kv_allocated_len=len(fill_ids), + cache_protected_len=cache_protected_len, ) def get_fill_ids(self): @@ -211,7 +212,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): cache.cache_unfinished_req(req) # Step 3: cache_finished_req with is_insert=True (dec lock) - cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len) + cache.cache_finished_req(req, kv_len_to_handle=req.kv.kv_committed_len) # Verify: all non-root nodes should have lock_ref == 0 # (root always has lock_ref == 1) @@ -260,7 +261,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): cache.cache_unfinished_req(req) # Step 3: cache_finished_req (dec leaf) - cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len) + cache.cache_finished_req(req, kv_len_to_handle=req.kv.kv_committed_len) # Root lock unchanged, all nodes unlocked self.assertEqual(cache.root_node.lock_ref, root_lock_before) @@ -304,7 +305,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): # Transfer fails -> cache_finished_req with is_insert=False # This frees delta tokens and dec_lock_ref on last_node cache.cache_finished_req( - req, is_insert=False, kv_len_to_handle=req.kv_committed_len + req, is_insert=False, kv_len_to_handle=req.kv.kv_committed_len ) # The prefix node should be unlocked (back to evictable) @@ -351,7 +352,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): # Transfer fails -> cache_finished_req with is_insert=False # dec_lock_ref(root) is a no-op cache.cache_finished_req( - req, is_insert=False, kv_len_to_handle=req.kv_committed_len + req, is_insert=False, kv_len_to_handle=req.kv.kv_committed_len ) # Root lock unchanged, nothing protected or evictable @@ -482,7 +483,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): ) cache.cache_unfinished_req(req) - cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len) + cache.cache_finished_req(req, kv_len_to_handle=req.kv.kv_committed_len) # After all iterations, root lock should be 1, no protected nodes self.assertEqual(cache.root_node.lock_ref, 1) diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index abe5f34d4..b9aae98ee 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -64,9 +64,11 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False): req_pool_idx=req_pool_idx, dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"), inflight_middle_chunks=1 if req_pool_idx is not None else 0, - kv_committed_len=len(prefix) if req_pool_idx is not None else 0, kv=SimpleNamespace( - kv_allocated_len=len(prefix) + block_size if req_pool_idx is not None else 0 + kv_committed_len=len(prefix) if req_pool_idx is not None else 0, + kv_allocated_len=( + len(prefix) + block_size if req_pool_idx is not None else 0 + ), ), ) diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index d1994f883..7b6337981 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -153,7 +153,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase): self.assertEqual(kwargs["swa_tail_len"], swa_tail_len) self.assertEqual(req.kv.swa_evicted_seqlen, fill_len - swa_tail_len) self.assertEqual(req.kv.kv_allocated_len, fill_len) - self.assertEqual(req.kv_committed_len, fill_len) + self.assertEqual(req.kv.kv_committed_len, fill_len) self.assertEqual(req.extend_range.length, fill_len) self.assertEqual(len(req_to_token_pool.writes), 1) coordinator.host_token_len.assert_called_once_with(fill_len) diff --git a/test/registered/unit/mem_cache/test_session_unified_radix_cache.py b/test/registered/unit/mem_cache/test_session_unified_radix_cache.py index 55dff04d4..60f7e95bd 100644 --- a/test/registered/unit/mem_cache/test_session_unified_radix_cache.py +++ b/test/registered/unit/mem_cache/test_session_unified_radix_cache.py @@ -166,7 +166,6 @@ def register(cache, token_ids, session_id, generation=None): ).last_device_node, origin_input_ids=array("q", token_ids), output_ids=array("q"), - kv_committed_len=len(token_ids), extra_key=None, ) ) 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 e002cac63..2bed0f237 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -68,8 +68,8 @@ class _FakeReq: _inflight=False, ) self.req_pool_idx = req_pool_idx - self.kv_committed_len = committed self.kv = SimpleNamespace( + kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0, cache_protected_len=0, @@ -113,9 +113,11 @@ def test_preabort_detaches_session_and_preserves_slot(): tree_cache = StreamingSession(inner) 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_committed_len=48, + kv_allocated_len=48, + swa_evicted_seqlen=0, + cache_protected_len=16, ), ) @@ -134,7 +136,7 @@ def test_preabort_detaches_session_and_preserves_slot(): # Slot untouched. slot = tree_cache.slots["session-a"] assert slot.req_pool_idx == 0 - assert slot.kv_committed_len == 48 + assert slot.kv.kv_committed_len == 48 assert slot.kv.kv_allocated_len == 48 assert len(result.device_indices) == 0 @@ -178,9 +180,11 @@ def test_nth_mid_abort_nukes_session_slot(): # Session already has a slot from a previous turn. tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, - kv_committed_len=50, kv=SimpleNamespace( - kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0 + kv_committed_len=50, + kv_allocated_len=50, + swa_evicted_seqlen=0, + cache_protected_len=0, ), last_node=None, ) @@ -217,9 +221,11 @@ def test_release_session_threads_mamba_skip_ids(): lock_node = SimpleNamespace(id=42) tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, - kv_committed_len=50, kv=SimpleNamespace( - kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0 + kv_committed_len=50, + kv_allocated_len=50, + swa_evicted_seqlen=0, + cache_protected_len=0, ), last_node=lock_node, skip_lock_node_ids={ComponentType.MAMBA: {42}}, @@ -265,7 +271,7 @@ def test_trim_overshoot_postcondition(): tree_cache._trim_overshoot(req, finished_len=12) target = 38 - assert req.kv_committed_len == target + assert req.kv.kv_committed_len == target assert req.kv.kv_allocated_len == target assert req.kv.swa_evicted_seqlen == target assert len(req.output_ids) == 12 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 314f5652a..527a8c3a6 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 @@ -643,7 +643,7 @@ def bench_cache_finished( ) req.last_node = node req.kv.cache_protected_len = matched_len - req.kv_committed_len = len(seq) + req.kv.kv_committed_len = len(seq) if hasattr(lr, "swa_uuid_for_lock"): req.swa_uuid_for_lock = lr.swa_uuid_for_lock env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices @@ -657,7 +657,7 @@ def bench_cache_finished( "cache_finished", lambda: req_items, lambda req: env.tree.cache_finished_req( - req, is_insert=True, kv_len_to_handle=req.kv_committed_len + req, is_insert=True, kv_len_to_handle=req.kv.kv_committed_len ), len(req_items) - warmup, env.avg_tokens, 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 db3530223..9855b093a 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 @@ -1226,7 +1226,7 @@ class UnifiedRadixCacheSuite: kv_len = len(input_ids) + len(output_ids) kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) - req.kv_committed_len = kv_len + req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -1266,7 +1266,7 @@ class UnifiedRadixCacheSuite: kv_len = req.extend_range.end kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) - req.kv_committed_len = kv_len + req.kv.kv_committed_len = kv_len req.kv.kv_allocated_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -1311,7 +1311,7 @@ class UnifiedRadixCacheSuite: kv_len = len(tokens) kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) - req.kv_committed_len = kv_len + req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -1346,7 +1346,7 @@ class UnifiedRadixCacheSuite: kv_len = len(tokens) kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) - req.kv_committed_len = kv_len + req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -1383,7 +1383,7 @@ class UnifiedRadixCacheSuite: req.set_extend_range(0, len(req.full_untruncated_fill_ids)) kv_indices = self._alloc(allocator, len(tokens)) req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) - req.kv_committed_len = len(tokens) + req.kv.kv_committed_len = len(tokens) req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -1485,7 +1485,7 @@ class UnifiedRadixCacheSuite: kv_len = len(input_ids) kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) - req.kv_committed_len = kv_len + req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -1610,7 +1610,7 @@ class UnifiedRadixCacheSuite: kv_len = len(tokens) fresh_value = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value) - req.kv_committed_len = kv_len + req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -2197,7 +2197,7 @@ class UnifiedRadixCacheSuite: req.set_extend_range(0, len(req.full_untruncated_fill_ids)) kv_indices = self._alloc(allocator, pre_len) req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) - req.kv_committed_len = pre_len + req.kv.kv_committed_len = pre_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -2286,7 +2286,7 @@ class UnifiedRadixCacheSuite: req.set_extend_range(0, len(req.full_untruncated_fill_ids)) kv_indices = self._alloc(allocator, pre_len) req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) - req.kv_committed_len = pre_len + req.kv.kv_committed_len = pre_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -6768,7 +6768,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): ) req_to_token_pool.alloc([req]) req.output_ids = array("q") - req.kv_committed_len = len(tokens) + req.kv.kv_committed_len = len(tokens) req.kv.kv_allocated_len = len(tokens) req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -8095,7 +8095,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase): req.set_extend_range(0, len(req.full_untruncated_fill_ids)) kv_indices = self._alloc_paged(allocator, seq_len) req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices) - req.kv_committed_len = seq_len + req.kv.kv_committed_len = seq_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 req.swa_uuid_for_lock = None diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 8aaf92bc3..c5cb85193 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -104,9 +104,7 @@ _OWNER_SITES = { "free_member_rows", "kv_allocated_len", ): 1, - # streaming session slot save/restore and tail trimming - (_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1, - (_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1, + # streaming session tail trimming (_SS, "StreamingSession._free_tail", "kv_committed_len"): 2, (_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2, (_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1, diff --git a/test/registered/xpu/test_lmcache_radix_cache.py b/test/registered/xpu/test_lmcache_radix_cache.py index 354767e2a..bf0100cd6 100644 --- a/test/registered/xpu/test_lmcache_radix_cache.py +++ b/test/registered/xpu/test_lmcache_radix_cache.py @@ -72,10 +72,9 @@ def _make_req(rid, req_pool_idx, token_ids, tree): extra_key=None, cache_salt=None, last_node=tree.root_node, - kv=SimpleNamespace(cache_protected_len=0), + kv=SimpleNamespace(cache_protected_len=0, kv_committed_len=len(token_ids)), priority=0, kv_committed_freed=False, - kv_committed_len=len(token_ids), ) req.pop_committed_kv_cache = lambda: len(token_ids) return req @@ -163,11 +162,7 @@ class TestLMCRadixCacheXPU(unittest.TestCase): # commit it as a finished request (inserts into radix + stores to # LMCache on tree.store_stream). req_pool_idx = req_to_token_pool.alloc( - [ - SimpleNamespace( - req_pool_idx=None, inflight_middle_chunks=0, kv_committed_len=0 - ) - ] + [SimpleNamespace(req_pool_idx=None, inflight_middle_chunks=0)] )[0] kv_slots = allocator.alloc(self.INPUT_LEN) self.assertIsNotNone(kv_slots)