[misc] Keep req.kv non-optional and key KV ownership on req_pool_idx (#36958)

python/sglang/srt/beam_search/fork.py

@@ -87,15 +87,14 @@ def free_member_rows(group, req_to_token_pool, token_to_kv_pool_allocator) -> No
         return
     leader = group.leader
     start = group.prompt_len
-    end = leader.kv.kv_allocated_len if leader.kv is not None else start
+    end = leader.kv.kv_allocated_len
     if end > start:
         # The rewind below is required: without it the leader's own per-Req
         # 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())
-        if leader.kv is not None:
-            leader.kv_committed_len = start
-            leader.kv.kv_allocated_len = start
+        leader.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
     group.member_rows_cpu = None

python/sglang/srt/disaggregation/decode.py

@@ -67,7 +67,6 @@ from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     NextBatchPlan,
-    ReqKvInfo,
     ScheduleBatch,
 )
 from sglang.srt.managers.schedule_policy import match_prefix_for_req
@@ -1906,10 +1905,7 @@ def alloc_for_decode_prealloc_hisparse(
     uses_swa_tail: bool,
     swa_tail_len: int,
 ) -> torch.Tensor:
-    if req.kv is None:
-        req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0)
-    else:
-        req.kv.kv_allocated_len = fill_len
+    req.kv.kv_allocated_len = fill_len
     device = allocator.device
     prefix_lens = torch.tensor([0], dtype=torch.int64, device=device)
     prefix_lens_cpu = torch.tensor([0], dtype=torch.int64)
@@ -1954,10 +1950,7 @@ def alloc_for_decode_prealloc(
     swa_tail_len: int,
     req_to_token_pool: Optional[ReqToTokenPool] = None,
 ) -> torch.Tensor:
-    if req.kv is None:
-        req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0)
-    else:
-        req.kv.kv_allocated_len = fill_len
+    req.kv.kv_allocated_len = fill_len
     if allocator.page_size == 1:
         kv_loc = allocator.alloc(delta_len)
     else:

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -281,7 +281,7 @@ class DecodeKVCacheOffloadManager:
             self.token_to_kv_pool_allocator.free(overalloc_indices)
 
         self.req_to_token_pool.free(req)
-        req.kv = None
+        req.kv.mark_released()
         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/disaggregation/prefill.py

@@ -1036,11 +1036,7 @@ class SchedulerDisaggregationPrefillMixin:
         else:
             logger.warning(error_message)
         req.time_stats.trace_ctx.abort(abort_info={"reason": error_message})
-        if (
-            req.req_pool_idx is not None
-            or req.kv is not None
-            or req.mamba_pool_idx is not None
-        ):
+        if req.is_holding_kv or req.mamba_pool_idx is not None:
             release_kv_cache(req, self.tree_cache)
         maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
         req.pending_bootstrap = False

python/sglang/srt/managers/schedule_batch.py

@@ -815,13 +815,23 @@ class ReqLogprob:
 
 @dataclasses.dataclass(slots=True, kw_only=True)
 class ReqKvInfo:
-    kv_allocated_len: int
+    # Device KV a request holds outside the prefix cache. Always present on the Req;
+    # whether any KV is held is `req.req_pool_idx is not None` (Req.is_holding_kv).
+    kv_allocated_len: int = 0
     # The length of KV that have been removed in swa cache.
     # SWA KV cache eviction behavior differs by cache type:
     # - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in
     #   `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction.
     # - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`.
-    swa_evicted_seqlen: int
+    swa_evicted_seqlen: int = 0
+
+    @property
+    def is_released(self) -> bool:
+        return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0
+
+    def mark_released(self) -> None:
+        self.kv_allocated_len = 0
+        self.swa_evicted_seqlen = 0
 
 
 class Req(ReqDllmMixin):
@@ -904,7 +914,7 @@ class Req(ReqDllmMixin):
 
         # For req-level memory management
         self.kv_committed_len = 0
-        self.kv: Optional[ReqKvInfo] = None
+        self.kv = ReqKvInfo()
         self.retraction_backup: Optional[RetractionBackup] = None
 
         # for cross-encoder model
@@ -1265,6 +1275,10 @@ class Req(ReqDllmMixin):
             or self.mamba_host_hit_length > 0
         )
 
+    @property
+    def is_holding_kv(self) -> bool:
+        return self.req_pool_idx is not None
+
     def effective_kv_committed_len(self) -> int:
         # Report only the prompt prefix so thinking + answer fall into the
         # overallocated range and are reclaimed by release_kv_cache. #22373.
@@ -1731,7 +1745,7 @@ class Req(ReqDllmMixin):
         self.mamba_cow_src_index = None
         self.mamba_needs_clear = False
         self.already_computed = 0
-        assert self.kv is None, "expect it is already released"
+        assert not self.is_holding_kv, "expect it is already released"
         self.kv_committed_len = 0
         self.extend_batch_idx = 0
         self.decode_batch_idx = 0
@@ -3481,7 +3495,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                     # seqlen progress is monotonic per KV handle.
                     if (
                         req.decode_batch_idx >= 1
-                        and req.kv is not None
+                        and req.is_holding_kv
                         and req.seqlen - 1 - sliding_window_size
                         >= req.kv.swa_evicted_seqlen + eviction_interval
                     ):

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

@@ -252,7 +252,7 @@ class SchedulerInvariantChecker:
         swa_uncached = 0
         for batch in batches:
             for req in batch.reqs:
-                if req.kv is None:
+                if not req.is_holding_kv:
                     continue
 
                 allocated_len = req.kv.kv_allocated_len
@@ -324,7 +324,7 @@ class SchedulerInvariantChecker:
         batch = self.get_last_batch()
         if batch is not None:
             for req in batch.reqs:
-                if req.kv is None:
+                if not req.is_holding_kv:
                     continue
                 _add_owner(
                     req,

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -723,7 +723,7 @@ class SchedulerPPMixin:
                 latencies.append(latency_ms)
 
                 # Release KV and Mamba cache
-                if req.req_pool_idx is not None:
+                if req.is_holding_kv:
                     kv_indices = self.req_to_token_pool.req_to_token[
                         req.req_pool_idx, : req.extend_range.end
                     ]
@@ -731,7 +731,7 @@ class SchedulerPPMixin:
                     if req.mamba_pool_idx is not None:
                         self.req_to_token_pool.free_mamba_cache(req)
                     self.req_to_token_pool.free(req)
-                    req.kv = None
+                    req.kv.mark_released()
 
             logger.info(
                 f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: "

python/sglang/srt/mem_cache/allocation.py

@@ -385,13 +385,8 @@ def alloc_for_extend(
             batch.seq_lens_cpu,
         )
 
-    from sglang.srt.managers.schedule_batch import ReqKvInfo
-
     for req, seq_len in zip(batch.reqs, batch.seq_lens_cpu.tolist()):
-        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
+        req.kv.kv_allocated_len = seq_len
 
     return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu
 
@@ -416,7 +411,7 @@ def _alloc_extend_loc_with_kv_reuse(
         retained_len = len(req.dllm_incomplete_ids)
         if extend_len != retained_len:
             raise RuntimeError("dLLM FDFO retained KV must be reused as a full block.")
-        if req.kv is None or prefix_len + extend_len > req.kv.kv_allocated_len:
+        if prefix_len + extend_len > req.kv.kv_allocated_len:
             raise RuntimeError("dLLM FDFO retained KV is missing.")
 
     alloc_extend_lens = [

python/sglang/srt/mem_cache/common.py

@@ -62,7 +62,7 @@ def free_swa_out_of_window_slots(
     is_chunk_cache: bool = False,
     retain_floor: int | None = None,
 ) -> None:
-    if req.kv is None:
+    if not req.is_holding_kv:
         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
@@ -200,10 +200,9 @@ def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> N
 
 
 def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
-    # the two resources currently have the same lifecycle, thus simplify logic below
-    assert (req.req_pool_idx is None) == (req.kv is None)
+    assert (not req.is_holding_kv) == req.kv.is_released
     # MambaRadixCache may alloc mamba state before alloc KV cache
-    if req.req_pool_idx is None:
+    if not req.is_holding_kv:
         assert (
             tree_cache.supports_mamba()
         ), "Only MambaRadixCache allow freeing before alloc"
@@ -224,8 +223,8 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
 
     # StreamingSession.cache_finished_req handles speculative tail trim
     # internally, then sets req_pool_idx = None.
-    assert (req.req_pool_idx is None) == (req.kv is None)
-    if req.req_pool_idx is None and req.kv is None:
+    assert (not req.is_holding_kv) == req.kv.is_released
+    if not req.is_holding_kv:
         return
 
     start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len
@@ -242,7 +241,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
     # The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the
     # c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here.
     tree_cache.req_to_token_pool.free(req)
-    req.kv = None
+    req.kv.mark_released()
 
 
 def _release_overallocated_kv_indices(

python/sglang/srt/session/streaming_session.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch
 
+from sglang.srt.managers.schedule_batch import ReqKvInfo
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
     DecLockRefParams,
@@ -21,7 +22,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
 from sglang.srt.utils.common import ceil_align, is_npu
 
 if TYPE_CHECKING:
-    from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
+    from sglang.srt.managers.schedule_batch import Req
 
 
 logger = logging.getLogger(__name__)
@@ -43,10 +44,10 @@ class SessionSlot:
 
     virtual_node: _VirtualNode = field(default_factory=_VirtualNode)
 
-    # KV pool state (None means no KV is currently held by this slot)
+    # KV pool state
     req_pool_idx: Optional[int] = None
     kv_committed_len: int = 0
-    kv: Optional[ReqKvInfo] = None
+    kv: ReqKvInfo = field(default_factory=ReqKvInfo)
 
     # First req's radix tree node (for dec_lock_ref on session close)
     last_node: Any = None
@@ -67,7 +68,7 @@ class SessionSlot:
     @property
     def is_holding_kv(self) -> bool:
         """Whether this slot currently holds KV pool resources."""
-        return self.kv is not None
+        return self.req_pool_idx is not None
 
     def save_from_req(self, req: Req, is_first: bool):
         """Save KV state from a finishing request into this slot."""
@@ -97,7 +98,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.kv = ReqKvInfo()
         req.mamba_pool_idx = None
         req.mamba_ping_pong_track_buffer = None
         req.mamba_next_track_idx = None
@@ -221,7 +222,7 @@ class StreamingSession(BasePrefixCache):
         if not _is_streaming(req):
             return None
         slot = self.slots.get(req.session.session_id)
-        if slot is None or slot.kv is None:
+        if slot is None or not slot.is_holding_kv:
             return None
         if req.to_finish is not None:
             req.session.abort_req()
@@ -350,7 +351,7 @@ class StreamingSession(BasePrefixCache):
             )
             self.release_session(session_id)
             req.req_pool_idx = None
-            req.kv = None
+            req.kv = ReqKvInfo()
             req.session.abort_req()
             return True
 

python/sglang/test/scripted_runtime/req_handle.py

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