fix(pd): make DSV4 retraction safe

python/sglang/srt/disaggregation/decode.py

@@ -94,11 +94,7 @@
     set_schedule_time_batch,
     set_time_batch,
 )
-from sglang.srt.runtime_context import (
-    get_disagg,
-    get_memory,
-    get_parallel,
-)
+from sglang.srt.runtime_context import get_disagg, get_memory, get_parallel
 from sglang.srt.utils import ceil_align, get_num_new_pages, is_npu
 from sglang.srt.utils.network import NetworkAddress
 from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
@@ -821,24 +817,34 @@ def resume_retracted_reqs(
             if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens:
                 break
 
-            resumed_reqs.append(req)
-            indices_to_remove.add(i)
             req.is_retracted = False
             self._pre_alloc(req)
-            full_allocatable_tokens -= full_required
-            if uses_swa_tail_prealloc:
-                swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
-                    count_retracted=False,
-                    extra_reserved_reqs=len(resumed_reqs),
-                )
-
-            retraction_restore(
+            restored = retraction_restore(
                 req,
                 self.tree_cache,
                 self.req_to_token_pool,
                 self.token_to_kv_pool_allocator,
                 get_disagg().disaggregation_decode_retraction_backup,
             )
+            indices_to_remove.add(i)
+            if not restored:
+                error_message = "Retraction KV restore failed. Aborting the request."
+                prepare_abort(
+                    req,
+                    error_message,
+                    status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                )
+                self.scheduler.output_streamer.stream_output([req], req.return_logprob)
+                release_kv_cache(req, self.tree_cache, is_insert=False)
+                continue
+
+            resumed_reqs.append(req)
+            full_allocatable_tokens -= full_required
+            if uses_swa_tail_prealloc:
+                swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
+                    count_retracted=False,
+                    extra_reserved_reqs=len(resumed_reqs),
+                )
 
         self.retracted_queue = [
             entry

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py

@@ -236,6 +236,10 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
     accessors instead).
     """
 
+    def supports_host_pool_retraction(self, is_speculative: bool = False) -> bool:
+        # Generic V4 HostPool sidecars do not implement kernel_ascend transfers.
+        return False
+
     def __init__(self, *args, **kwargs):
         c128_page_size = get_schedule().c128_page_size
         if c128_page_size <= 0 or c128_page_size % 16 != 0:

python/sglang/srt/managers/schedule_batch.py

@@ -87,10 +87,7 @@
 from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import (
     NewTokenRatioTracker,
 )
-from sglang.srt.mem_cache.allocation import (
-    alloc_for_decode,
-    alloc_for_extend,
-)
+from sglang.srt.mem_cache.allocation import alloc_for_decode, alloc_for_extend
 from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import (
@@ -2851,16 +2848,15 @@ def retract_decode(
             if self.release_req(idx, len(sorted_indices), server_args):
                 retracted_reqs.append(req)
             else:
-                # The retraction host pool could not hold the backup and the
-                # device KV is already freed, so the request cannot resume.
+                # The device KV is already freed and the backup was not
+                # resumable, so the request cannot continue.
                 req.to_finish = FINISH_ABORT(
-                    "Retraction host KV pool exhausted. Aborting the request.",
+                    "Retraction KV backup failed. Aborting the request.",
                     status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                 )
                 reqs_to_abort.append(req)
                 logger.warning(
-                    "retract_decode: aborted request %s, retraction host pool "
-                    "exhausted",
+                    "retract_decode: aborted request %s, KV backup failed",
                     req.rid,
                 )
 

python/sglang/srt/mem_cache/base_swa_memory_pool.py

@@ -16,6 +16,10 @@ class BaseSWAKVPool(KVCache):
 
     swa_kv_pool: KVCache
 
+    def supports_host_pool_retraction(self, is_speculative: bool = False) -> bool:
+        """Whether the pool's complete live state is represented in HiCache."""
+        return False
+
     @abc.abstractmethod
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor) -> None:
         raise NotImplementedError()

python/sglang/srt/mem_cache/common.py

@@ -35,6 +35,7 @@ class RetractionBackup(NamedTuple):
     cpu_tensors: Any = None
     host_indices: Optional[torch.Tensor] = None
     pool_transfers: Optional[list[PoolTransfer]] = None
+    swa_window_start: Optional[int] = None
 
 
 def kv_to_page_indices(kv_indices: torch.Tensor, page_size: int) -> np.ndarray:
@@ -148,7 +149,16 @@ def retraction_backup(
     """Returns False when the host pool cannot hold the backup; the caller
     aborts the request since its KV cannot be preserved."""
     if backend == "cpu_tensor":
-        req.offload_kv_cache(req_to_token_pool, token_to_kv_pool_allocator)
+        try:
+            req.offload_kv_cache(req_to_token_pool, token_to_kv_pool_allocator)
+        except NotImplementedError:
+            req.retraction_backup = None
+            logger.error(
+                "CPU-tensor retraction backup is unsupported for request %s; "
+                "aborting the request instead of crashing the scheduler.",
+                req.rid,
+            )
+            return False
         return True
     if backend != "host_pool":
         raise ValueError(f"Unknown retraction backup backend: {backend}")
@@ -166,19 +176,20 @@ def retraction_restore(
     req_to_token_pool: ReqToTokenPool,
     token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
     backend: str,
-) -> None:
+) -> bool:
     if backend == "cpu_tensor":
         req.load_kv_cache(req_to_token_pool, token_to_kv_pool_allocator)
-        return
+        return True
     if backend != "host_pool":
         raise ValueError(f"Unknown retraction backup backend: {backend}")
     if req.seqlen <= 1:
-        return
+        return True
 
     unified_cache = cast("UnifiedRadixCache", tree_cache)
     assert req.retraction_backup is not None
-    unified_cache.retraction_restore(req, req.retraction_backup)
+    restored = unified_cache.retraction_restore(req, req.retraction_backup)
     req.retraction_backup = None
+    return restored
 
 
 def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> None:

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -664,6 +664,16 @@ def get_unified_kv(self, layer_id: int) -> torch.Tensor:
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
         self.full_to_swa_index_mapping = full_to_swa_index_mapping
 
+    def supports_host_pool_retraction(self, is_speculative: bool = False) -> bool:
+        # Unified KV has no separate SWA host entry. HiSparse owns most C4
+        # state outside the HiCache group. Speculative request-indexed state has
+        # not yet been admitted to the retraction transfer contract.
+        return (
+            not is_speculative
+            and self.swa_kv_pool is not None
+            and not isinstance(self.c4_kv_pool, HiSparseC4DevicePool)
+        )
+
     def get_ring_size(self, compress_ratio: int) -> int:
         is_speculative = get_spec().speculative_algorithm is not None
         return get_compress_state_ring_size(compress_ratio, is_speculative)

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -34,10 +34,10 @@ class KVCacheBuildResult:
 from sglang.srt.environ import envs
 from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.managers.mm_schedule import init_mm_embedding_cache
+from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
-from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool
 from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache
-from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 from sglang.srt.model_loader.utils import get_resolved_model_impl
 from sglang.srt.runtime_context import (
@@ -46,6 +46,7 @@ class KVCacheBuildResult:
     get_memory,
     get_parallel,
     get_schedule,
+    get_spec,
 )
 
 if TYPE_CHECKING:
@@ -103,10 +104,7 @@ def _register_legacy_hicache_draft(
     server_args: ServerArgs,
     page_size: int,
 ) -> None:
-    from sglang.srt.mem_cache.memory_pool import (
-        MHATokenToKVPool,
-        MLATokenToKVPool,
-    )
+    from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
     from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
     from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
 
@@ -146,6 +144,21 @@ def _register_legacy_hicache_draft(
 BACKUP_ONLY_HICACHE_RATIO = 0.2
 
 
+def _supports_host_pool_retraction(
+    kv_cache: KVCache,
+    full_tokens_per_layer: Optional[int],
+    is_speculative: bool = False,
+) -> bool:
+    if isinstance(kv_cache, MHATokenToKVPool):
+        return True
+    return (
+        isinstance(kv_cache, BaseSWAKVPool)
+        and full_tokens_per_layer is not None
+        and full_tokens_per_layer > 0
+        and kv_cache.supports_host_pool_retraction(is_speculative)
+    )
+
+
 def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
     """Resolve the retraction backend onto the config bags and return it.
 
@@ -165,8 +178,10 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
             if tp_worker.is_hybrid_swa
             else None
         )
-        supports_host_pool = isinstance(kv_cache, MHATokenToKVPool) or (
-            isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0
+        supports_host_pool = _supports_host_pool_retraction(
+            kv_cache,
+            full_tokens_per_layer,
+            get_spec().speculative_algorithm is not None,
         )
         schedule = get_schedule()
         priority_preemption = (

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -123,6 +123,9 @@ def __init__(
             f"SWAKVPool {'VA upper bound' if self.post_capture_active else 'mem usage'}: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}"
         )
 
+    def supports_host_pool_retraction(self, is_speculative: bool = False) -> bool:
+        return self.swa_kv_pool is not None
+
     @property
     def post_capture_active(self) -> bool:
         """True iff the sub-pools took the post-capture VA-backed path (both share the flag)."""