[Fix] Unify pinned host pool release on graceful shutdown (#32029)

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -105,6 +105,9 @@ class DecodeKVCacheOffloadManager:
         self.offload_inflight = {}
         logger.info("Enable offload kv cache for decode side")
 
+    def release_host_resources(self) -> None:
+        self.decode_host_mem_pool.destroy()
+
     def _mark_offload_started(self, rid):
         self.offload_inflight[rid] = self.offload_inflight.get(rid, 0) + 1
 

python/sglang/srt/managers/scheduler.py

@@ -1465,11 +1465,9 @@ class Scheduler(
         # HostKVCache.destroy. Called from run_scheduler_process's finally.
         if self.hisparse_coordinator is not None:
             self.hisparse_coordinator.destroy()
-        # A plain HiRadixCache (no hisparse) also holds a large pinned host KV
-        # pool; unregister it here too, else the kernel unpins it during reclaim.
-        host_pool = getattr(self.tree_cache, "token_to_kv_pool_host", None)
-        if host_pool is not None:
-            host_pool.destroy()
+        self.tree_cache.release_host_resources()
+        if self.decode_offload_manager is not None:
+            self.decode_offload_manager.release_host_resources()
 
     def run_event_loop(self) -> None:
         """Run the scheduler's event loop.

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -236,6 +236,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
             )
             self.metrics_collector.increment_eviction_num_tokens(num_evicted)
 
+    def release_host_resources(self) -> None:
+        """Release pinned host buffers in userspace on graceful shutdown.
+
+        Kernel-side unpinning during process reclaim can stall teardown for
+        tens of seconds (see HostKVCache.destroy). Idempotent.
+        """
+
     @abstractmethod
     def reset(self):
         pass

python/sglang/srt/mem_cache/hi_mamba_radix_cache.py

@@ -207,6 +207,9 @@ class HiMambaRadixCache(MambaRadixCache):
         )
         super().reset()
 
+    def release_host_resources(self) -> None:
+        self.host_pool_group.destroy()
+
     def write_backup(self, node: TreeNode, write_back=False) -> int:
         # Backup invariant (for write-through mode): backed-up nodes must form a
         # contiguous prefix from root — no gaps.  Skip if parent isn't backed

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -775,6 +775,10 @@ class HiRadixCache(RadixCache):
         self.evictable_host_leaves.clear()
         super().reset()
 
+    def release_host_resources(self) -> None:
+        if self.token_to_kv_pool_host is not None:
+            self.token_to_kv_pool_host.destroy()
+
     def get_height(self, node: TreeNode):
         height = 0
         while node != self.root_node:

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -1573,6 +1573,10 @@ class HostPoolGroup:
         for entry in self.entries:
             entry.host_pool.clear()
 
+    def destroy(self) -> None:
+        for entry in self.entries:
+            entry.host_pool.destroy()
+
     def available_size(self):
         return self.anchor_entry.host_pool.available_size()
 

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -371,6 +371,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
 
         # HiCache D↔H defaults (overridden by init_hicache)
         self.cache_controller: Optional[HybridCacheController] = None
+        self.host_pool_group = None  # set by attach_hybrid_pool_to_unified_cache
         self.write_through_threshold = 256
         self.prefetch_stop_policy = "best_effort"
         self.prefetch_threshold = 256
@@ -568,6 +569,10 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
     def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None:
         self.sidecar_pool_specs.append(spec)
 
+    def release_host_resources(self) -> None:
+        if self.host_pool_group is not None:
+            self.host_pool_group.destroy()
+
     def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
         result = self.session.try_match_prefix(params)
         if result is not None: