[DCP] [HiCache] Support role-local DCP storage in prefill/decode serving (upstream #41695)

docs/docs/advanced_features/hicache_best_practices.mdx

@@ -100,7 +100,18 @@ general DCP and model-specific attention-backend restrictions still apply.
 Use `--hicache-host-memory-mode cache`; DCP L3 rejects `buffer_only` and
 speculative decoding. DSPARK with DCP remains supported at L1/L2 without L3.
 
-Use aggregated serving; P/D is not supported with DCP L3.
+In P/D deployments, keep persisted cache namespaces role-local when the roles
+use different layouts. Use separate file directories or Mooncake
+`extra_backend_tag` values. Live P-to-D transfer remains responsible for moving
+the request's KV and current recurrent state; L3 does not convert the stored
+objects between the roles' topologies. Decode-side radix reuse requires
+`--disaggregation-decode-enable-radix-cache`.
+For decode with DCP L3, also set
+`--hicache-storage-prefetch-policy wait_complete`. The default `timeout` and
+`best_effort` policies can stop before fulfilling the prefix promised to
+prefill, so startup and runtime attachment reject them on decode. Prefill and
+aggregated servers may use all three policies. Decode KV offload
+(`--disaggregation-decode-enable-offload-kvcache`) is not supported with DCP L3.
 
 DCP keys include shard and layout identity. Recurrent file keys also preserve
 TP ownership independently of MLA replication, including when long filenames

python/sglang/srt/arg_groups/hicache_hook.py

@@ -154,8 +154,6 @@ def validate_hicache_dcp_storage(
     cfg = resolving_view(server_args)
     if cfg.dcp_size <= 1:
         return
-    if cfg.disaggregation_mode != "null":
-        raise NotImplementedError("HiCache L3 with DCP requires aggregated serving.")
     if not use_mla_backend(server_args):
         raise NotImplementedError("HiCache L3 with DCP requires MLA.")
     if cfg.hicache_host_memory_mode == "buffer_only":
@@ -165,6 +163,22 @@ def validate_hicache_dcp_storage(
         raise NotImplementedError(
             "HiCache L3 with DCP does not support speculative draft storage."
         )
+    if cfg.disaggregation_mode == "decode":
+        if cfg.disaggregation_decode_enable_offload_kvcache:
+            raise NotImplementedError(
+                "HiCache L3 with DCP does not support decode offload."
+            )
+        # Decode promises the probed L3 span to prefill before the read completes.
+        policy = (
+            cfg.hicache_storage_prefetch_policy
+            if prefetch_policy is None
+            else prefetch_policy
+        )
+        if policy != "wait_complete":
+            raise NotImplementedError(
+                "Decode HiCache L3 with DCP requires "
+                "--hicache-storage-prefetch-policy wait_complete."
+            )
     if (storage_backend or cfg.hicache_storage_backend) not in ("file", "mooncake"):
         raise NotImplementedError(
             "HiCache L3 with DCP requires file or Mooncake storage."

python/sglang/srt/disaggregation/decode.py

@@ -1349,6 +1349,18 @@ def pop_preallocated(
                 retractable_tokens=retractable_tokens, count_retracted=True
             )
         reserved_restore_tokens = self._hicache_pending_restore_tokens()
+        reserved_restore_states = 0
+        if (
+            getattr(self.req_to_token_pool, "mamba_allocator", None) is not None
+            and self.scheduler.enable_decode_hicache
+        ):
+            reserved_restore_states = sum(
+                dr.prefix_match.needs_local_restore
+                for dr in self.transfer_queue.queue
+                if dr.prefix_match is not None
+                and dr.hicache_restore_status == HiCacheRestoreResult.PENDING
+                and dr.hicache_restore_lock is None
+            )
         full_allocatable_tokens -= reserved_restore_tokens
         # Sort by priority before any index-based bookkeeping so that both the
         # abort-scan loop and the preallocation loop operate on the same order.
@@ -1451,10 +1463,12 @@ def pop_preallocated(
             if self.req_to_token_pool.available_size() <= 0:
                 break
 
-            # Reserve enough Mamba capacity for radix COW and tracking buffers
-            # before prefix matching can bind state or lock cached entries.
+            # Pending HiCache restores also need canonical recurrent slots.
             if not ensure_mamba_capacity(
-                self.req_to_token_pool, [decode_req.req], self.tree_cache
+                self.req_to_token_pool,
+                [decode_req.req],
+                self.tree_cache,
+                reserved_slots=reserved_restore_states,
             ):
                 break
 
@@ -1472,6 +1486,17 @@ def pop_preallocated(
             if use_decode_radix_cache:
                 # Match prefix against decode's radix cache.
                 prefix_match = self._match_prefix_and_lock(decode_req.req)
+                # A hybrid restore materializes one canonical checkpoint in
+                # addition to the request's live state received from prefill.
+                # If it cannot fit, promise only the device-resident prefix.
+                if prefix_match.needs_local_restore and not ensure_mamba_capacity(
+                    self.req_to_token_pool,
+                    [decode_req.req],
+                    self.tree_cache,
+                    reserved_slots=reserved_restore_states + 1,
+                ):
+                    prefix_match.l2_host_hit_length = 0
+                    prefix_match.l3_storage_hit_length = 0
                 prefix_indices = prefix_match.prefix_indices
                 # prefix_len: tokens already on device (L1 hit).
                 # total_prefix_len: full prefix promised to prefill
@@ -1573,20 +1598,28 @@ def pop_preallocated(
                     failed_reqs.append(decode_req)
                     indices_to_remove.add(i)
                     continue
+            if self.scheduler.enable_decode_hicache:
+                self._start_hicache_prefetch(decode_req.req, prefix_match)
+                if prefix_match is not None:
+                    # A probe is only a candidate hit. Admission may decline
+                    # prefetch, so allocate/advertise the resulting prefix.
+                    # Preserve any smaller SWA transfer cap selected above.
+                    total_prefix_len = min(
+                        total_prefix_len, prefix_match.decode_prefix_len
+                    )
             dst_kv_indices = self._pre_alloc(
                 decode_req.req,
                 prefix_indices,
                 prefix_len,
                 total_prefix_len,
             )
             decode_req.prefix_match = prefix_match
-            if self.scheduler.enable_decode_hicache:
-                self._start_hicache_prefetch(decode_req.req, prefix_match)
             hisparse_req_budget -= 1
             # Recompute from actual pool state for the next queue entry.
             # This accounts for page rounding and newly locked evictable cache.
             if prefix_match is not None:
                 reserved_restore_tokens += prefix_match.restore_token_count
+                reserved_restore_states += int(prefix_match.needs_local_restore)
             full_allocatable_tokens = self._allocatable_token_budgets(
                 retractable_tokens=retractable_tokens,
                 count_retracted=True,

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -143,6 +143,8 @@ def _start_hicache_prefetch(
             prefix_match.prefetch_registered = self.tree_cache.has_ongoing_prefetch(
                 req.cache_request_handle
             )
+            if not prefix_match.prefetch_registered:
+                prefix_match.l3_storage_hit_length = 0
         except Exception as e:
             logger.warning(
                 "HiCache L3 prefetch failed for rid=%s: %s; falling back to L2-only LoadingBack",
@@ -223,6 +225,7 @@ def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool:
                 best_match_node=rematch.best_match_node,
                 host_hit_length=rematch.host_hit_length,
                 req=dr.req,
+                restore_request_state=False,
             )
         )
         # The rematch repointed req.last_node to feed init_load_back's device

python/sglang/srt/managers/cache_controller.py

@@ -589,6 +589,7 @@ def attach_storage_backend(
         prefetch_threshold: int = 256,
         model_name: Optional[str] = None,
         storage_backend_extra_config: Optional[dict] = None,
+        prefetch_policy: Optional[str] = None,
     ):
         """Attach (enable) storage backend at runtime.
 
@@ -605,6 +606,7 @@ def attach_storage_backend(
             validate_hicache_dcp_storage(
                 get_server_args(),
                 storage_backend=storage_backend,
+                prefetch_policy=prefetch_policy,
             )
             self._validate_dcp_storage_pools(
                 self.storage_host_pool,

python/sglang/srt/mem_cache/allocation.py

@@ -292,13 +292,15 @@ def ensure_mamba_capacity(
     req_to_token_pool: ReqToTokenPool,
     reqs: list[Req],
     tree_cache: BasePrefixCache | None,
+    *,
+    reserved_slots: int = 0,
 ) -> bool:
-    """Reclaim the Mamba slots needed before binding any request rows."""
+    """Reclaim request slots while preserving pending checkpoint restore capacity."""
     if not isinstance(req_to_token_pool, HybridReqToTokenPool):
         return True
 
     # Count the slots HybridReqToTokenPool.alloc will allocate for each request.
-    needed = 0
+    needed = reserved_slots
     for req in reqs:
         if not req.kv.holds_mamba:
             needed += 1

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -264,6 +264,9 @@ class InitLoadBackParams:
     host_hit_length: int
     mem_quota: Optional[int] = None
     req: Optional[Req] = None
+    # Decode in P/D receives its live state from prefill. Restore canonical
+    # cache state without copying an older checkpoint into that request slot.
+    restore_request_state: bool = True
 
 
 class MatchResult(NamedTuple):

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -303,6 +303,7 @@ def attach_storage_backend(
         model_name: Optional[str] = None,
         storage_backend_extra_config: Optional[dict] = None,
         host_pools: Optional[list[PoolEntry]] = None,
+        prefetch_policy: Optional[str] = None,
     ):
         if isinstance(
             self.storage_host_pool, UnifiedPageEnvelopeHostPool
@@ -334,6 +335,7 @@ def attach_storage_backend(
                 prefetch_threshold=prefetch_threshold,
                 model_name=model_name,
                 storage_backend_extra_config=storage_backend_extra_config,
+                prefetch_policy=prefetch_policy,
             )
         except Exception:
             if self.pp_prefetch_command_group is not None:

python/sglang/srt/mem_cache/unified_cache/storage_attachment.py

@@ -167,6 +167,7 @@ def attach(
                 model_name=served_model_name,
                 storage_backend_extra_config=extra_config,
                 host_pools=controller.mem_pool_host.entries,
+                prefetch_policy=cache.prefetch_stop_policy,
             )
         except Exception as e:
             (

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -3487,7 +3487,10 @@ def init_load_back(
                 and (req.swa_host_hit_length > 0 or req.mamba_host_hit_length > 0)
             )
         ):
-            if self.load_back(best_match_node_id, mem_quota, req=req):
+            # Keep req for prefix/lock bookkeeping, but P/D must not queue an
+            # older checkpoint into the live state slot owned by P's transfer.
+            state_req = req if params.restore_request_state else None
+            if self.load_back(best_match_node_id, mem_quota, req=state_req):
                 new_indices = self.tree_core.collect_full_device_indices(
                     best_match_node_id, last_best_match_device_node_id
                 )