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

test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py

@@ -8,7 +8,6 @@ Requires: torch, sglang (run in an environment with sglang installed)
 """
 
 import unittest
-from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import torch
@@ -18,6 +17,7 @@ from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
 )
 from sglang.srt.disaggregation.kv_events import OffloadedState
 from sglang.srt.managers.cache_controller import HiCacheAck
+from sglang.srt.managers.schedule_batch import ReqKvInfo
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=8, suite="base-a-test-cpu")
@@ -35,7 +35,7 @@ def _make_mock_req(
     req.rid = rid
     req.req_pool_idx = req_pool_idx
     req.kv_committed_len = kv_committed_len
-    req.kv = SimpleNamespace(kv_allocated_len=kv_allocated_len)
+    req.kv = ReqKvInfo(kv_allocated_len=kv_allocated_len)
     req.prefix_indices = list(range(prefix_indices_len))
     req.effective_kv_committed_len = lambda: req.kv_committed_len
     return req

test/registered/unit/managers/test_kv_page_invariants.py

@@ -49,6 +49,7 @@ class _FakeReq:
         self.req_pool_idx = rpi
         self.kv_committed_len = committed
         self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
+        self.is_holding_kv = True
 
 
 class _FakeSlot:

test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py

@@ -65,10 +65,8 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False):
         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 None
+        kv=SimpleNamespace(
+            kv_allocated_len=len(prefix) + block_size if req_pool_idx is not None else 0
         ),
     )
 

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -6,6 +6,7 @@ from unittest.mock import MagicMock
 import numpy as np
 import torch
 
+from sglang.srt.managers.schedule_batch import ReqKvInfo
 from sglang.srt.mem_cache.allocator.hisparse import (
     DeepSeekV4HiSparseTokenToKVPoolAllocator,
 )
@@ -91,7 +92,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
             rid="req-0",
             origin_input_ids=list(range(fill_len)),
             output_ids=[],
-            kv=None,
+            kv=ReqKvInfo(),
         )
 
         def set_extend_range(start, end):

test/registered/unit/mem_cache/test_swa_eviction_boundary.py

@@ -104,6 +104,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
     """Mock Req with fields needed by _evict_swa and cache_finished_req."""
     req = SimpleNamespace(
         req_pool_idx=req_pool_idx,
+        is_holding_kv=True,
         origin_input_ids=token_ids,
         output_ids=[],
         cache_protected_len=cache_protected_len,

test/registered/unit/mem_cache/test_unified_radix_cache_bench.py

@@ -240,7 +240,7 @@ def create_bench_cache(
     _rid = [0]
 
     def make_req():
-        from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
+        from sglang.srt.managers.schedule_batch import Req
         from sglang.srt.sampling.sampling_params import SamplingParams
 
         req = Req(
@@ -251,8 +251,6 @@ def create_bench_cache(
         )
         _rid[0] += 1
         req_to_token_pool.alloc([req])
-        # fabricated reqs bypass alloc_for_extend, the normal creator of req.kv
-        req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
         return req
 
     return tree, allocator, req_to_token_pool, make_req

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -24,7 +24,7 @@ from sglang.srt.disaggregation.kv_events import (
     StorageMedium,
 )
 from sglang.srt.environ import envs
-from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
+from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import (
@@ -1030,7 +1030,6 @@ class UnifiedRadixCacheSuite:
         )
         self._rid += 1
         req_to_token_pool.alloc([req])
-        req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
         return req
 
     def _apply_match_to_req(self, req, match):
@@ -1265,7 +1264,7 @@ class UnifiedRadixCacheSuite:
         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 = ReqKvInfo(kv_allocated_len=kv_len, swa_evicted_seqlen=0)
+        req.kv.kv_allocated_len = kv_len
         req.last_node = cache.root_node_handle()
         req.cache_protected_len = 0
         req.swa_uuid_for_lock = None
@@ -2200,7 +2199,6 @@ class UnifiedRadixCacheSuite:
         req.cache_protected_len = 0
         req.swa_uuid_for_lock = None
         req.extra_key = None
-        req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
 
         swa_avail_before = allocator.swa_attn_allocator.available_size()
 
@@ -2290,7 +2288,6 @@ class UnifiedRadixCacheSuite:
         req.cache_protected_len = 0
         req.swa_uuid_for_lock = None
         req.extra_key = None
-        req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
 
         with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
             cache.cache_unfinished_req(req)
@@ -6603,7 +6600,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
         req_to_token_pool.alloc([req])
         req.output_ids = array("q")
         req.kv_committed_len = len(tokens)
-        req.kv = ReqKvInfo(kv_allocated_len=len(tokens), swa_evicted_seqlen=0)
+        req.kv.kv_allocated_len = len(tokens)
         req.cache_protected_len = 0
         req.swa_uuid_for_lock = None
         req.extra_key = None
@@ -7934,7 +7931,6 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
         req.cache_protected_len = 0
         req.swa_uuid_for_lock = None
         req.extra_key = None
-        req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
 
         with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
             cache.cache_unfinished_req(req)