[mem_cache] Move kv_committed_len into ReqKvInfo (#37078)

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

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

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)
 

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

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.

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:

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,
                     )
 

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

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)

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"
 

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

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

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]
 

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)
 

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: