Introduce req.kv container for coupled owned kv field lifecycle (#29427)

python/sglang/srt/disaggregation/decode.py

@@ -1420,7 +1420,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_allocated_len = fill_len
+        req.kv.kv_allocated_len = fill_len
         req.kv_committed_len = fill_len
 
         if prefix_len > 0:
@@ -1525,7 +1525,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     extend_num_tokens=fill_len,
                     swa_tail_len=self._swa_tail_len(fill_len),
                 )
-                req.swa_evicted_seqlen = fill_len - self._swa_tail_len(fill_len)
+                req.kv.swa_evicted_seqlen = fill_len - self._swa_tail_len(fill_len)
             else:
                 kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
                     prefix_lens=torch.tensor(

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py

@@ -663,7 +663,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_allocated_len)
+        kv_len = max(req.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/hisparse_coordinator.py

@@ -268,7 +268,7 @@ class HiSparseCoordinator:
         """
         self.alloc_device_buffer(req)
 
-        host_len = self.host_token_len(req.kv_allocated_len)
+        host_len = self.host_token_len(req.kv.kv_allocated_len)
         if host_len <= self.device_buffer_size:
             # Short sequences (seq_len <= device_buffer_size): the kernel fast path
             # returns device_buffer_locs directly without any host loading, so we
@@ -293,7 +293,7 @@ class HiSparseCoordinator:
 
     def _preload_to_device_buffer(self, req: Req) -> None:
         """Preload all tokens from host pool into the device buffer."""
-        n = self.host_token_len(req.kv_allocated_len)
+        n = self.host_token_len(req.kv.kv_allocated_len)
         host_indices = self.req_to_host_pool[req.req_pool_idx, :n]
         device_locs = self.req_to_device_buffer[req.req_pool_idx, :n]
 
@@ -311,7 +311,7 @@ class HiSparseCoordinator:
             allocated_len = req.extend_range.end
             alloc_size = self.padded_buffer_size
         else:
-            allocated_len = req.kv_allocated_len
+            allocated_len = req.kv.kv_allocated_len
             page_size = self.mem_pool_device.page_size
             # Allocate only enough for current tokens (page-aligned).
             # When prefill already fills device_buffer_size, include the reserved page.
@@ -766,7 +766,7 @@ class HiSparseCoordinator:
         # we just freed via free_hisparse_indices(all_hi). If left set, the
         # subsequent release_kv_cache -> allocator.free -> free_hisparse path
         # re-frees them (double-free into the page allocator's free list).
-        allocated_len = req.kv_allocated_len
+        allocated_len = req.kv.kv_allocated_len
 
         # release memory -- only free actually-allocated buffer indices
         current_cap = int(self.req_device_buffer_size[req.req_pool_idx])

python/sglang/srt/managers/schedule_batch.py

@@ -661,6 +661,17 @@ class ReqLogprob:
     output_token_ids_logprobs_idx: Optional[list] = None
 
 
+@dataclasses.dataclass(slots=True, kw_only=True)
+class ReqKvInfo:
+    kv_allocated_len: int
+    # 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
+
+
 class Req(ReqDllmMixin):
     """The input and output status of a request."""
 
@@ -737,19 +748,13 @@ class Req(ReqDllmMixin):
 
         # For req-level memory management
         self.kv_committed_len = 0
-        self.kv_allocated_len = 0
+        self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
         self.kv_committed_freed = False
         self.kv_overallocated_freed = False
 
         # for cross-encoder model
         self.token_type_ids = token_type_ids
 
-        # 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`.
-        self.swa_evicted_seqlen = 0
         # 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
@@ -1094,9 +1099,9 @@ class Req(ReqDllmMixin):
         # e.g., speculative decoding may allocate more KV cache than actually used.
         assert (
             not self.kv_overallocated_freed
-        ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv_allocated_len=}"
+        ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv.kv_allocated_len=}"
         self.kv_overallocated_freed = True
-        return self._cache_commit_len(), self.kv_allocated_len
+        return self._cache_commit_len(), self.kv.kv_allocated_len
 
     def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
         """Update the speculative decoding acceptance histogram.
@@ -1518,11 +1523,11 @@ class Req(ReqDllmMixin):
         self.mamba_cow_src_index = None
         self.mamba_needs_clear = False
         self.already_computed = 0
-        self.kv_allocated_len = 0
+        self.kv.kv_allocated_len = 0
         self.kv_committed_len = 0
         self.kv_committed_freed = False
         self.kv_overallocated_freed = False
-        self.swa_evicted_seqlen = 0
+        self.kv.swa_evicted_seqlen = 0
         self.extend_batch_idx = 0
         self.decode_batch_idx = 0
 
@@ -2190,7 +2195,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
             # update req-level memory management fields
             req.kv_committed_len = seq_len
-            req.kv_allocated_len = seq_len
+            req.kv.kv_allocated_len = seq_len
 
             # If input_embeds are available, store them
             if req.input_embeds is not None:
@@ -2558,8 +2563,8 @@ 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_allocated_len)
-            cur = r.kv_allocated_len
+            x = max(0, r.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)
         return total
@@ -2790,7 +2795,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         for req in self.reqs:
             req.decode_batch_idx += 1
             req.kv_committed_len += 1
-            req.kv_allocated_len += 1
+            req.kv.kv_allocated_len += 1
 
         # New-tensor avoids racing model_worker_batch refs queued for
         # overlap forward.

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -250,7 +250,7 @@ class SchedulerInvariantChecker:
                 if req.kv_committed_freed or req.req_pool_idx is None:
                     continue
 
-                allocated_len = req.kv_allocated_len
+                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
@@ -258,7 +258,7 @@ class SchedulerInvariantChecker:
                 full_uncached += allocated_len - req.cache_protected_len
                 if self.is_hybrid_swa:
                     swa_uncached += allocated_len - max(
-                        req.cache_protected_len, req.swa_evicted_seqlen
+                        req.cache_protected_len, req.kv.swa_evicted_seqlen
                     )
 
         return full_uncached, swa_uncached
@@ -321,7 +321,7 @@ class SchedulerInvariantChecker:
                     f"req {req.rid}",
                     req.req_pool_idx,
                     req.kv_committed_len,
-                    req.kv_allocated_len,
+                    req.kv.kv_allocated_len,
                 )
         sess = getattr(self.tree_cache, "slots", None)
         if sess:
@@ -332,7 +332,7 @@ class SchedulerInvariantChecker:
                         f"slot {sid[:8]}",
                         slot.req_pool_idx,
                         slot.kv_committed_len,
-                        slot.kv_allocated_len,
+                        slot.kv.kv_allocated_len,
                     )
 
         active = [

python/sglang/srt/mem_cache/chunk_cache.py

@@ -156,7 +156,7 @@ class PureSWAChunkCache(SWAChunkCache):
             req.req_pool_idx, :kv_committed_len
         ]
         evict_floor = req.swa_evict_floor
-        evicted_seqlen = req.swa_evicted_seqlen
+        evicted_seqlen = req.kv.swa_evicted_seqlen
         if evicted_seqlen > evict_floor:
             parts = []
             if evict_floor > 0:

python/sglang/srt/mem_cache/common.py

@@ -77,7 +77,7 @@ def free_swa_out_of_window_slots(
     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.swa_evicted_seqlen = max(req.swa_evicted_seqlen, evict_floor)
+    req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, evict_floor)
 
     if is_chunk_cache:
         # Chunk cache builds no radix tree, so no tombstone-leaf concern; evict
@@ -90,22 +90,22 @@ def free_swa_out_of_window_slots(
         # No extra page margin is needed.
         evict_threshold = pre_len - max(sliding_window_size, page_size)
     new_swa_evicted_seqlen = max(
-        req.swa_evicted_seqlen,
+        req.kv.swa_evicted_seqlen,
         evict_threshold,
     )
 
     if page_size > 1:
         new_swa_evicted_seqlen = (new_swa_evicted_seqlen // page_size) * page_size
 
-    if new_swa_evicted_seqlen > req.swa_evicted_seqlen:
+    if new_swa_evicted_seqlen > req.kv.swa_evicted_seqlen:
         free_slots = req_to_token_pool.req_to_token[
-            req.req_pool_idx, req.swa_evicted_seqlen : new_swa_evicted_seqlen
+            req.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen
         ]
         token_to_kv_pool_allocator.free_swa(free_slots)
         maybe_evict_dsv4_state_on_swa(
             token_to_kv_pool_allocator, req_to_token_pool, req, new_swa_evicted_seqlen
         )
-        req.swa_evicted_seqlen = new_swa_evicted_seqlen
+        req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
 
 
 def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
@@ -661,7 +661,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
     if spec_algo is None and not global_server_args.strip_thinking_cache:
         assert (
             start_p == end_p
-        ), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"
+        ), f"Unexpected overallocated KV cache, {req.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/pure_swa_radix_cache.py

@@ -93,7 +93,7 @@ class PureSWARadixCache(RadixCache):
 
         old_prefix_len = req.cache_protected_len
         swa_evict_floor = req.swa_evict_floor
-        swa_evicted_seqlen = req.swa_evicted_seqlen
+        swa_evicted_seqlen = req.kv.swa_evicted_seqlen
 
         if self.page_size > 1 and swa_evict_floor > 0:
             swa_evict_floor = -(-swa_evict_floor // self.page_size) * self.page_size

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -486,7 +486,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
                     key=radix_key,
                     value=values,
                     prev_prefix_len=old_prefix_len,
-                    swa_evicted_seqlen=req.swa_evicted_seqlen,
+                    swa_evicted_seqlen=req.kv.swa_evicted_seqlen,
                 )
             )
         else:

python/sglang/srt/mem_cache/unified_cache_components/swa_component.py

@@ -589,7 +589,7 @@ class SWAComponent(TreeComponent):
     ) -> Optional[int]:
         # Unfinished requests can already have an SWA-evicted prefix; preserve
         # that boundary so insertion creates a tombstone instead of live SWA KV.
-        insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
+        insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
         return None
 
     def free_out_of_window_slots(
@@ -604,7 +604,7 @@ class SWAComponent(TreeComponent):
                 req_to_token_pool=self.cache.req_to_token_pool,
                 token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
             )
-        insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
+        insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
 
     # ---- HiCache Hooks ----
 

python/sglang/srt/session/streaming_session.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import logging
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, Optional
@@ -20,7 +21,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
 from sglang.srt.utils.common import ceil_align
 
 if TYPE_CHECKING:
-    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
 
 
 logger = logging.getLogger(__name__)
@@ -36,6 +37,12 @@ class _VirtualNode:
     pass
 
 
+def _new_kv() -> ReqKvInfo:
+    from sglang.srt.managers.schedule_batch import ReqKvInfo
+
+    return ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
+
+
 @dataclass
 class SessionSlot:
     """Holds KV state between streaming session turns."""
@@ -45,16 +52,13 @@ class SessionSlot:
     # KV pool state (None means no KV is currently held by this slot)
     req_pool_idx: Optional[int] = None
     kv_committed_len: int = 0
-    kv_allocated_len: int = 0
+    kv: ReqKvInfo = field(default_factory=_new_kv)
 
     # 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
 
-    # SWA state
-    swa_evicted_seqlen: int = 0
-
     # Mamba states
     mamba_pool_idx: Any = None
     mamba_ping_pong_track_buffer: Any = None
@@ -71,8 +75,7 @@ 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_allocated_len = req.kv_allocated_len
-        self.swa_evicted_seqlen = req.swa_evicted_seqlen
+        self.kv = copy.copy(req.kv)
 
         if is_first:
             self.last_node = req.last_node
@@ -104,8 +107,7 @@ class SessionSlot:
         """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_allocated_len = self.kv_allocated_len
-        req.swa_evicted_seqlen = self.swa_evicted_seqlen
+        req.kv = copy.copy(self.kv)
         req.swa_uuid_for_lock = self.swa_uuid_for_lock
 
         req.mamba_pool_idx = self.mamba_pool_idx
@@ -305,7 +307,7 @@ class StreamingSession(BasePrefixCache):
                 # the mamba pool; otherwise the abort orphans them.
                 slot = SessionSlot(
                     req_pool_idx=req.req_pool_idx,
-                    kv_allocated_len=req.kv_allocated_len,
+                    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,
@@ -317,7 +319,9 @@ class StreamingSession(BasePrefixCache):
                 # the abort fall-through doesn't double-free.
                 req.mamba_pool_idx = None
                 req.mamba_ping_pong_track_buffer = None
-            slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len)
+            slot.kv.kv_allocated_len = max(
+                slot.kv.kv_allocated_len, req.kv.kv_allocated_len
+            )
             self.release_session(session_id)
             req.req_pool_idx = None
             req.session.abort_req()
@@ -339,7 +343,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_allocated_len)
+        slot.kv_committed_len = min(target, slot.kv.kv_allocated_len)
 
         # Update req_nodes to this successfully finished request.
         req.session.finish_req(req)
@@ -411,7 +415,9 @@ class StreamingSession(BasePrefixCache):
         protected_len = slot.cache_protected_len
         lock_node = slot.last_node
         tokens_freed = (
-            max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0
+            max(0, slot.kv.kv_allocated_len - protected_len)
+            if slot.is_holding_kv
+            else 0
         )
         logger.info(
             "Session KV released: %s (%d tokens freed)", session_id, tokens_freed
@@ -428,7 +434,7 @@ class StreamingSession(BasePrefixCache):
 
         if slot.is_holding_kv:
             start = protected_len
-            end = slot.kv_allocated_len
+            end = slot.kv.kv_allocated_len
             if start < end:
                 kv_indices = self.req_to_token_pool.req_to_token[
                     slot.req_pool_idx, start:end
@@ -454,7 +460,7 @@ class StreamingSession(BasePrefixCache):
                 active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
             )
             if slot.is_holding_kv and not in_batch:
-                allocated = ceil_align(slot.kv_allocated_len, self.page_size)
+                allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
                 total += allocated - slot.cache_protected_len
         return total
 
@@ -470,9 +476,9 @@ class StreamingSession(BasePrefixCache):
                 active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
             )
             if slot.is_holding_kv and not in_batch:
-                allocated = ceil_align(slot.kv_allocated_len, self.page_size)
+                allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
                 total += allocated - max(
-                    slot.cache_protected_len, slot.swa_evicted_seqlen
+                    slot.cache_protected_len, slot.kv.swa_evicted_seqlen
                 )
         return total
 
@@ -527,13 +533,13 @@ class StreamingSession(BasePrefixCache):
         decoding pushes allocated above committed, or when retract retry's
         logit-reserve pulls prefix_len below committed.
         """
-        self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len)
-        slot.kv_allocated_len = prefix_len
+        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.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len)
-        req.kv_allocated_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.swa_evicted_seqlen = min(req.swa_evicted_seqlen, 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:
         """Trim slot KV to finished_len boundary. Spec v2 may overshoot
@@ -542,10 +548,10 @@ class StreamingSession(BasePrefixCache):
         be released to avoid token/KV mismatch.
         """
         target = len(req.origin_input_ids) + finished_len
-        self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
-        req.kv_allocated_len = min(req.kv_allocated_len, target)
+        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.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target)
+        req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
         req.output_ids = req.output_ids[:finished_len]
 
     def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None:

python/sglang/srt/speculative/dflash_info_v2.py

@@ -155,7 +155,7 @@ class DFlashDraftInputV2(SpecInput):
         for i, req in enumerate(batch.reqs):
             committed_len = int(req.kv_committed_len)
             # Read the allocation watermark from the req object like EAGLE.
-            cur_alloc_len = int(req.kv_allocated_len)
+            cur_alloc_len = int(req.kv.kv_allocated_len)
             reserved_len = max(cur_alloc_len, committed_len + 2 * block_size)
             top_k = int(req.sampling_params.top_k)
 
@@ -233,7 +233,9 @@ class DFlashDraftInputV2(SpecInput):
         # This request-side high-water mark is what release_kv_cache() uses to
         # reclaim any DFLASH over-allocation if the request finishes later.
         for i, req in enumerate(batch.reqs):
-            req.kv_allocated_len = max(req.kv_allocated_len, int(nxt_kv_lens_cpu_t[i]))
+            req.kv.kv_allocated_len = max(
+                req.kv.kv_allocated_len, int(nxt_kv_lens_cpu_t[i])
+            )
 
         # Seed committed; overlap's resolve overwrites it with the published value.
         batch.seq_lens_cpu = batch_seq_lens_cpu_t

python/sglang/srt/speculative/eagle_utils.py

@@ -833,7 +833,7 @@ def eagle_prepare_for_decode(batch: ScheduleBatch):
     nxt_kv_lens = [0] * bs
     num_needed_tokens = 0
     for i, r in enumerate(batch.reqs):
-        cur = r.kv_allocated_len
+        cur = r.kv.kv_allocated_len
         # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur.
         # kv_committed_len is honest (bonus committed in resolve, not here),
         # so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs.
@@ -841,7 +841,7 @@ def eagle_prepare_for_decode(batch: ScheduleBatch):
         cur_kv_lens[i] = cur
         nxt_kv_lens[i] = nxt
         num_needed_tokens += nxt - cur
-        r.kv_allocated_len = nxt
+        r.kv.kv_allocated_len = nxt
         r.decode_batch_idx += 1
 
     cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")

python/sglang/test/scripted_runtime/req_handle.py

@@ -45,7 +45,7 @@ class ScriptedReqHandle:
         if req is None or req.req_pool_idx is None:
             return 0
         page_size = self.context.scheduler.page_size
-        return (req.kv_allocated_len + page_size - 1) // page_size
+        return (req.kv.kv_allocated_len + page_size - 1) // page_size
 
     @property
     def lock_refs(self) -> int: