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.

python/sglang/srt/disaggregation/prefill.py

@@ -853,9 +853,18 @@ class SchedulerDisaggregationPrefillMixin:
             )
             return
         prefill_queue = getattr(self, "disagg_prefill_bootstrap_queue", None)
+        has_draft_pool = (
+            getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None
+        )
+        prefix_len = len(getattr(req, "prefix_indices", ()))
+        host_hit_length = int(getattr(req, "host_hit_length", 0) or 0)
+        draft_prefix_overlap = max(0, min(end_idx, prefix_len) - start_idx)
         _cp_draft_shared_kv_debug(
             "prefill_send_kv_chunk rid=%s room=%s start_idx=%s end_idx=%s "
-            "last_chunk=%s page_size=%s pages=%s state_pages=%s has_draft_pool=%s",
+            "last_chunk=%s page_size=%s pages=%s state_pages=%s "
+            "has_draft_pool=%s prefix_len=%s host_hit_length=%s "
+            "cache_protected_len=%s extend_input_len=%s fill_len=%s "
+            "origin_input_len=%s already_computed=%s draft_prefix_overlap=%s",
             req.rid,
             req.bootstrap_room,
             start_idx,
@@ -864,6 +873,27 @@ class SchedulerDisaggregationPrefillMixin:
             page_size,
             _seq_summary(page_indices),
             _seq_summary(state_indices),
-            getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None,
+            has_draft_pool,
+            prefix_len,
+            host_hit_length,
+            getattr(req, "cache_protected_len", None),
+            getattr(req, "extend_input_len", None),
+            len(req.fill_ids),
+            len(req.origin_input_ids),
+            getattr(req, "already_computed", None),
+            draft_prefix_overlap,
         )
+        if has_draft_pool and draft_prefix_overlap > 0:
+            _cp_draft_shared_kv_debug(
+                "prefill_send_cachehit_draft_prefix rid=%s room=%s "
+                "draft_prefix_overlap=%s prefix_len=%s host_hit_length=%s "
+                "start_idx=%s end_idx=%s note=transfer_reads_draft_pool_for_cached_prefix",
+                req.rid,
+                req.bootstrap_room,
+                draft_prefix_overlap,
+                prefix_len,
+                host_hit_length,
+                start_idx,
+                end_idx,
+            )
         req.disagg_kv_sender.send(page_indices, state_indices)

python/sglang/srt/managers/cache_controller.py

@@ -70,10 +70,11 @@ class LayerDoneCounter:
     def __init__(self, num_layers: int):
         self.num_layers = num_layers
         # extra producer and consumer counters for overlap mode
-        self.num_counters = 3
+        self.num_counters = 5
         self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)]
         self.producer_index = -1
         self.consumer_index = -1
+        self.consumer_indices: List[int] = []
 
     def update_producer(self):
         self.producer_index = (self.producer_index + 1) % self.num_counters
@@ -84,17 +85,26 @@ class LayerDoneCounter:
         )
         return self.producer_index
 
-    def set_consumer(self, index: int):
-        self.consumer_index = index
+    def set_consumer(self, index):
+        if isinstance(index, (list, tuple, set)):
+            self.consumer_indices = [int(i) for i in index if int(i) >= 0]
+            self.consumer_index = (
+                self.consumer_indices[0] if self.consumer_indices else -1
+            )
+            return
+        self.consumer_index = int(index)
+        self.consumer_indices = [self.consumer_index] if self.consumer_index >= 0 else []
 
     def wait_until(self, threshold: int):
-        if self.consumer_index < 0:
+        if not self.consumer_indices:
             return
-        self.events[self.consumer_index].wait(threshold)
+        for consumer_index in self.consumer_indices:
+            self.events[consumer_index].wait(threshold)
 
     def reset(self):
         self.producer_index = -1
         self.consumer_index = -1
+        self.consumer_indices = []
 
 
 class CacheOperation:
@@ -275,6 +285,8 @@ class HiCacheController:
         pp_size: int = 1,
         enable_storage_metrics: bool = False,
         cp_shared_kv_layout: Optional[CpSharedKVLayout] = None,
+        draft_mem_pool_host: Optional["HostKVCache"] = None,
+        draft_mem_pool_device=None,
     ):
         self.tp_group = tp_group
         self.mem_pool_device_allocator = token_to_kv_pool_allocator
@@ -285,6 +297,10 @@ class HiCacheController:
             mem_pool_device = mem_pool_device.full_kv_pool
         self.mem_pool_device = mem_pool_device
         self.cp_shared_kv_layout = cp_shared_kv_layout
+
+        self.has_draft = False
+        self.mem_pool_device_draft = None
+        self.mem_pool_host_draft = None
         self.uses_cp_hicache = cp_shared_kv_layout is not None
         if self.uses_cp_hicache and not isinstance(
             self.mem_pool_device, NSATokenToKVPool
@@ -317,20 +333,28 @@ class HiCacheController:
         self.layer_num = self.mem_pool_device.layer_num
         self.layer_done_counter = LayerDoneCounter(self.layer_num)
         self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter)
+        self.draft_mem_pool_host = None
+        self.draft_mem_pool_device = None
 
         if write_policy not in [
             "write_through",
             "write_through_selective",
             "write_back",
+            "write_behind",
         ]:
             raise ValueError(f"Invalid write policy: {write_policy}")
 
         # self.write_queue = PriorityQueue[CacheOperation]()
         self.load_queue: List[CacheOperation] = []
         self.write_queue: List[CacheOperation] = []
+        self.draft_load_queue: List[CacheOperation] = []
+        self.draft_write_queue: List[CacheOperation] = []
         self.ack_load_queue: List[HiCacheAck] = []
         self.ack_write_queue: List[HiCacheAck] = []
 
+        if draft_mem_pool_host is not None or draft_mem_pool_device is not None:
+            self.attach_draft_pool(draft_mem_pool_device, draft_mem_pool_host)
+
         self.stop_event = threading.Event()
         self.write_buffer = TransferBuffer(self.stop_event)
         self.load_buffer = TransferBuffer(
@@ -354,6 +378,27 @@ class HiCacheController:
                 # Preserve the historical error shape on init for unknown backends.
                 raise ValueError(f"Failed to create storage backend: {e}") from e
 
+    @property
+    def has_draft_hicache(self) -> bool:
+        return (
+            self.draft_mem_pool_host is not None
+            and self.draft_mem_pool_device is not None
+        )
+
+    def attach_draft_pool(self, draft_mem_pool_device, draft_mem_pool_host) -> None:
+        if draft_mem_pool_device is None and draft_mem_pool_host is None:
+            return
+        if draft_mem_pool_device is None or draft_mem_pool_host is None:
+            raise ValueError(
+                "draft_mem_pool_device and draft_mem_pool_host must be provided together"
+            )
+        self.draft_mem_pool_device = draft_mem_pool_device
+        self.draft_mem_pool_host = draft_mem_pool_host
+        if hasattr(draft_mem_pool_device, "register_layer_transfer_counter"):
+            draft_mem_pool_device.register_layer_transfer_counter(
+                self.layer_done_counter
+            )
+
     def _start_storage_threads(self):
         """Start storage prefetch/backup threads and their queues.
 
@@ -652,6 +697,8 @@ class HiCacheController:
 
         self.write_queue.clear()
         self.load_queue.clear()
+        self.draft_write_queue.clear()
+        self.draft_load_queue.clear()
         self.write_buffer.clear()
         self.load_buffer.clear()
         self.ack_write_queue.clear()
@@ -706,6 +753,10 @@ class HiCacheController:
         self.write_queue.append(
             CacheOperation(host_indices, device_indices, node_id, priority)
         )
+        if self.has_draft and not self.uses_cp_hicache:
+            self.draft_write_queue.append(
+                CacheOperation(host_indices, device_indices, node_id, priority)
+            )
         self.start_writing()
         logger.info(
             "[CacheCtrl-write] write non-CP submitted: node_id=%d len=%d",
@@ -715,12 +766,28 @@ class HiCacheController:
         return host_indices
 
     def start_writing(self) -> None:
-        if len(self.write_queue) == 0:
+        if len(self.write_queue) == 0 and len(self.draft_write_queue) == 0:
             return
 
-        op = CacheOperation.merge_ops(self.write_queue)
-        host_indices, device_indices = self.move_indices(op)
+        op = CacheOperation.merge_ops(self.write_queue) if self.write_queue else None
+        draft_op = (
+            CacheOperation.merge_ops(self.draft_write_queue)
+            if self.draft_write_queue
+            else None
+        )
+        node_ids = op.node_ids if op is not None else draft_op.node_ids
+        if op is not None:
+            host_indices, device_indices = self.move_indices(op, self.mem_pool_host)
+        else:
+            host_indices = device_indices = None
+        if draft_op is not None:
+            draft_host_indices, draft_device_indices = self.move_indices(
+                draft_op, self.draft_mem_pool_host
+            )
+        else:
+            draft_host_indices = draft_device_indices = None
         self.write_queue.clear()
+        self.draft_write_queue.clear()
 
         start_event = device_module.Event()
         finish_event = device_module.Event()
@@ -728,19 +795,31 @@ class HiCacheController:
         start_event.record()
         with device_module.stream(self.write_stream):
             start_event.wait(self.write_stream)
-            self.mem_pool_host.backup_from_device_all_layer(
-                self.mem_pool_device, host_indices, device_indices, self.io_backend
-            )
+            if op is not None:
+                self.mem_pool_host.backup_from_device_all_layer(
+                    self.mem_pool_device, host_indices, device_indices, self.io_backend
+                )
+            if draft_op is not None:
+                self.draft_mem_pool_host.backup_from_device_all_layer(
+                    self.draft_mem_pool_device,
+                    draft_host_indices,
+                    draft_device_indices,
+                    self.io_backend,
+                )
             finish_event.record()
             # NOTE: We must save the host indices and device indices here,
             # this is because we need to guarantee that these tensors are
             # still alive when the write stream is executing.
-            if host_indices.is_cuda:
+            if host_indices is not None and host_indices.is_cuda:
                 host_indices.record_stream(self.write_stream)
-            if device_indices.is_cuda:
+            if device_indices is not None and device_indices.is_cuda:
                 device_indices.record_stream(self.write_stream)
+            if draft_host_indices is not None and draft_host_indices.is_cuda:
+                draft_host_indices.record_stream(self.write_stream)
+            if draft_device_indices is not None and draft_device_indices.is_cuda:
+                draft_device_indices.record_stream(self.write_stream)
 
-        self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids))
+        self.ack_write_queue.append(HiCacheAck(start_event, finish_event, node_ids))
 
     def _append_completed_write_ack(self, node_id: int) -> None:
         event = device_module.Event()
@@ -757,10 +836,10 @@ class HiCacheController:
         host_indices: torch.Tensor,
         device_indices: torch.Tensor,
     ) -> None:
-        validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
         validate_page_aligned_token_indices(
             device_indices, self.page_size, "physical_device_indices"
         )
+        validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
 
     def _write_cp(
         self,
@@ -786,6 +865,11 @@ class HiCacheController:
                     logical_len=logical_len,
                     owned_positions=owned_positions,
                     host_indices=torch.empty((0,), dtype=torch.int64),
+                    draft_host_indices=(
+                        torch.empty((0,), dtype=torch.int64)
+                        if self.has_draft_hicache
+                        else None
+                    ),
                 )
             )
 
@@ -800,31 +884,75 @@ class HiCacheController:
                 owned_positions.numel(),
             )
             return HiCacheWriteFailure(required_host_slots=len(physical_device_indices))
+        draft_host_indices = None
+        if self.has_draft_hicache:
+            draft_host_indices = self.draft_mem_pool_host.alloc(
+                len(physical_device_indices)
+            )
+            if draft_host_indices is None:
+                self.mem_pool_host.free(host_indices)
+                logger.info(
+                    "[CacheCtrl-write] _write_cp FAILED (draft host full): node_id=%d logical_len=%d owned=%d",
+                    node_id,
+                    logical_len,
+                    owned_positions.numel(),
+                )
+                return HiCacheWriteFailure(
+                    required_host_slots=len(physical_device_indices)
+                )
         try:
             self._validate_cp_hicache_page_indices(host_indices, physical_device_indices)
+            if draft_host_indices is not None:
+                self._validate_cp_hicache_page_indices(
+                    draft_host_indices, physical_device_indices
+                )
         except Exception:
             self.mem_pool_host.free(host_indices)
+            if draft_host_indices is not None:
+                self.draft_mem_pool_host.free(draft_host_indices)
             raise
 
         self.write_queue.append(
             CacheOperation(host_indices, physical_device_indices, node_id, priority)
         )
+        if draft_host_indices is not None:
+            self.draft_write_queue.append(
+                CacheOperation(
+                    draft_host_indices, physical_device_indices, node_id, priority
+                )
+            )
         self.start_writing()
         logger.info(
-            "[CacheCtrl-write] _write_cp submitted: node_id=%d logical_len=%d owned=%d physical=%d",
+            "[CacheCtrl-write] _write_cp submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
             node_id,
             logical_len,
             owned_positions.numel(),
             len(physical_device_indices),
+            draft_host_indices is not None,
         )
         return HiCacheWriteResult(
             metadata=CpHiCacheNodeMetadata(
                 logical_len=logical_len,
                 owned_positions=owned_positions,
                 host_indices=host_indices.cpu(),
+                draft_host_indices=(
+                    draft_host_indices.cpu() if draft_host_indices is not None else None
+                ),
             )
         )
 
+    def set_draft_kv_pool(self, draft_device_pool, draft_host_pool) -> None:
+        """Register draft KV pools so L2 ops piggyback draft transfers."""
+        self.has_draft = True
+        self.mem_pool_device_draft = draft_device_pool
+        self.mem_pool_host_draft = draft_host_pool
+        self.attach_draft_pool(draft_device_pool, draft_host_pool)
+        logger.info(
+            "HiCache draft KV registered: %s (host %d slots)",
+            type(draft_device_pool).__name__,
+            draft_host_pool.size,
+        )
+
     def load(
         self,
         host_indices: torch.Tensor,
@@ -840,6 +968,10 @@ class HiCacheController:
         self.load_queue.append(
             CacheOperation(host_indices, device_indices, node_id, priority)
         )
+        if self.has_draft and not self.uses_cp_hicache:
+            self.draft_load_queue.append(
+                CacheOperation(host_indices, device_indices, node_id, priority)
+            )
         return device_indices
 
     def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]:
@@ -849,6 +981,7 @@ class HiCacheController:
             return None
 
         host_chunks = []
+        draft_host_chunks = []
         physical_chunks = []
         offset = 0
         for node in nodes_to_load:
@@ -862,9 +995,31 @@ class HiCacheController:
                 self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
             )
             host_chunks.append(node.cp_hicache.host_indices)
+            if self.has_draft_hicache:
+                draft_host_indices = getattr(
+                    node.cp_hicache, "draft_host_indices", None
+                )
+                if draft_host_indices is None:
+                    self.mem_pool_device_allocator.free(device_indices)
+                    raise RuntimeError(
+                        "CP HiCache draft KV restore requested but node metadata "
+                        "does not contain draft_host_indices"
+                    )
+                draft_host_chunks.append(draft_host_indices)
 
         if not host_chunks:
-            self._append_completed_load_ack(node_id)
+            # Keep CP load ACK rows identical across ranks.  A zero-owned rank
+            # still queues a zero-length op with the logical node id so a later
+            # batched start_loading() merges the same node_ids as owning ranks.
+            self.load_queue.append(
+                CacheOperation(
+                    torch.empty((0,), dtype=torch.int64),
+                    torch.empty(
+                        (0,), dtype=device_indices.dtype, device=device_indices.device
+                    ),
+                    node_id,
+                )
+            )
             return device_indices
 
         host_indices = torch.cat(host_chunks)
@@ -874,6 +1029,17 @@ class HiCacheController:
         except Exception:
             self.mem_pool_device_allocator.free(device_indices)
             raise
+        draft_host_indices = None
+        if self.has_draft_hicache:
+            draft_host_indices = torch.cat(draft_host_chunks)
+            try:
+                self._validate_cp_hicache_page_indices(
+                    draft_host_indices, physical_device_indices
+                )
+            except Exception:
+                self.mem_pool_device_allocator.free(device_indices)
+                raise
+
         self.load_queue.append(
             CacheOperation(
                 host_indices,
@@ -881,10 +1047,19 @@ class HiCacheController:
                 node_id,
             )
         )
+        if draft_host_indices is not None:
+            self.draft_load_queue.append(
+                CacheOperation(
+                    draft_host_indices,
+                    physical_device_indices,
+                    node_id,
+                )
+            )
 
         return device_indices
 
-    def move_indices(self, op: CacheOperation):
+    def move_indices(self, op: CacheOperation, mem_pool_host=None):
+        mem_pool_host = mem_pool_host or self.mem_pool_host
         host_indices, device_indices = op.host_indices, op.device_indices
         # move indices to GPU if using kernels, to host if using direct indexing
         if self.io_backend == "kernel":
@@ -892,15 +1067,15 @@ class HiCacheController:
                 host_indices = host_indices.to(self.device, non_blocking=True)
             return host_indices, device_indices
         elif self.io_backend == "direct":
-            if self.mem_pool_host.layout == "layer_first":
+            if mem_pool_host.layout == "layer_first":
                 device_indices = device_indices.cpu()
                 host_indices, idx = host_indices.sort()
                 return host_indices, device_indices.index_select(0, idx)
-            elif self.mem_pool_host.layout == "page_first_direct":
+            elif mem_pool_host.layout == "page_first_direct":
                 return host_indices, device_indices.cpu()
             else:
                 raise ValueError(
-                    f"Unsupported layout {self.mem_pool_host.layout!r} for io backend 'direct'"
+                    f"Unsupported layout {mem_pool_host.layout!r} for io backend 'direct'"
                 )
         elif self.io_backend == "kernel_ascend":
             return host_indices, device_indices.cpu()
@@ -908,40 +1083,76 @@ class HiCacheController:
             raise ValueError(f"Unsupported io backend")
 
     def start_loading(self) -> int:
-        if len(self.load_queue) == 0:
+        if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0:
             return -1
 
         producer_id = self.layer_done_counter.update_producer()
-        op = CacheOperation.merge_ops(self.load_queue)
-        host_indices, device_indices = self.move_indices(op)
+        op = CacheOperation.merge_ops(self.load_queue) if self.load_queue else None
+        draft_op = (
+            CacheOperation.merge_ops(self.draft_load_queue)
+            if self.draft_load_queue
+            else None
+        )
+        node_ids = op.node_ids if op is not None else draft_op.node_ids
+        if op is not None:
+            host_indices, device_indices = self.move_indices(op, self.mem_pool_host)
+        else:
+            host_indices = device_indices = None
+        if draft_op is not None:
+            draft_host_indices, draft_device_indices = self.move_indices(
+                draft_op, self.draft_mem_pool_host
+            )
+        else:
+            draft_host_indices = draft_device_indices = None
         self.load_queue.clear()
+        self.draft_load_queue.clear()
         producer_event = self.layer_done_counter.events[producer_id]
         producer_event.start_event.record()
 
         with device_module.stream(self.load_stream):
             producer_event.start_event.wait(self.load_stream)
-            for i in range(self.layer_num):
-                self.mem_pool_host.load_to_device_per_layer(
-                    self.mem_pool_device,
-                    host_indices,
-                    device_indices,
-                    i,
-                    self.io_backend,
-                )
-                producer_event.complete(i)
+            draft_layer_num = (
+                self.draft_mem_pool_device.layer_num if draft_op is not None else 0
+            )
+            for i in range(max(self.layer_num, draft_layer_num)):
+                if draft_op is not None and i < draft_layer_num:
+                    if len(draft_host_indices) > 0:
+                        self.draft_mem_pool_host.load_to_device_per_layer(
+                            self.draft_mem_pool_device,
+                            draft_host_indices,
+                            draft_device_indices,
+                            i,
+                            self.io_backend,
+                        )
+                if op is not None and i < self.layer_num:
+                    if len(host_indices) > 0:
+                        self.mem_pool_host.load_to_device_per_layer(
+                            self.mem_pool_device,
+                            host_indices,
+                            device_indices,
+                            i,
+                            self.io_backend,
+                        )
+                    producer_event.complete(i)
+                elif op is None and i < self.layer_num:
+                    producer_event.complete(i)
             # NOTE: We must save the host indices and device indices here,
             # this is because we need to guarantee that these tensors are
             # still alive when the load stream is executing.
-            if host_indices.is_cuda:
+            if host_indices is not None and host_indices.is_cuda:
                 host_indices.record_stream(self.load_stream)
-            if device_indices.is_cuda:
+            if device_indices is not None and device_indices.is_cuda:
                 device_indices.record_stream(self.load_stream)
+            if draft_host_indices is not None and draft_host_indices.is_cuda:
+                draft_host_indices.record_stream(self.load_stream)
+            if draft_device_indices is not None and draft_device_indices.is_cuda:
+                draft_device_indices.record_stream(self.load_stream)
 
         self.ack_load_queue.append(
             HiCacheAck(
                 start_event=producer_event.start_event,
                 finish_event=producer_event.finish_event,
-                node_ids=op.node_ids,
+                node_ids=node_ids,
             )
         )
         return producer_id
@@ -957,6 +1168,16 @@ class HiCacheController:
         self.mem_pool_host.free(host_indices)
         return len(host_indices)
 
+    def evict_cp_host(self, metadata, backup_only: bool = True) -> int:
+        if not backup_only:
+            raise ValueError("Other eviction policies are not supported yet.")
+
+        freed = self.evict_host(metadata.host_indices, backup_only=backup_only)
+        draft_host_indices = getattr(metadata, "draft_host_indices", None)
+        if self.has_draft_hicache and draft_host_indices is not None:
+            self.draft_mem_pool_host.free(draft_host_indices)
+        return freed
+
     def prefetch(
         self,
         request_id: str,

python/sglang/srt/managers/scheduler.py

@@ -966,6 +966,10 @@ class Scheduler(
             _cp_draft_pool_summary(draft_token_to_kv_pool),
             _cp_draft_pool_summary(self.token_to_kv_pool_allocator.get_kvcache()),
         )
+        if draft_token_to_kv_pool is not None and hasattr(
+            self.tree_cache, "attach_draft_kv_pool"
+        ):
+            self.tree_cache.attach_draft_kv_pool(draft_token_to_kv_pool)
 
         if (
             self.disaggregation_mode == DisaggregationMode.DECODE

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -57,6 +57,7 @@ class CpHiCacheNodeMetadata:
     logical_len: int
     owned_positions: torch.Tensor
     host_indices: torch.Tensor
+    draft_host_indices: Optional[torch.Tensor] = None
 
     def __post_init__(self):
         if self.logical_len < 0:
@@ -67,11 +68,23 @@ class CpHiCacheNodeMetadata:
         self.host_indices = self.host_indices.to(
             device="cpu", dtype=torch.int64
         ).detach().clone()
+        if self.draft_host_indices is not None:
+            self.draft_host_indices = self.draft_host_indices.to(
+                device="cpu", dtype=torch.int64
+            ).detach().clone()
         if self.owned_positions.numel() != self.host_indices.numel():
             raise ValueError(
                 "owned_positions and host_indices must have same length, got "
                 f"{self.owned_positions.numel()} and {self.host_indices.numel()}"
             )
+        if (
+            self.draft_host_indices is not None
+            and self.owned_positions.numel() != self.draft_host_indices.numel()
+        ):
+            raise ValueError(
+                "draft_host_indices and owned_positions must have same length, got "
+                f"{self.draft_host_indices.numel()} and {self.owned_positions.numel()}"
+            )
         if self.owned_positions.numel() > 0:
             if torch.any(self.owned_positions < 0) or torch.any(
                 self.owned_positions >= self.logical_len
@@ -96,19 +109,88 @@ class CpHiCacheNodeMetadata:
                 logical_len=split_len,
                 owned_positions=self.owned_positions[parent_mask],
                 host_indices=self.host_indices[parent_mask],
+                draft_host_indices=(
+                    self.draft_host_indices[parent_mask]
+                    if self.draft_host_indices is not None
+                    else None
+                ),
             ),
             CpHiCacheNodeMetadata(
                 logical_len=self.logical_len - split_len,
                 owned_positions=self.owned_positions[child_mask] - split_len,
                 host_indices=self.host_indices[child_mask],
+                draft_host_indices=(
+                    self.draft_host_indices[child_mask]
+                    if self.draft_host_indices is not None
+                    else None
+                ),
             ),
         )
 
 
 class HiRadixCache(RadixCache):
 
+    def _create_token_to_kv_pool_host(self, kv_cache, server_args: ServerArgs):
+        from sglang.srt.mem_cache.memory_pool_host import (
+            MHATokenToKVPoolHost,
+            MLATokenToKVPoolHost,
+            NSATokenToKVPoolHost,
+        )
+
+        if isinstance(kv_cache, MHATokenToKVPool):
+            return MHATokenToKVPoolHost(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+                allocator_type=server_args.hicache_storage_backend,
+            )
+        if isinstance(kv_cache, NSATokenToKVPool):
+            return NSATokenToKVPoolHost(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+                allocator_type=server_args.hicache_storage_backend,
+            )
+        if isinstance(kv_cache, MLATokenToKVPool):
+            return MLATokenToKVPoolHost(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+                allocator_type=server_args.hicache_storage_backend,
+            )
+        raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools")
+
+    def attach_draft_kv_pool(self, draft_token_to_kv_pool) -> None:
+        if not self._uses_cp_hicache or draft_token_to_kv_pool is None:
+            return
+        if (
+            getattr(self.cache_controller, "draft_mem_pool_device", None)
+            is draft_token_to_kv_pool
+        ):
+            return
+        draft_token_to_kv_pool_host = self._create_token_to_kv_pool_host(
+            draft_token_to_kv_pool, self._server_args
+        )
+        self.cache_controller.attach_draft_pool(
+            draft_token_to_kv_pool, draft_token_to_kv_pool_host
+        )
+        self.draft_token_to_kv_pool_host = draft_token_to_kv_pool_host
+        logger.info(
+            "[HiCache-draft] attached CP draft KV host pool: pool=%s host_pool=%s page_size=%d",
+            draft_token_to_kv_pool.__class__.__name__,
+            draft_token_to_kv_pool_host.__class__.__name__,
+            self.page_size,
+        )
+
     def __init__(self, params: CacheInitParams, server_args: ServerArgs):
         self._enable_metrics_flag = params.enable_metrics
+        self._server_args = server_args
 
         if not server_args.disable_hicache_numa_detect:
             bind_to_closest_numa_node_cuda()
@@ -130,43 +212,10 @@ class HiRadixCache(RadixCache):
                     "CP shared KV HiCache host integration requires NSATokenToKVPool."
                 )
         self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
-
-        from sglang.srt.mem_cache.memory_pool_host import (
-            MHATokenToKVPoolHost,
-            MLATokenToKVPoolHost,
-            NSATokenToKVPoolHost,
+        self.token_to_kv_pool_host = self._create_token_to_kv_pool_host(
+            self.kv_cache, server_args
         )
 
-        if isinstance(self.kv_cache, MHATokenToKVPool):
-            self.token_to_kv_pool_host = MHATokenToKVPoolHost(
-                self.kv_cache,
-                server_args.hicache_ratio,
-                server_args.hicache_size,
-                self.page_size,
-                server_args.hicache_mem_layout,
-                allocator_type=server_args.hicache_storage_backend,
-            )
-        elif isinstance(self.kv_cache, NSATokenToKVPool):
-            self.token_to_kv_pool_host = NSATokenToKVPoolHost(
-                self.kv_cache,
-                server_args.hicache_ratio,
-                server_args.hicache_size,
-                self.page_size,
-                server_args.hicache_mem_layout,
-                allocator_type=server_args.hicache_storage_backend,
-            )
-        elif isinstance(self.kv_cache, MLATokenToKVPool):
-            self.token_to_kv_pool_host = MLATokenToKVPoolHost(
-                self.kv_cache,
-                server_args.hicache_ratio,
-                server_args.hicache_size,
-                self.page_size,
-                server_args.hicache_mem_layout,
-                allocator_type=server_args.hicache_storage_backend,
-            )
-        else:
-            raise ValueError(f"HiRadixCache only supports MHA and MLA yet")
-
         self.tp_group = params.tp_cache_group
         self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
         self.pp_rank = params.pp_rank
@@ -1261,7 +1310,12 @@ class HiRadixCache(RadixCache):
 
             self._record_remove_event(x)
             if physical_count > 0:
-                num_evicted += self.cache_controller.evict_host(host_indices)
+                if self._uses_cp_hicache and hasattr(
+                    self.cache_controller, "evict_cp_host"
+                ):
+                    num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache)
+                else:
+                    num_evicted += self.cache_controller.evict_host(host_indices)
             x.host_len = 0
             x.cp_hicache = None
             x.host_value = None

python/sglang/srt/models/deepseek_nextn.py

@@ -184,6 +184,12 @@ class DeepseekModelNextN(nn.Module):
         local_input_ids = pad_cp_local_input_ids_for_embedding(
             forward_batch, local_input_ids
         )
+        self._debug_cp_draft_shared_kv(
+            "local_embedding_path "
+            f"full_tokens={full_num_tokens} "
+            f"local_tokens={local_num_tokens} "
+            f"padded_tokens={padded_token_count}"
+        )
 
         hidden_states = self.embed_tokens(local_input_ids)
         if hidden_states.shape[0] != local_num_tokens:
@@ -230,6 +236,11 @@ class DeepseekModelNextN(nn.Module):
                     if hidden_states is None:
                         # Conservative compatibility fallback: embed full input
                         # so all TP ranks all-reduce the same shape, then CP-split.
+                        self._debug_cp_draft_shared_kv(
+                            "full_embedding_fallback "
+                            f"full_tokens={input_ids.shape[0]} "
+                            f"local_tokens={local_num_tokens}"
+                        )
                         hidden_states = cp_split_and_rebuild_data(
                             forward_batch, self.embed_tokens(input_ids)
                         )