Let the presence of req.kv indicate the existence of owned kv resources (#29429)

python/sglang/srt/disaggregation/decode.py

@@ -1420,7 +1420,15 @@ 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.kv_allocated_len = fill_len
+        # TODO(th4): co-locate this req.kv bookkeeping with the real KV
+        # allocation; the pool alloc above and the kv_allocated_len assignment
+        # below should become a single owned-kv allocation step.
+        if req.kv is None:
+            from sglang.srt.managers.schedule_batch import ReqKvInfo
+
+            req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0)
+        else:
+            req.kv.kv_allocated_len = fill_len
         req.kv_committed_len = fill_len
 
         if prefix_len > 0:

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -286,6 +286,7 @@ class DecodeKVCacheOffloadManager:
             self.token_to_kv_pool_allocator.free(overalloc_indices)
 
         self.req_to_token_pool.free(req)
+        req.kv = None
         self.tree_cache.protected_size_ -= len(req.prefix_indices)
         if req.rid in self.offloaded_state:
             del self.offloaded_state[req.rid]

python/sglang/srt/managers/schedule_batch.py

@@ -748,7 +748,7 @@ class Req(ReqDllmMixin):
 
         # For req-level memory management
         self.kv_committed_len = 0
-        self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
+        self.kv: Optional[ReqKvInfo] = None
 
         # for cross-encoder model
         self.token_type_ids = token_type_ids
@@ -1501,9 +1501,8 @@ class Req(ReqDllmMixin):
         self.mamba_cow_src_index = None
         self.mamba_needs_clear = False
         self.already_computed = 0
-        self.kv.kv_allocated_len = 0
+        assert self.kv is None, "expect it is already released"
         self.kv_committed_len = 0
-        self.kv.swa_evicted_seqlen = 0
         self.extend_batch_idx = 0
         self.decode_batch_idx = 0
 
@@ -2170,8 +2169,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req.extend_batch_idx += 1
 
             # update req-level memory management fields
+            # 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_allocated_len = seq_len
+            if req.kv is None:
+                req.kv = ReqKvInfo(kv_allocated_len=seq_len, swa_evicted_seqlen=0)
+            else:
+                req.kv.kv_allocated_len = seq_len
 
             # If input_embeds are available, store them
             if req.input_embeds is not None:

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

@@ -315,6 +315,8 @@ class SchedulerInvariantChecker:
         batch = self.get_last_batch()
         if batch is not None:
             for req in batch.reqs:
+                if req.kv is None:
+                    continue
                 _add_owner(
                     req,
                     f"req {req.rid}",

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -721,6 +721,7 @@ class SchedulerPPMixin:
                     ]
                     self.token_to_kv_pool_allocator.free(kv_indices)
                     self.req_to_token_pool.free(req)
+                    req.kv = None
 
             logger.info(
                 f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: "

python/sglang/srt/mem_cache/common.py

@@ -70,6 +70,9 @@ def free_swa_out_of_window_slots(
     token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
     is_chunk_cache: bool = False,
 ) -> None:
+    if req.kv is None:
+        return
+
     # 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
@@ -683,6 +686,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
         tree_cache.req_to_token_pool.free_mamba_cache(req)
     # DSV4-NPU's free() also releases c4/c128 state pages; no-op for others.
     tree_cache.req_to_token_pool.free(req)
+    req.kv = None
 
 
 def available_and_evictable_str(tree_cache: BasePrefixCache) -> str:

python/sglang/srt/session/streaming_session.py

@@ -37,12 +37,6 @@ 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."""
@@ -52,7 +46,7 @@ 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: ReqKvInfo = field(default_factory=_new_kv)
+    kv: Optional[ReqKvInfo] = None
 
     # First req's radix tree node (for dec_lock_ref on session close)
     last_node: Any = None
@@ -97,6 +91,7 @@ class SessionSlot:
         # the slot's tensor to be reused by a new req and leaked when
         # the slot is later freed.
         req.req_pool_idx = None
+        req.kv = None
         req.mamba_pool_idx = None
         req.mamba_ping_pong_track_buffer = None
         req.mamba_next_track_idx = None
@@ -324,6 +319,7 @@ class StreamingSession(BasePrefixCache):
             )
             self.release_session(session_id)
             req.req_pool_idx = None
+            req.kv = None
             req.session.abort_req()
             return True