Preserve draft KV across CP HiCache hits

Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix.  CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool.

Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV.

Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank.

Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint.

Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance.

Confidence: medium

Scope-risk: moderate

Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length.

Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings

Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery.

test/registered/unit/managers/test_hicache_controller_cp.py

@@ -5,8 +5,72 @@ from unittest.mock import patch
 
 import torch
 
-if "sgl_kernel" not in sys.modules:
-    sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
+try:
+    import sgl_kernel  # noqa: F401
+    import sgl_kernel.kvcacheio  # noqa: F401
+except (ImportError, RuntimeError):
+    if "sgl_kernel" not in sys.modules:
+        sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
+    sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
+    sys.modules["sgl_kernel"].__path__ = []
+    if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):
+
+        def _sgl_kernel_getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *args, **kwargs: None
+            setattr(sys.modules["sgl_kernel"], name, fn)
+            return fn
+
+        sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
+    if "sgl_kernel.quantization" not in sys.modules:
+        quantization_stub = types.ModuleType("sgl_kernel.quantization")
+        quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"
+
+        def _quantization_getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *args, **kwargs: None
+            setattr(quantization_stub, name, fn)
+            return fn
+
+        quantization_stub.__getattr__ = _quantization_getattr
+        for name in (
+            "ggml_dequantize",
+            "ggml_moe_a8",
+            "ggml_moe_a8_vec",
+            "ggml_moe_get_block_size",
+            "ggml_mul_mat_a8",
+            "ggml_mul_mat_vec_a8",
+        ):
+            setattr(quantization_stub, name, lambda *args, **kwargs: None)
+        sys.modules["sgl_kernel.quantization"] = quantization_stub
+    for name in (
+        "sgl_per_token_group_quant_8bit",
+        "sgl_per_token_group_quant_fp8",
+        "sgl_per_token_quant_fp8",
+        "fp8_blockwise_scaled_mm",
+        "fp8_scaled_mm",
+        "silu_and_mul",
+    ):
+        if not hasattr(sys.modules["sgl_kernel"], name):
+            setattr(sys.modules["sgl_kernel"], name, lambda *args, **kwargs: None)
+    _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
+    for _schema in (
+        "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+        "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+        "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
+        "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
+        "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+    ):
+        try:
+            _sgl_kernel_lib.define(_schema)
+        except RuntimeError as exc:
+            if (
+                "already" not in str(exc).lower()
+                and "duplicate" not in str(exc).lower()
+            ):
+                raise
 if "sgl_kernel.kvcacheio" not in sys.modules:
     kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio")
     for name in (
@@ -56,12 +120,14 @@ class FakeHostPool:
     def backup_from_device_all_layer(
         self, device_pool, host_indices, device_indices, io_backend
     ):
-        self.backups.append((host_indices.clone(), device_indices.clone()))
+        self.backups.append((host_indices.clone(), device_indices.clone(), device_pool))
 
     def load_to_device_per_layer(
         self, device_pool, host_indices, device_indices, layer_id, io_backend
     ):
-        self.loads.append((host_indices.clone(), device_indices.clone(), layer_id))
+        self.loads.append(
+            (host_indices.clone(), device_indices.clone(), layer_id, device_pool)
+        )
 
     def free(self, indices):
         self.frees.append(indices.clone())
@@ -72,6 +138,10 @@ class FakeDevicePool:
     device = "cpu"
     layer_num = 1
 
+    def __init__(self, name="target", layer_num=1):
+        self.name = name
+        self.layer_num = layer_num
+
     def register_layer_transfer_counter(self, counter):
         self.counter = counter
 
@@ -80,11 +150,13 @@ class FakeAllocator:
     def __init__(self, alloc_result=None):
         self.alloc_result = alloc_result
         self.alloc_calls = []
+        self.frees = []
         self.cp_size = 4
         self.cp_rank = 1
+        self.device_pool = FakeDevicePool()
 
     def get_kvcache(self):
-        return FakeDevicePool()
+        return self.device_pool
 
     def alloc(self, need_size):
         self.alloc_calls.append(need_size)
@@ -92,6 +164,10 @@ class FakeAllocator:
             return None
         return self.alloc_result[:need_size].clone()
 
+    def free(self, indices):
+        self.frees.append(indices.clone())
+        return len(indices)
+
 
 class HostIndicesTensor(torch.Tensor):
     @staticmethod
@@ -166,7 +242,14 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
         self.addCleanup(self.device_module_patcher.stop)
         self.addCleanup(self.nsa_pool_patcher.stop)
 
-    def make_controller(self, host_pool, allocator=None, cp_rank=1):
+    def make_controller(
+        self,
+        host_pool,
+        allocator=None,
+        cp_rank=1,
+        draft_host_pool=None,
+        draft_mem_pool_device=None,
+    ):
         allocator = allocator or FakeAllocator()
         controller = HiCacheController(
             token_to_kv_pool_allocator=allocator,
@@ -178,6 +261,8 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
             cp_shared_kv_layout=CpSharedKVLayout(
                 page_size=4, cp_size=4, cp_rank=cp_rank
             ),
+            draft_mem_pool_host=draft_host_pool,
+            draft_mem_pool_device=draft_mem_pool_device,
         )
         controller.layer_done_counter = DummyLayerDoneCounter()
         return controller
@@ -200,7 +285,7 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
         controller = self.make_controller(host_pool, cp_rank=1)
         logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
 
-        with self.assertRaisesRegex(ValueError, "physical_device_indices.*whole pages"):
+        with self.assertRaisesRegex(ValueError, "host_indices.*whole pages"):
             controller.write(logical_locs, node_id=21)
 
     def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
@@ -242,6 +327,53 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
 
         self.assertEqual(result.required_host_slots, 4)
 
+    def test_cp_write_with_draft_pool_backs_target_and_draft_locs(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        draft_host_pool = FakeHostPool(
+            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
+        )
+        draft_device_pool = FakeDevicePool("draft")
+        controller = self.make_controller(
+            host_pool,
+            cp_rank=1,
+            draft_host_pool=draft_host_pool,
+            draft_mem_pool_device=draft_device_pool,
+        )
+        logical_locs = torch.arange(4, 20, dtype=torch.int64)
+
+        result = controller.write(logical_locs, node_id=77)
+
+        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
+        self.assertEqual(
+            result.metadata.draft_host_indices.tolist(), [200, 201, 202, 203]
+        )
+        self.assertEqual(host_pool.alloc_calls, [4])
+        self.assertEqual(draft_host_pool.alloc_calls, [4])
+        self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
+        self.assertEqual(draft_host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
+        self.assertIs(draft_host_pool.backups[0][2], draft_device_pool)
+        self.assertEqual(len(controller.ack_write_queue), 1)
+        self.assertEqual(controller.ack_write_queue[0].node_ids, [77])
+
+    def test_cp_write_draft_allocation_failure_rolls_back_target_host(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        draft_host_pool = FakeHostPool(None)
+        controller = self.make_controller(
+            host_pool,
+            cp_rank=1,
+            draft_host_pool=draft_host_pool,
+            draft_mem_pool_device=FakeDevicePool("draft"),
+        )
+        logical_locs = torch.arange(4, 20, dtype=torch.int64)
+
+        result = controller.write(logical_locs, node_id=78)
+
+        self.assertIsNone(result.metadata)
+        self.assertEqual(result.required_host_slots, 4)
+        self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
+        self.assertEqual(host_pool.backups, [])
+        self.assertEqual(draft_host_pool.backups, [])
+
     def test_generate_storage_config_constructs_config_at_runtime(self):
         controller = HiCacheController.__new__(HiCacheController)
         controller.mem_pool_device = FakeDevicePool()
@@ -300,6 +432,39 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         self.assertEqual(allocator.alloc_calls, [16])
         self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
 
+    def test_cp_load_with_draft_pool_restores_target_and_draft_locs(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        draft_host_pool = FakeHostPool(
+            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
+        )
+        draft_device_pool = FakeDevicePool("draft")
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
+        controller = self.make_controller(
+            host_pool,
+            allocator=allocator,
+            cp_rank=1,
+            draft_host_pool=draft_host_pool,
+            draft_mem_pool_device=draft_device_pool,
+        )
+        node = TreeNode()
+        node.host_len = 16
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=16,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
+        )
+
+        device_indices = controller.load_cp([node], node_id=14)
+        controller.start_loading()
+
+        self.assertEqual(device_indices.tolist(), list(range(64, 80)))
+        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
+        self.assertEqual(draft_host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
+        self.assertIs(draft_host_pool.loads[0][3], draft_device_pool)
+        self.assertEqual(len(controller.ack_load_queue), 1)
+        self.assertEqual(controller.ack_load_queue[0].node_ids, [14])
+
     def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self):
         host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
         allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
@@ -313,6 +478,7 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         )
 
         device_indices = controller.load_cp([node], node_id=12)
+        controller.start_loading()
 
         self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
         self.assertEqual(host_pool.loads, [])

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -1,14 +1,80 @@
 import sys
+import types
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import torch
 
-# Stub out sgl_kernel before any sglang import so this CPU unit test does not
-# require CUDA extension libraries to be installed.
-for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
-    if _mod not in sys.modules:
-        sys.modules[_mod] = MagicMock()
+# Prefer the real sgl_kernel package when the test image provides it so custom
+# Torch operators are registered. Fall back to stubs on local CPU-only hosts.
+try:
+    import sgl_kernel  # noqa: F401
+    import sgl_kernel.kvcacheio  # noqa: F401
+except (ImportError, RuntimeError):
+    if "sgl_kernel" not in sys.modules:
+        sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
+    sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
+    sys.modules["sgl_kernel"].__path__ = []
+    if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):
+
+        def _sgl_kernel_getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *args, **kwargs: None
+            setattr(sys.modules["sgl_kernel"], name, fn)
+            return fn
+
+        sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
+    if "sgl_kernel.quantization" not in sys.modules:
+        quantization_stub = types.ModuleType("sgl_kernel.quantization")
+        quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"
+
+        def _quantization_getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *args, **kwargs: None
+            setattr(quantization_stub, name, fn)
+            return fn
+
+        quantization_stub.__getattr__ = _quantization_getattr
+        for _name in (
+            "ggml_dequantize",
+            "ggml_moe_a8",
+            "ggml_moe_a8_vec",
+            "ggml_moe_get_block_size",
+            "ggml_mul_mat_a8",
+            "ggml_mul_mat_vec_a8",
+        ):
+            setattr(quantization_stub, _name, lambda *args, **kwargs: None)
+        sys.modules["sgl_kernel.quantization"] = quantization_stub
+    for _name in (
+        "sgl_per_token_group_quant_8bit",
+        "sgl_per_token_group_quant_fp8",
+        "sgl_per_token_quant_fp8",
+        "fp8_blockwise_scaled_mm",
+        "fp8_scaled_mm",
+        "silu_and_mul",
+    ):
+        if not hasattr(sys.modules["sgl_kernel"], _name):
+            setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None)
+    _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
+    for _schema in (
+        "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+        "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+        "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
+        "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
+        "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+    ):
+        try:
+            _sgl_kernel_lib.define(_schema)
+        except RuntimeError as exc:
+            if (
+                "already" not in str(exc).lower()
+                and "duplicate" not in str(exc).lower()
+            ):
+                raise
+if "sgl_kernel.kvcacheio" not in sys.modules:
+    sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
 
 from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
 from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
@@ -70,6 +136,22 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
         self.assertEqual(child.owned_positions.tolist(), [0, 4])
         self.assertEqual(child.host_indices.tolist(), [22, 23])
 
+    def test_split_keeps_draft_host_indices_aligned_with_owned_positions(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=10,
+            owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
+            host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
+            draft_host_indices=torch.tensor([120, 121, 122, 123], dtype=torch.int64),
+        )
+
+        parent, child = metadata.split(5)
+
+        self.assertEqual(parent.host_indices.tolist(), [20, 21])
+        self.assertEqual(parent.draft_host_indices.tolist(), [120, 121])
+        self.assertEqual(child.owned_positions.tolist(), [0, 4])
+        self.assertEqual(child.host_indices.tolist(), [22, 23])
+        self.assertEqual(child.draft_host_indices.tolist(), [122, 123])
+
     def test_zero_owned_metadata_is_valid(self):
         metadata = CpHiCacheNodeMetadata(
             logical_len=64,
@@ -150,6 +232,15 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
                 host_indices=torch.tensor([9], dtype=torch.int64),
             )
 
+    def test_draft_host_length_mismatch_raises(self):
+        with self.assertRaisesRegex(ValueError, "draft_host_indices.*same length"):
+            CpHiCacheNodeMetadata(
+                logical_len=4,
+                owned_positions=torch.tensor([1, 2], dtype=torch.int64),
+                host_indices=torch.tensor([9, 10], dtype=torch.int64),
+                draft_host_indices=torch.tensor([109], dtype=torch.int64),
+            )
+
     def test_out_of_range_positions_raise(self):
         with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
             CpHiCacheNodeMetadata(
@@ -222,6 +313,11 @@ class FakeEvictDeviceController:
         return len(device_indices)
 
 
+class FakeTokenAllocator:
+    def available_size(self):
+        return 0
+
+
 class TestHiRadixCacheCPBackup(CustomTestCase):
     def test_node_backuped_uses_cp_metadata(self):
         cache = HiRadixCache.__new__(HiRadixCache)
@@ -341,6 +437,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         cache.evictable_leaves = set()
         cache.evictable_host_leaves = set()
         cache.eviction_strategy = FakeEvictionStrategy()
+        cache.token_to_kv_pool_allocator = FakeTokenAllocator()
         cache.evictable_size_ = 4
         cache.protected_size_ = 0
         cache._record_remove_event = lambda node: (_ for _ in ()).throw(
@@ -602,14 +699,18 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         cache.root_node.children[1] = node
         cache.evictable_host_leaves.add(node)
 
-        def mark_not_done(done_tensor, op=None, group=None):
-            self.assertIs(group, cache.tp_group)
-            done_tensor.fill_(0)
+        all_done_states = iter([False, True])
+        cache._cp_all_ranks_true = lambda done: next(all_done_states)
+        cache._cp_broadcast_node_ids = lambda node_ids, max_ids: node_ids[:max_ids]
+        cache._cp_filter_all_ranks_safe_node_ids = (
+            lambda node_ids, is_safe, **_kwargs: [
+                node_id
+                for node_id in node_ids
+                if is_safe(cache._cp_node_by_id(node_id))
+            ]
+        )
 
-        with patch("torch.distributed.all_reduce", side_effect=mark_not_done):
-            physical_freed = cache._evict_host_for_physical_slots(
-                0, synchronize_across_ranks=True
-            )
+        physical_freed = cache._cp_evict_host_for_physical_slots(0)
 
         self.assertEqual(physical_freed, 0)
         self.assertEqual(freed, [])