diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index b73eea2e9..9ddd167d6 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 94340105e..967e1120e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 341c874b3..cbc6a93f1 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index b9b7a7074..efcd42366 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index bedf2f740..57166f23a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 6b7ae4fcd..268976191 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 24adcb48b..6e80925f0 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: