Fix bookkeeping fields not encapsulated with real allocations in normal alloc, PD pre-alloc, DFlash and EAGLE (#29432)

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -85,6 +85,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
             rid="req-0",
             origin_input_ids=list(range(fill_len)),
             output_ids=[],
+            kv=None,
         )
 
         def set_extend_range(start, end):
@@ -143,8 +144,8 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
         _, kwargs = allocator.alloc_extend_swa_tail.call_args
         self.assertEqual(kwargs["extend_num_tokens"], fill_len)
         self.assertEqual(kwargs["swa_tail_len"], swa_tail_len)
-        self.assertEqual(req.swa_evicted_seqlen, fill_len - swa_tail_len)
-        self.assertEqual(req.kv_allocated_len, fill_len)
+        self.assertEqual(req.kv.swa_evicted_seqlen, fill_len - swa_tail_len)
+        self.assertEqual(req.kv.kv_allocated_len, fill_len)
         self.assertEqual(req.kv_committed_len, fill_len)
         self.assertEqual(req.extend_range.length, fill_len)
         self.assertEqual(len(req_to_token_pool.writes), 1)

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -4071,7 +4071,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
         req_to_token_pool.alloc([req])
         req.output_ids = array("q")
         req.kv_committed_len = len(tokens)
-        req.kv_allocated_len = len(tokens)
+        req.kv = ReqKvInfo(kv_allocated_len=len(tokens), swa_evicted_seqlen=0)
         req.cache_protected_len = 0
         req.swa_uuid_for_lock = None
         req.extra_key = None
@@ -4085,7 +4085,9 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
         req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
         req.last_node = cache.root_node
 
-        cache.cache_finished_req(req, is_insert=True)
+        cache.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len()
+        )
 
     def test_finished_req_stores_radix_mamba_state_in_int8_pool(self):
         cache, allocator, req_to_token_pool = build_fixture(self.cfg)

test/registered/unit/spec/test_decode_bookkeeping_ownership.py

@@ -50,32 +50,36 @@ _OWNER_SITES = {
     # non-spec scheduler
     (_SB, "ScheduleBatch.prepare_for_decode", "decode_batch_idx"): 1,
     (_SB, "ScheduleBatch.prepare_for_decode", "kv_committed_len"): 1,
-    (_SB, "ScheduleBatch.prepare_for_decode", "kv_allocated_len"): 1,
     (_SB, "ScheduleBatch.prepare_for_extend", "extend_batch_idx"): 1,
     (_SB, "ScheduleBatch.prepare_for_extend", "kv_committed_len"): 1,
-    (_SB, "ScheduleBatch.prepare_for_extend", "kv_allocated_len"): 1,
+    # kv_allocated_len is settled inside the owned-kv alloc functions (op28).
     ("mem_cache/allocation.py", "alloc_for_extend", "evict"): 1,
+    ("mem_cache/allocation.py", "alloc_for_extend", "kv_allocated_len"): 1,
     ("mem_cache/allocation.py", "alloc_for_decode", "evict"): 1,
+    ("mem_cache/allocation.py", "alloc_for_decode", "kv_allocated_len"): 1,
     # spec v2: no pre-claim; resolve commits the full accepted run uniformly.
+    # kv_allocated_len for spec v2 draft decode (eagle + dflash) is settled
+    # inside the owned-kv alloc_for_spec_decode function (op42).
     (*_EAGLE_DECODE, "decode_batch_idx"): 1,
     (*_EAGLE_DECODE, "evict"): 1,
-    (*_EAGLE_DECODE, "kv_allocated_len"): 1,
-    (*_RESOLVE, "kv_committed_len"): 1,
-    (*_RESOLVE, "spec_verify_ct"): 1,
     (
-        "speculative/dflash_info_v2.py",
-        "DFlashDraftInputV2.prepare_for_decode",
+        "mem_cache/allocation.py",
+        "alloc_for_spec_decode",
         "kv_allocated_len",
     ): 1,
-    # disaggregation decode prealloc
+    (*_RESOLVE, "kv_committed_len"): 1,
+    (*_RESOLVE, "spec_verify_ct"): 1,
+    # disaggregation decode prealloc: kv_allocated_len is settled inside the
+    # owned-kv alloc_for_decode_prealloc(_hisparse) functions (op42).
     (
         "disaggregation/decode.py",
         "DecodePreallocQueue._pre_alloc",
         "kv_committed_len",
     ): 1,
+    ("disaggregation/decode.py", "alloc_for_decode_prealloc", "kv_allocated_len"): 1,
     (
         "disaggregation/decode.py",
-        "DecodePreallocQueue._pre_alloc",
+        "alloc_for_decode_prealloc_hisparse",
         "kv_allocated_len",
     ): 1,
     # streaming session slot save/restore and tail trimming