Preserve CP HiCache valid tails while padding physical pages

CP HiCache now keeps radix and scheduler-visible lengths as valid tokens while host/device transfers reserve and replay the padded physical page span. Exact valid-tail write, insertion, and match paths no longer fall back to page-flooring; the physical owner-lane contract still uses padded page metadata.

Constraint: Scheduler prefix indices must never include padded tail locs.
Constraint: Host/device transfer and owner-lane admission remain page-based.
Rejected: Pad to cp_size or 2*cp_size pages | wastes KV and recreates short-tail fallback behavior.
Rejected: Expose padded locs through load_cp return | would leak fake tokens into req.prefix_indices.
Confidence: medium
Scope-risk: moderate
Directive: Do not implement split-inside-tail by duplicating page_owners without a page-sharing/refcount design.
Tested: local py_compile for touched CP HiCache/radix/controller files and tests.
Tested: remote g0034 CP HiCache impacted suites: 143 passed, 5 warnings.
Tested: remote g0034 CP shared KV C1-C5 suite: 122 passed, 5 warnings.
Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson/starlette.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>

test/registered/unit/managers/test_hicache_controller_cp.py

@@ -701,15 +701,20 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
         self.assertEqual(host_pool.backups, [])
         self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
 
-    def test_cp_write_rejects_incomplete_owned_physical_page(self):
+    def test_cp_write_accepts_valid_tail_and_pads_owned_physical_page(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         controller = self.make_controller(host_pool, cp_rank=1)
         logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
 
-        with self.assertRaisesRegex(
-            ValueError, "_write_cp expects page-aligned device_indices"
-        ):
-            controller.write(logical_locs, node_id=21)
+        result = controller.write(logical_locs, node_id=21)
+
+        self.assertEqual(result.metadata.logical_len, 3)
+        self.assertEqual(result.metadata.valid_len, 3)
+        self.assertEqual(result.metadata.padded_len, 4)
+        self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3])
+        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
+        self.assertEqual(host_pool.alloc_calls, [4])
+        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
 
     def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
@@ -1052,6 +1057,29 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
         self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
 
+    def test_cp_load_returns_valid_locs_while_transferring_padded_tail_page(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 72, dtype=torch.int64))
+        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
+        node = TreeNode()
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            padded_len=8,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            page_owners=torch.tensor([3, 0], dtype=torch.int8),
+            page_size=4,
+        )
+
+        device_indices = controller.load_cp([node], node_id=112)
+        controller.start_loading()
+
+        self.assertEqual(device_indices.tolist(), list(range(64, 70)))
+        self.assertEqual(allocator.owner_alloc_calls, [[3, 0]])
+        self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
+        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
+
     def test_cp_load_frees_unexpected_owner_allocator_length(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))

test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py

@@ -232,6 +232,26 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase):
         self.assertEqual(plan.deficit_by_owner, [0, 1, 0, 0])
         self.assertEqual(plan.host_hit_len, 16)
 
+    def test_load_back_plan_keeps_valid_hit_len_while_planning_padded_pages(self):
+        allocator = _make_allocator(page_size=4, cp_size=4)
+        cache = _make_cache(allocator)
+        node = TreeNode(id=12)
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            padded_len=8,
+            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
+            host_indices=torch.arange(4, dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=4,
+        )
+
+        plan = cache._build_cp_load_back_plan([node], node_id=node.id)
+
+        self.assertEqual(plan.host_hit_len, 6)
+        self.assertEqual(plan.page_owners, [0, 1])
+        self.assertEqual(plan.required_by_owner, [1, 1, 0, 0])
+
     def test_load_back_plan_fails_closed_without_cp_metadata(self):
         allocator = _make_allocator()
         cache = _make_cache(allocator)

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -109,7 +109,7 @@ from sglang.srt.mem_cache.hiradix_cache import (
     PreparedCpHiCacheBackup,
     _compute_shared_hicache_token_capacities,
 )
-from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
+from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, _key_match_paged
 from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -199,6 +199,11 @@ class TestHiRadixCacheCPDraftHostPool(CustomTestCase):
 
 
 class TestCpHiCacheNodeMetadata(CustomTestCase):
+    def test_paged_key_match_returns_valid_tail_length_not_next_page(self):
+        key = RadixKey(list(range(6)))
+
+        self.assertEqual(_key_match_paged(key, key, page_size=4), 6)
+
     def test_split_zero_len_moves_all_positions_to_child(self):
         metadata = CpHiCacheNodeMetadata(
             logical_len=8,
@@ -268,6 +273,21 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
         self.assertEqual(metadata.owned_positions.dtype, torch.int64)
         self.assertEqual(metadata.host_indices.dtype, torch.int64)
 
+    def test_valid_length_can_be_shorter_than_physical_padded_length(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=100,
+            padded_len=128,
+            owned_positions=torch.tensor([0, 63, 100, 127], dtype=torch.int64),
+            host_indices=torch.tensor([10, 11, 12, 13], dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=64,
+        )
+
+        self.assertEqual(metadata.logical_len, 100)
+        self.assertEqual(metadata.valid_len, 100)
+        self.assertEqual(metadata.padded_len, 128)
+        self.assertEqual(metadata.page_owners.tolist(), [0, 1])
+
     def test_non_int64_inputs_are_converted(self):
         metadata = CpHiCacheNodeMetadata(
             logical_len=4,
@@ -1162,6 +1182,80 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
             [{"catch_up_all_layers": False}],
         )
 
+    def test_prepare_write_backup_for_req_keeps_valid_tail_length(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.disable = False
+        cache._uses_cp_hicache = True
+        cache.is_eagle = False
+        cache.page_size = 4
+        cache.req_to_token_pool = types.SimpleNamespace(
+            req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
+        )
+        cache.cache_controller = FakeReserveWriteController(
+            [
+                lambda device_indices, node_id: make_write_reservation(
+                    device_indices, node_id=node_id, host_start=170
+                )
+            ]
+        )
+        cache.maybe_bigram_convert = lambda key: (key, None)
+        cache.root_node = TreeNode()
+        cache.root_node.children = {}
+        cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
+
+        req = types.SimpleNamespace(
+            rid="rid-tail-prepare",
+            fill_ids=list(range(6)),
+            cache_protected_len=0,
+            req_pool_idx=0,
+            is_chunked=0,
+            cp_hicache_prepared_backup=None,
+        )
+
+        cache.prepare_write_backup_for_req(req)
+
+        self.assertIsNotNone(req.cp_hicache_prepared_backup)
+        self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 6)
+        self.assertEqual(cache.cache_controller.reservations[0][0].tolist(), list(range(6)))
+
+    def test_cache_finished_req_keeps_cp_valid_tail_insert_key(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.disable_finished_insert = False
+        cache.disable = False
+        cache._uses_cp_hicache = True
+        cache.is_eagle = False
+        cache.page_size = 4
+        cache.req_to_token_pool = types.SimpleNamespace(
+            req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
+        )
+        freed = []
+        cache.token_to_kv_pool_allocator = types.SimpleNamespace(
+            free=lambda indices: freed.append(indices.clone())
+        )
+        inserted = []
+        cache.insert = lambda params: inserted.append(params) or types.SimpleNamespace(
+            prefix_len=0
+        )
+        cache.dec_lock_ref = lambda node: None
+
+        req = types.SimpleNamespace(
+            origin_input_ids=list(range(6)),
+            output_ids=[],
+            req_pool_idx=0,
+            extra_key=None,
+            cache_protected_len=0,
+            last_node=TreeNode(),
+            cp_hicache_prepared_backup=None,
+            pop_committed_kv_cache=lambda: 6,
+        )
+
+        cache.cache_finished_req(req)
+
+        self.assertEqual(len(inserted), 1)
+        self.assertEqual(inserted[0].key.token_ids, list(range(6)))
+        self.assertEqual(inserted[0].value.tolist(), list(range(6)))
+        self.assertEqual([indices.tolist() for indices in freed], [[], []])
+
     def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache.disable = False
@@ -2136,6 +2230,49 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         self.assertIs(result.last_host_node, cache.root_node)
         self.assertEqual(result.host_hit_length, 0)
 
+    def test_cp_match_prefix_reports_valid_tail_host_hit(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.device = "cpu"
+        cache.disable = False
+        cache.page_size = 4
+        cache.ongoing_write_through = {}
+        cache.pending_host_backups = {}
+        cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
+        cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
+        cache.key_match_fn = lambda child_key, key: _key_match_paged(
+            child_key, key, page_size=4
+        )
+        cache.maybe_bigram_convert = lambda key: (key, None)
+        root = TreeNode()
+        root.key = RadixKey([])
+        root.value = torch.empty((0,), dtype=torch.int64)
+        root.host_len = 0
+        cache.root_node = root
+        node = TreeNode()
+        node.id = 140
+        node.parent = root
+        node.key = RadixKey(list(range(6)))
+        node.value = None
+        node.host_value = None
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            padded_len=8,
+            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
+            host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=4,
+        )
+        root.children[(0, 1, 2, 3)] = node
+
+        result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6)))))
+
+        self.assertEqual(result.device_indices.tolist(), [])
+        self.assertEqual(result.host_hit_length, 6)
+        self.assertIs(result.last_device_node, root)
+        self.assertIs(result.last_host_node, node)
+
     def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = False