From 38f32c38abd41a296bd21186b5d8c5e6af1f8997 Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Tue, 26 May 2026 22:38:00 +0800 Subject: [PATCH] [UnifiedRadixTree]: Support L3 HiStorage framework (#26062) --- .../sglang/srt/mem_cache/base_prefix_cache.py | 4 + .../hybrid_cache/hybrid_cache_controller.py | 131 ++++ .../hybrid_cache/hybrid_pool_assembler.py | 68 +- .../full_component.py | 31 +- .../mamba_component.py | 113 ++- .../unified_cache_components/swa_component.py | 11 +- .../tree_component.py | 26 +- .../srt/mem_cache/unified_radix_cache.py | 690 +++++++++++++++++- .../test_unified_radix_cache_kl_hicache.py | 2 +- ..._unified_radix_cache_kl_hicache_nightly.py | 59 +- ...st_unified_radix_cache_kl_hicache_part2.py | 77 ++ .../test_unified_radix_cache_unittest.py | 7 - 12 files changed, 1148 insertions(+), 71 deletions(-) create mode 100644 test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 6eaefc714..e0ba5b4a9 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -75,6 +75,7 @@ class InsertResult: prefix_len: int mamba_exist: bool = False + inserted_host_node: Any = None @dataclasses.dataclass @@ -101,6 +102,7 @@ class IncLockRefResult: delta: Optional[int] = None swa_uuid_for_lock: Optional[int] = None + swa_uuid_for_host_lock: Optional[int] = None # Component nodes that were tombstones at acquire time. Replaying this set # at release prevents a short-lived lock from consuming a later load-back or # request lock after that tombstone becomes a valid device value. @@ -112,6 +114,7 @@ class IncLockRefResult: """Convert to the corresponding DecLockRefParams for dec_lock_ref.""" return DecLockRefParams( swa_uuid_for_lock=self.swa_uuid_for_lock, + swa_uuid_for_host_lock=self.swa_uuid_for_host_lock, skip_lock_node_ids={ component_type: set(node_ids) for component_type, node_ids in self.skip_lock_node_ids.items() @@ -124,6 +127,7 @@ class DecLockRefParams: """Parameters for dec_lock_ref operation.""" swa_uuid_for_lock: Optional[int] = None + swa_uuid_for_host_lock: Optional[int] = None skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field( default_factory=dict ) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index a339eef5d..1b09a5b96 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -1,8 +1,11 @@ from __future__ import annotations +import json import logging +import os import threading import time +from queue import Queue from typing import TYPE_CHECKING, Any, Callable, List, Optional import torch @@ -171,6 +174,7 @@ class HybridCacheController(BaseHiCacheController): enable_storage_metrics: bool = False, ): startup_storage_backend = storage_backend + self.extra_host_mem_release_queues: dict[PoolName, Queue] = {} super().__init__( token_to_kv_pool_allocator=token_to_kv_pool_allocator, mem_pool_host=mem_pool_host, @@ -204,6 +208,10 @@ class HybridCacheController(BaseHiCacheController): host_pools=getattr(mem_pool_host, "entries", None), ) + def _start_storage_threads(self): + super()._start_storage_threads() + self._init_extra_host_mem_release_queues() + def attach_storage_backend( self, storage_backend: str, @@ -222,10 +230,133 @@ class HybridCacheController(BaseHiCacheController): for entry in host_pools or []: self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name) + @staticmethod + def parse_storage_backend_extra_config( + storage_backend_extra_config: Optional[str], + ) -> tuple[dict, int, float, float, bool]: + extra_config = {} + if storage_backend_extra_config: + if storage_backend_extra_config.startswith("@"): + path = storage_backend_extra_config[1:] + ext = os.path.splitext(path)[1].lower() + with open(path, "rb" if ext == ".toml" else "r") as f: + if ext == ".json": + extra_config = json.load(f) + elif ext == ".toml": + import tomllib + + extra_config = tomllib.load(f) + elif ext in (".yaml", ".yml"): + import yaml + + extra_config = yaml.safe_load(f) + else: + raise ValueError( + f"Unsupported config file {path} (config format: {ext})" + ) + else: + extra_config = json.loads(storage_backend_extra_config) + + prefetch_threshold = extra_config.pop("prefetch_threshold", 256) + prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1) + prefetch_timeout_per_ki_token = extra_config.pop( + "prefetch_timeout_per_ki_token", 0.25 + ) + hicache_storage_pass_prefix_keys = extra_config.pop( + "hicache_storage_pass_prefix_keys", False + ) + + if not isinstance(prefetch_threshold, int): + raise ValueError( + f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}" + ) + if not isinstance(prefetch_timeout_base, (int, float)): + raise ValueError( + f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}" + ) + if not isinstance(prefetch_timeout_per_ki_token, (int, float)): + raise ValueError( + "prefetch_timeout_per_ki_token must be number, got " + f"{type(prefetch_timeout_per_ki_token).__name__}" + ) + if not isinstance(hicache_storage_pass_prefix_keys, bool): + raise ValueError( + "hicache_storage_pass_prefix_keys must be bool, got " + f"{type(hicache_storage_pass_prefix_keys).__name__}" + ) + + return ( + extra_config, + prefetch_threshold, + float(prefetch_timeout_base), + float(prefetch_timeout_per_ki_token), + hicache_storage_pass_prefix_keys, + ) + + def clear_storage_backend(self) -> bool: + if not self.enable_storage: + logger.warning("Hierarchical cache storage backend is not enabled.") + return False + if not hasattr(self.storage_backend, "clear"): + logger.warning( + "Storage backend %s does not support clear operation.", + type(self.storage_backend).__name__, + ) + return False + self.storage_backend.clear() + return True + + def _init_extra_host_mem_release_queues(self) -> None: + self.extra_host_mem_release_queues = {} + entries = getattr(self.mem_pool_host, "entries", None) or [] + anchor_entry = getattr(self.mem_pool_host, "anchor_entry", None) + for entry in entries: + if entry is anchor_entry or entry.is_primary_index_anchor: + continue + self.extra_host_mem_release_queues[entry.name] = Queue() + + def _append_host_mem_release_pages( + self, release_queue: Queue, host_indices: torch.Tensor, page_size: int + ) -> None: + if host_indices.numel() == 0: + return + for page in host_indices.split(page_size): + release_queue.put(page) + + def append_host_mem_release( + self, + host_indices: Optional[torch.Tensor] = None, + extra_pools: Optional[list[PoolTransfer]] = None, + ): + if host_indices is not None: + self._append_host_mem_release_pages( + self.host_mem_release_queue, + host_indices, + self.mem_pool_host.page_size, + ) + for transfer in extra_pools or []: + if transfer.host_indices is None or transfer.host_indices.numel() == 0: + continue + entry = self.mem_pool_host.entry_map.get(transfer.name) + if ( + entry is None + or entry.is_primary_index_anchor + or transfer.indices_from_pool is not None + ): + continue + release_queue = self.extra_host_mem_release_queues.get(transfer.name) + if release_queue is None: + continue + self._append_host_mem_release_pages( + release_queue, transfer.host_indices, entry.host_pool.page_size + ) + def reset(self): super().reset() if self.enable_storage: self.host_mem_release_queue.queue.clear() + for release_queue in self.extra_host_mem_release_queues.values(): + release_queue.queue.clear() self.prefetch_tokens_occupied = 0 def write( diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 0090932a9..e12c9d350 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -679,6 +679,11 @@ class StackStrategy: load_cache_event, attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, + storage_backend: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + enable_storage_metrics: bool = False, ) -> StackBuildResult: raise NotImplementedError @@ -704,6 +709,11 @@ class _DeepSeekV4Strategy(StackStrategy): load_cache_event, attn_cp_group=None, attn_tp_group=None, + storage_backend=None, + storage_backend_extra_config=None, + prefetch_threshold=256, + model_name=None, + enable_storage_metrics=False, ): from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -716,11 +726,15 @@ class _DeepSeekV4Strategy(StackStrategy): load_cache_event=load_cache_event, attn_cp_group=attn_cp_group, attn_tp_group=attn_tp_group, - storage_backend=None, + storage_backend=storage_backend, host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA), device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)), + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + enable_storage_metrics=enable_storage_metrics, ) sidecars = [ SidecarPoolSpec(pool_name=name, indices_from_pool=src) @@ -766,6 +780,11 @@ class _MambaStrategy(StackStrategy): load_cache_event, attn_cp_group=None, attn_tp_group=None, + storage_backend=None, + storage_backend_extra_config=None, + prefetch_threshold=256, + model_name=None, + enable_storage_metrics=False, ): from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -783,12 +802,16 @@ class _MambaStrategy(StackStrategy): load_cache_event=load_cache_event, attn_cp_group=attn_cp_group, attn_tp_group=attn_tp_group, - storage_backend=None, + storage_backend=storage_backend, use_mla=kvcache.use_mla, host_mamba_evict_fn=lambda n: cache.evict_host(n, ComponentType.MAMBA), device_mamba_evict_fn=lambda n: cache.evict(EvictParams(mamba_num=n)), + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + enable_storage_metrics=enable_storage_metrics, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -834,6 +857,11 @@ class _SwaStrategy(StackStrategy): load_cache_event, attn_cp_group=None, attn_tp_group=None, + storage_backend=None, + storage_backend_extra_config=None, + prefetch_threshold=256, + model_name=None, + enable_storage_metrics=False, ): from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -850,12 +878,16 @@ class _SwaStrategy(StackStrategy): load_cache_event=load_cache_event, attn_cp_group=attn_cp_group, attn_tp_group=attn_tp_group, - storage_backend=None, + storage_backend=storage_backend, use_mla=False, host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA), device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)), + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + enable_storage_metrics=enable_storage_metrics, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -887,6 +919,11 @@ class _DsaStrategy(StackStrategy): load_cache_event, attn_cp_group=None, attn_tp_group=None, + storage_backend=None, + storage_backend_extra_config=None, + prefetch_threshold=256, + model_name=None, + enable_storage_metrics=False, ): from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool @@ -904,7 +941,7 @@ class _DsaStrategy(StackStrategy): load_cache_event=load_cache_event, attn_cp_group=attn_cp_group, attn_tp_group=attn_tp_group, - storage_backend=None, + storage_backend=storage_backend, use_mla=use_mla, override_kv_cache_dim=full_kv_pool.kv_cache_dim, sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost( @@ -913,8 +950,12 @@ class _DsaStrategy(StackStrategy): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ), + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + enable_storage_metrics=enable_storage_metrics, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -961,6 +1002,11 @@ class _PlainKvStrategy(StackStrategy): load_cache_event, attn_cp_group=None, attn_tp_group=None, + storage_backend=None, + storage_backend_extra_config=None, + prefetch_threshold=256, + model_name=None, + enable_storage_metrics=False, ): from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool @@ -977,10 +1023,14 @@ class _PlainKvStrategy(StackStrategy): load_cache_event=load_cache_event, attn_cp_group=attn_cp_group, attn_tp_group=attn_tp_group, - storage_backend=None, + storage_backend=storage_backend, use_mla=use_mla, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + enable_storage_metrics=enable_storage_metrics, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -1057,6 +1107,9 @@ def attach_hybrid_pool_to_unified_cache( load_cache_event, attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, + storage_backend: Optional[str] = None, + storage_extra_config: Optional[dict] = None, + storage_prefetch_threshold: int = 256, ) -> None: """Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache.""" try: @@ -1071,6 +1124,11 @@ def attach_hybrid_pool_to_unified_cache( load_cache_event=load_cache_event, attn_cp_group=attn_cp_group, attn_tp_group=attn_tp_group, + storage_backend=storage_backend, + storage_backend_extra_config=storage_extra_config, + prefetch_threshold=storage_prefetch_threshold, + model_name=server_args.served_model_name, + enable_storage_metrics=cache._enable_metrics_flag, ) _apply_stack_result(cache, kvcache, params, result) except Exception: diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index bccb866a7..5ed89e3b4 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -65,7 +65,7 @@ class FullComponent(TreeComponent): # last_device_node, summing host_value lengths of evicted nodes. ct = self.component_type kv_host_hit = 0 - node = result.last_host_node + node = result.best_match_node root_node = self.cache.root_node while node is not result.last_device_node and node is not root_node: full_host = node.component_data[ct].host_value @@ -155,9 +155,22 @@ class FullComponent(TreeComponent): heapq.heappush(heap, (x.parent.last_access_time, x.parent)) def acquire_component_lock( - self, node: UnifiedTreeNode, result: IncLockRefResult + self, + node: UnifiedTreeNode, + result: IncLockRefResult, + lock_host: bool = False, ) -> IncLockRefResult: ct = self.component_type + + # Only the last host node needs to be protected. + if lock_host: + cd = node.component_data[ct] + if cd.host_value is None: + return result + cd.host_lock_ref += 1 + self.cache._update_evictable_leaf_sets(node) + return result + root = self.cache.root_node cur = node @@ -185,9 +198,20 @@ class FullComponent(TreeComponent): return result def release_component_lock( - self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + self, + node: UnifiedTreeNode, + params: Optional[DecLockRefParams], + lock_host: bool = False, ) -> None: ct = self.component_type + if lock_host: + cd = node.component_data[ct] + if cd.host_value is None or cd.host_lock_ref == 0: + return + cd.host_lock_ref -= 1 + self.cache._update_evictable_leaf_sets(node) + return + root = self.cache.root_node skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () cur = node @@ -255,6 +279,7 @@ class FullComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), + **kw, ) -> None: ct = self.component_type diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index 7ecfec0a2..fac322005 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -13,7 +13,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, @@ -213,34 +213,59 @@ class MambaComponent(TreeComponent): x = x_next def acquire_component_lock( - self, node: UnifiedTreeNode, result: IncLockRefResult + self, + node: UnifiedTreeNode, + result: IncLockRefResult, + lock_host: bool = False, ) -> IncLockRefResult: ct = self.component_type + if node is self.cache.root_node: + return result cd = node.component_data[ct] - value = cd.value + value = cd.host_value if lock_host else cd.value # A node in skip_lock_node_ids was a tombstone when this lock was acquired. if value is None: result.skip_lock_node_ids.setdefault(ct, set()).add(node.id) return result - if cd.lock_ref == 0: - vlen = len(value) - self.cache.component_evictable_size_[ct] -= vlen - self.cache.component_protected_size_[ct] += vlen - cd.lock_ref += 1 + if lock_host: + if cd.host_lock_ref == 0: + host_lru = self.cache.host_lru_lists[ct] + if host_lru.in_list(node): + host_lru.remove_node(node) + cd.host_lock_ref += 1 + else: + if cd.lock_ref == 0: + vlen = len(value) + self.cache.component_evictable_size_[ct] -= vlen + self.cache.component_protected_size_[ct] += vlen + cd.lock_ref += 1 return result def release_component_lock( - self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + self, + node: UnifiedTreeNode, + params: Optional[DecLockRefParams], + lock_host: bool = False, ) -> None: ct = self.component_type + if node is self.cache.root_node: + return cd = node.component_data[ct] skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () if node.id in skip_lock_node_ids: return - value = cd.value - if value is not None and cd.lock_ref > 0: + value = cd.host_value if lock_host else cd.value + if lock_host: + cd.host_lock_ref -= 1 + if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None: + host_lru = self.cache.host_lru_lists[ct] + if not host_lru.in_list(node): + host_lru.insert_mru(node) + return + + if cd.lock_ref > 0: if cd.lock_ref == 1: vlen = len(value) self.cache.component_evictable_size_[ct] += vlen @@ -392,6 +417,35 @@ class MambaComponent(TreeComponent): return transfers if transfers else None + if phase == CacheTransferPhase.BACKUP_STORAGE: + cd = node.component_data[ct] + if cd.host_value is None or not node.hash_value: + return None + return [ + PoolTransfer( + name=PoolName.MAMBA, + host_indices=cd.host_value, + keys=[node.hash_value[-1]], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + + if phase == CacheTransferPhase.PREFETCH: + host_indices = self._mamba_pool_host.alloc(1) + if host_indices is None: + self.cache.evict_host(1, ComponentType.MAMBA) + host_indices = self._mamba_pool_host.alloc(1) + if host_indices is None: + return [] + return [ + PoolTransfer( + name=PoolName.MAMBA, + host_indices=host_indices, + keys=["__placeholder__"], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + return None def commit_hicache_transfer( @@ -399,6 +453,7 @@ class MambaComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), + **kw, ) -> None: ct = self.component_type @@ -423,6 +478,41 @@ class MambaComponent(TreeComponent): self.cache.lru_lists[ct].insert_mru(node) self.cache.component_evictable_size_[ct] += count + elif phase == CacheTransferPhase.PREFETCH: + if not transfers: + return + transfer = transfers[0] + host_indices = transfer.host_indices + insert_result = kw.get("insert_result") + pool_storage_result = kw.get("pool_storage_result") + loaded = ( + pool_storage_result is not None + and pool_storage_result.extra_pool_hit_pages.get(PoolName.MAMBA, 0) >= 1 + ) + target_node = ( + insert_result.inserted_host_node if insert_result is not None else None + ) + if ( + host_indices is None + or target_node is None + or not loaded + or target_node.component_data[ct].host_value is not None + ): + self.cache.cache_controller.append_host_mem_release( + extra_pools=[transfer] + ) + if insert_result is not None: + insert_result.mamba_exist = True + return + + target_node.component_data[ct].host_value = host_indices.clone() + if target_node.component_data[ct].value is None: + host_lru = self.cache.host_lru_lists[ct] + if not host_lru.in_list(target_node): + host_lru.insert_mru(target_node) + if insert_result is not None: + insert_result.mamba_exist = False + def drive_host_eviction( self, num_tokens: int, tracker: dict[ComponentType, int] ) -> None: @@ -445,4 +535,5 @@ class MambaComponent(TreeComponent): x, self, target=EvictLayer.HOST, tracker=tracker ) self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST) + self.cache._update_evictable_leaf_sets(x) x = x_next diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 63223625e..0ba006079 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -350,7 +350,10 @@ class SWAComponent(TreeComponent): x = x_next def acquire_component_lock( - self, node: UnifiedTreeNode, result: IncLockRefResult + self, + node: UnifiedTreeNode, + result: IncLockRefResult, + lock_host: bool = False, ) -> IncLockRefResult: ct = self.component_type root = self.cache.root_node @@ -384,7 +387,10 @@ class SWAComponent(TreeComponent): return result def release_component_lock( - self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + self, + node: UnifiedTreeNode, + params: Optional[DecLockRefParams], + lock_host: bool = False, ) -> None: ct = self.component_type root = self.cache.root_node @@ -484,6 +490,7 @@ class SWAComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), + **kw, ) -> None: ct = self.component_type diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index ae6f71167..2b9b03b88 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -276,9 +276,12 @@ class TreeComponent(ABC): @abstractmethod def acquire_component_lock( - self, node: UnifiedTreeNode, result: IncLockRefResult + self, + node: UnifiedTreeNode, + result: IncLockRefResult, + lock_host: bool = False, ) -> IncLockRefResult: - """Increment lock_ref for this component, protecting nodes from + """Increment component lock refs, protecting nodes from eviction. Updates evictable → protected size on first lock. - Full: path-lock — walks from node up to root, incrementing lock_ref on every ancestor. @@ -286,21 +289,31 @@ class TreeComponent(ABC): sliding window is filled; records a component_uuid at the boundary for release_component_lock to know where to stop. - Mamba: single-node lock — only increments lock_ref on the - node itself (mamba state is per-leaf, not per-path).""" + node itself (mamba state is per-leaf, not per-path). + + When ``lock_host`` is True, the lock applies to host-side state: + - Full: single-node host lock. + - SWA: host window-lock with a dedicated host UUID boundary. + - Mamba: single-node host lock with host LRU detach.""" ... @abstractmethod def release_component_lock( - self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + self, + node: UnifiedTreeNode, + params: Optional[DecLockRefParams], + lock_host: bool = False, ) -> None: - """Decrement lock_ref for this component, un-protecting nodes. + """Decrement component lock refs, un-protecting nodes. Updates protected → evictable size when lock_ref drops to 0. - Full: path-unlock — walks from node up to root, decrementing lock_ref on every ancestor. - SWA: path-unlock — walks upward, stopping at the node whose component_uuid matches the one recorded during acquire. - Mamba: single-node unlock — only decrements lock_ref on the - node itself.""" + node itself. + + When ``lock_host`` is True, the inverse host-side semantics apply.""" ... def prepare_for_caching_req( @@ -351,6 +364,7 @@ class TreeComponent(ABC): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), + **kw, ) -> None: """Post-transfer bookkeeping: store host indices, update LRU, etc.""" pass diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index fa298b053..a42c5d0c3 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -5,7 +5,8 @@ import threading import time from array import array from collections import defaultdict -from functools import partial +from functools import lru_cache, partial +from queue import Empty from typing import TYPE_CHECKING, Any, Optional import torch @@ -28,6 +29,9 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransfer, SidecarPoolSpec, ) +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache_components import ( _NUM_COMPONENT_TYPES, @@ -42,6 +46,8 @@ from sglang.srt.mem_cache.unified_cache_components import ( TreeComponent, get_and_increase_time_counter, ) +from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value +from sglang.srt.observability.metrics_collector import StorageMetricsCollector from sglang.srt.session.streaming_session import StreamingSession if TYPE_CHECKING: @@ -93,6 +99,18 @@ class UnifiedTreeNode: def __lt__(self, other: UnifiedTreeNode): return self.last_access_time < other.last_access_time + def get_last_hash_value(self) -> Optional[str]: + if self.hash_value is None or len(self.hash_value) == 0: + return None + return self.hash_value[-1] + + @lru_cache(maxsize=1) + def get_prefix_hash_values(self, node: UnifiedTreeNode) -> list[str]: + if node is None or node.hash_value is None: + return [] + + return node.get_prefix_hash_values(node.parent) + node.hash_value + class UnifiedLRUList: def __init__( @@ -220,6 +238,10 @@ class UnifiedRadixCache(BasePrefixCache): if params.enable_metrics: self.init_metrics_collector() + self._enable_metrics_flag = params.enable_metrics + self.enable_storage_metrics = False + self.storage_metrics_collector: Optional[StorageMetricsCollector] = None + self.extra_metric_labels = None assert params.tree_components is not None self.tree_components = tuple(params.tree_components) @@ -248,6 +270,11 @@ class UnifiedRadixCache(BasePrefixCache): # HiCache D↔H defaults (overridden by init_hicache) self.cache_controller = None self.write_through_threshold = 256 + self.prefetch_stop_policy = "best_effort" + self.prefetch_threshold = 256 + self.prefetch_timeout_base = 1.0 + self.prefetch_timeout_per_page = 0.25 + self.hicache_storage_pass_prefix_keys = False self.reset() logger.info(f"Init Unified RadixTree with components {self.tree_components}") @@ -260,6 +287,7 @@ class UnifiedRadixCache(BasePrefixCache): self.root_node = UnifiedTreeNode(self.tree_components) self.root_node.key = RadixKey(array("q"), None) self.root_node.component_data[BASE_COMPONENT_TYPE].value = [] + self.root_node.hash_value = [] for ct in self.tree_components: self.root_node.component_data[ct].lock_ref = 1 self.component_evictable_size_ = {ct: 0 for ct in self.tree_components} @@ -281,12 +309,14 @@ class UnifiedRadixCache(BasePrefixCache): ] = {} self.ongoing_load_back: dict[int, tuple[UnifiedTreeNode, DecLockRefParams]] = {} self.enable_storage = False + self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {} self.ongoing_prefetch: dict = {} self.ongoing_backup: dict = {} if self.cache_controller is not None: self.cache_controller.reset() self.cache_controller.mem_pool_host.clear() + self.enable_storage = self.cache_controller.enable_storage self._empty_match_result = MatchResult( device_indices=torch.empty( @@ -316,6 +346,26 @@ class UnifiedRadixCache(BasePrefixCache): self.load_cache_event = threading.Event() self.sidecar_pool_specs.clear() + self.extra_metric_labels = server_args.extra_metric_labels + + # Parse storage config once, share with assembler and tree + storage_backend = server_args.hicache_storage_backend + storage_extra_config = None + storage_prefetch_threshold = 256 + prefetch_timeout_base = 1.0 + prefetch_timeout_per_ki_token = 0.25 + hicache_storage_pass_prefix_keys = False + if storage_backend is not None: + ( + storage_extra_config, + storage_prefetch_threshold, + prefetch_timeout_base, + prefetch_timeout_per_ki_token, + hicache_storage_pass_prefix_keys, + ) = HybridCacheController.parse_storage_backend_extra_config( + server_args.hicache_storage_backend_extra_config + ) + attach_hybrid_pool_to_unified_cache( self, params, @@ -323,6 +373,9 @@ class UnifiedRadixCache(BasePrefixCache): load_cache_event=self.load_cache_event, attn_cp_group=params.attn_cp_cache_group, attn_tp_group=params.attn_tp_cache_group, + storage_backend=storage_backend, + storage_extra_config=storage_extra_config, + storage_prefetch_threshold=storage_prefetch_threshold, ) # State initialization @@ -330,14 +383,19 @@ class UnifiedRadixCache(BasePrefixCache): 1 if server_args.hicache_write_policy == "write_through" else 2 ) self.load_back_threshold = 256 + self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy - logger.info( - f"HiCache D\u2194H initialized: " - f"host_pool_size={self.host_pool_group.size}, " - f"write_policy={server_args.hicache_write_policy}, " - f"tp_world_size={self.tp_world_size}, " - f"transfer_layer_num={self.cache_controller.layer_num}" - ) + if storage_backend is not None: + self._apply_storage_runtime_config( + storage_backend=storage_backend, + prefetch_threshold=storage_prefetch_threshold, + prefetch_timeout_base=prefetch_timeout_base, + prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token, + hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys, + enable_storage=self.cache_controller.enable_storage, + enable_storage_metrics=self._enable_metrics_flag, + extra_metric_labels=self.extra_metric_labels, + ) def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None: self.sidecar_pool_specs.append(spec) @@ -435,6 +493,29 @@ class UnifiedRadixCache(BasePrefixCache): # TODO: delta is not aggregated from components; no caller uses it yet. return DecLockRefResult() + def inc_host_lock_ref(self, node: Any) -> IncLockRefResult: + if self.disable: + return IncLockRefResult() + result = IncLockRefResult() + for component in self._components_tuple: + result = component.acquire_component_lock( + node=node, result=result, lock_host=True + ) + + self._update_evictable_leaf_sets(node) + return result + + def dec_host_lock_ref( + self, node: Any, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: + if self.disable: + return DecLockRefResult() + for component in self._components_tuple: + component.release_component_lock(node=node, params=params, lock_host=True) + + self._update_evictable_leaf_sets(node) + return DecLockRefResult() + def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs) -> None: if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs): return @@ -703,13 +784,15 @@ class UnifiedRadixCache(BasePrefixCache): cur_time -= 0.00001 node_update = node_update.parent - # Walk up to find last_host_node for full component. - if self.cache_controller is None: - last_host_node = best_match_device_node - else: - last_host_node = best_match_node - while last_host_node is not self.root_node and not last_host_node.backuped: - last_host_node = last_host_node.parent + # last_host_node will be used as the starting node for the subsequent + # `prefetch_from_storage` flow. We directly use best_match_node here, + # because best_match_node represents the node where all components + # have reached consensus on both device & host availability. + last_host_node = ( + best_match_node + if self.cache_controller is not None + else best_match_device_node + ) if best_match_device_value_len > 0: device_indices = torch.cat(value[:best_match_device_value_len]) @@ -744,6 +827,9 @@ class UnifiedRadixCache(BasePrefixCache): child.parent = new_node child.key = child.key[split_len:] + new_node.hash_value, child.hash_value = split_node_hash_value( + child.hash_value, split_len, self.page_size + ) for component in self._components_tuple: component.redistribute_on_node_split(new_parent=new_node, child=child) @@ -778,6 +864,8 @@ class UnifiedRadixCache(BasePrefixCache): new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone() parent.children[key.child_key(self.page_size)] = new_node self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value) + if self.enable_storage: + new_node.hash_value = compute_node_hash_values(new_node, self.page_size) self._update_evictable_leaf_sets(new_node) self._update_evictable_leaf_sets(parent) @@ -894,6 +982,58 @@ class UnifiedRadixCache(BasePrefixCache): self._inc_hit_count(target_node, params.chunked) return result + def _insert_helper_host( + self, + node: UnifiedTreeNode, + key: RadixKey, + host_value: torch.Tensor, + hash_value: list[str], + ) -> InsertResult: + total_len = len(key) + self._touch_node(node) + if total_len == 0: + return InsertResult(prefix_len=0, mamba_exist=True) + + child_key = key.child_key(self.page_size) + matched_length = 0 + while len(key) > 0 and child_key in node.children: + node = node.children[child_key] + self._touch_node(node) + prefix_len = node.key.match(key, page_size=self.page_size) + + key = key[prefix_len:] + host_value = host_value[prefix_len:] + hash_value = hash_value[prefix_len // self.page_size :] + matched_length += prefix_len + + if prefix_len < len(node.key): + node = self._split_node(node.key, node, prefix_len) + + if len(key): + child_key = key.child_key(self.page_size) + + result = InsertResult( + prefix_len=matched_length, + ) + if len(key) == 0: + if ( + node is not self.root_node + and node.component_data[BASE_COMPONENT_TYPE].host_value is not None + ): + result.inserted_host_node = node + return result + + new_node = UnifiedTreeNode(self.tree_components) + new_node.parent = node + new_node.key = key + new_node.hash_value = hash_value + new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone() + node.children[child_key] = new_node + self._update_evictable_leaf_sets(new_node) + self._update_evictable_leaf_sets(node) + result.inserted_host_node = new_node + return result + # ---- Evict Helpers ---- def _cascade_evict( @@ -1126,6 +1266,7 @@ class UnifiedRadixCache(BasePrefixCache): and self.cache_controller.write_policy == "write_back" ): self.write_backup(node, write_back=True) + self.writing_check(write_back=True) self._evict_to_host(node, tracker) return else: @@ -1368,6 +1509,511 @@ class UnifiedRadixCache(BasePrefixCache): if not node.backuped and node.hit_count >= self.write_through_threshold: self.write_backup(node) + def write_backup_storage(self, node: UnifiedTreeNode) -> None: + if ( + not self.enable_storage + or self.cache_controller is None + or not node.backuped + ): + return + + prefix_keys = None + if self.hicache_storage_pass_prefix_keys: + prefix_keys = node.get_prefix_hash_values(node.parent) + + comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} + for comp in self._components_tuple: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + transfers = comp.build_hicache_transfers( + node, + CacheTransferPhase.BACKUP_STORAGE, + ) + if transfers: + comp_xfers[comp.component_type] = transfers + + kv_xfer = PoolTransfer( + name=PoolName.KV, + host_indices=node.component_data[BASE_COMPONENT_TYPE].host_value, + keys=node.hash_value, + ) + sidecar_xfers = self._build_sidecar_transfers( + CacheTransferPhase.BACKUP_STORAGE, kv_xfer, comp_xfers + ) + aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + aux_xfers.extend(sidecar_xfers) + + operation_id = self.cache_controller.write_storage( + node.component_data[BASE_COMPONENT_TYPE].host_value, + node.key.token_ids, + node.hash_value, + prefix_keys, + extra_pools=aux_xfers or None, + ) + self.ongoing_backup[operation_id] = ( + node, + self.inc_host_lock_ref(node).to_dec_params(), + ) + + def prefetch_from_storage( + self, + req_id: str, + last_host_node: UnifiedTreeNode, + new_input_tokens: list[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[list[str]] = None, + ) -> None: + if not self.enable_storage or self.cache_controller is None: + return + + extra_key = last_host_node.key.extra_key if last_host_node.key else None + prefetch_key = RadixKey( + new_input_tokens, + extra_key=extra_key, + is_bigram=self.is_eagle, + ).page_aligned(self.page_size) + prefetch_length = len(prefetch_key) + if ( + prefetch_length < self.prefetch_threshold + or self.cache_controller.prefetch_rate_limited() + ): + return + + anchor_lock_params = self.inc_host_lock_ref(last_host_node).to_dec_params() + host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) + if host_indices is None: + self.evict_host(prefetch_length) + host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) + if host_indices is None: + available_size = self.cache_controller.mem_pool_host.available_size() + prefetch_length = available_size - (available_size % self.page_size) + if prefetch_length >= self.prefetch_threshold: + prefetch_key = prefetch_key[:prefetch_length] + host_indices = self.cache_controller.mem_pool_host.alloc( + prefetch_length + ) + else: + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + return + if host_indices is None: + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + return + + comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} + alloc_failed = False + for comp in self._components_tuple: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + transfers = comp.build_hicache_transfers( + last_host_node, + CacheTransferPhase.PREFETCH, + token_ids=prefetch_key.token_ids, + prefetch_tokens=len(prefetch_key), + last_hash=last_hash, + ) + if transfers == []: + alloc_failed = True + break + if transfers: + comp_xfers[comp.component_type] = transfers + kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) + sidecar_xfers = self._build_sidecar_transfers( + CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers + ) + if alloc_failed: + self.cache_controller.append_host_mem_release( + host_indices=host_indices, + extra_pools=[x for xfers in comp_xfers.values() for x in xfers], + ) + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + return + + aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + aux_xfers.extend(sidecar_xfers) + operation = self.cache_controller.prefetch( + req_id, + host_indices, + prefetch_key.token_ids, + last_hash, + prefix_keys, + extra_pools=aux_xfers or None, + ) + self.ongoing_prefetch[req_id] = ( + last_host_node, + prefetch_key, + host_indices, + operation, + anchor_lock_params, + comp_xfers, + ) + self.cache_controller.prefetch_tokens_occupied += len(prefetch_key) + + def _prefetch_timeout_check_linear_func(self, operation) -> bool: + return ( + time.monotonic() - operation.start_time + > self.prefetch_timeout_base + + len(operation.hash_value) * self.prefetch_timeout_per_page + ) + + def can_terminate_prefetch(self, operation) -> bool: + if self.prefetch_stop_policy == "best_effort": + return True + + if len(operation.hash_value) == 0: + completed = False + else: + completed = ( + operation.completed_tokens == len(operation.hash_value) * self.page_size + ) + + if self.prefetch_stop_policy == "wait_complete": + can_terminate = completed + elif self.prefetch_stop_policy == "timeout": + can_terminate = completed or self._prefetch_timeout_check_linear_func( + operation + ) + else: + return True + + operation_terminated = operation.is_terminated() + states = torch.tensor( + [1 - int(can_terminate), int(operation_terminated)], + dtype=torch.int, + ) + if self.tp_world_size > 1: + torch.distributed.all_reduce( + states, op=torch.distributed.ReduceOp.MAX, group=self.tp_group + ) + can_terminate = states[0].item() == 0 + operation_terminated = states[1].item() == 1 + return can_terminate or operation_terminated + + def check_prefetch_progress(self, req_id: str) -> bool: + if req_id not in self.ongoing_prefetch: + return True + + ( + last_host_node, + prefetch_key, + host_indices, + operation, + anchor_lock_params, + comp_xfers, + ) = self.ongoing_prefetch[req_id] + if operation.host_indices is None: + return True + if not self.can_terminate_prefetch(operation): + return False + + completed_tokens, hash_value = self.cache_controller.terminate_prefetch( + operation + ) + min_completed_tokens = completed_tokens + if self.tp_world_size > 1: + completed_tokens_tensor = torch.tensor( + min_completed_tokens, dtype=torch.int + ) + torch.distributed.all_reduce( + completed_tokens_tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, + ) + min_completed_tokens = int(completed_tokens_tensor.item()) + + fetched_key = prefetch_key[:min_completed_tokens] + insert_result = self._insert_helper_host( + last_host_node, + fetched_key, + host_indices[:min_completed_tokens], + hash_value[: min_completed_tokens // self.page_size], + ) + + for ct, xfers in comp_xfers.items(): + self.components[ct].commit_hicache_transfer( + last_host_node, + CacheTransferPhase.PREFETCH, + xfers, + insert_result=insert_result, + pool_storage_result=operation.pool_storage_result, + ) + + self.cache_controller.mem_pool_host.free( + host_indices[: insert_result.prefix_len] + ) + self.cache_controller.append_host_mem_release( + host_indices[min_completed_tokens:completed_tokens] + ) + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + del self.ongoing_prefetch[req_id] + self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) + + loaded_from_storage = min_completed_tokens - insert_result.prefix_len + self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage + logger.info( + "HiCache prefetch success req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d tail_release=%d occupied=%d", + req_id, + completed_tokens, + min_completed_tokens, + insert_result.prefix_len, + loaded_from_storage, + completed_tokens - min_completed_tokens, + self.cache_controller.prefetch_tokens_occupied, + ) + if self.enable_storage_metrics and self.storage_metrics_collector is not None: + self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage) + return True + + def terminate_prefetch(self, req_id: str) -> None: + if req_id not in self.ongoing_prefetch: + return + _, _, _, operation, _, _ = self.ongoing_prefetch[req_id] + if operation.host_indices is None: + return + operation.mark_terminate() + + def pop_prefetch_loaded_tokens(self, req_id: str) -> int: + return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) + + def release_aborted_request(self, rid: str) -> None: + self.prefetch_loaded_tokens_by_reqid.pop(rid, None) + if rid not in self.ongoing_prefetch: + return + + ( + last_host_node, + prefetch_key, + host_indices, + operation, + anchor_lock_params, + comp_xfers, + ) = self.ongoing_prefetch[rid] + if operation.host_indices is None: + return + + completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) + if self.tp_world_size > 1: + torch.distributed.barrier(group=self.tp_group) + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + del self.ongoing_prefetch[rid] + self.cache_controller.append_host_mem_release( + host_indices=host_indices[:completed_tokens], + extra_pools=[x for xfers in comp_xfers.values() for x in xfers], + ) + self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) + + def _drain_storage_control_queues_impl( + self, + n_revoke: Optional[int], + n_backup: Optional[int], + n_release: Optional[int], + extra_release_counts: Optional[dict[PoolName, int]], + log_metrics: bool, + ) -> None: + cc = self.cache_controller + + def _drain_queue(q, limit: Optional[int]): + drained = 0 + while limit is None or drained < limit: + try: + item = q.get_nowait() + except Empty: + break + drained += 1 + yield item + + def _drain_revoke(): + drained = 0 + for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke): + info = self.ongoing_prefetch.pop(req_id, None) + if info is None: + continue + drained += 1 + ( + last_host_node, + prefetch_key, + _host_indices, + _operation, + anchor_lock_params, + comp_xfers, + ) = info + cc.append_host_mem_release( + extra_pools=[x for xfers in comp_xfers.values() for x in xfers] + ) + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + cc.prefetch_tokens_occupied -= len(prefetch_key) + if cc.prefetch_tokens_occupied < 0: + cc.prefetch_tokens_occupied = 0 + return drained + + def _drain_backup(): + drained = 0 + for operation in _drain_queue(cc.ack_backup_queue, n_backup): + drained += 1 + entry = self.ongoing_backup.pop(operation.id, None) + if entry is not None: + node, lock_params = entry + self.dec_host_lock_ref(node, lock_params) + if ( + log_metrics + and self.enable_storage_metrics + and self.storage_metrics_collector is not None + ): + self.storage_metrics_collector.log_backuped_tokens( + operation.completed_tokens + ) + return drained + + def _drain_release(): + host_indices_list = [] + released_tokens = 0 + for host_indices in _drain_queue(cc.host_mem_release_queue, n_release): + host_indices_list.append(host_indices) + released_tokens += len(host_indices) + if host_indices_list: + cc.mem_pool_host.free(torch.cat(host_indices_list, dim=0)) + return len(host_indices_list), released_tokens + + def _drain_extra_release(): + drained: dict[PoolName, tuple[int, int]] = {} + if not extra_release_counts: + return drained + for pool_name, limit in extra_release_counts.items(): + release_queue = cc.extra_host_mem_release_queues.get(pool_name) + if release_queue is None: + continue + host_indices_list = [] + released_tokens = 0 + for host_indices in _drain_queue(release_queue, limit): + host_indices_list.append(host_indices) + released_tokens += len(host_indices) + if host_indices_list: + entry = cc.mem_pool_host.entry_map.get(pool_name) + if entry is not None: + entry.host_pool.free(torch.cat(host_indices_list, dim=0)) + drained[pool_name] = (len(host_indices_list), released_tokens) + return drained + + _drain_revoke() + _drain_backup() + _drain_release() + _drain_extra_release() + + def drain_storage_control_queues(self) -> None: + cc = self.cache_controller + extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {}) + extra_pool_names = list(extra_release_queues) + local_qsize_list = [ + cc.prefetch_revoke_queue.qsize(), + cc.ack_backup_queue.qsize(), + cc.host_mem_release_queue.qsize(), + *[ + extra_release_queues[pool_name].qsize() + for pool_name in extra_pool_names + ], + ] + qsizes = torch.tensor( + local_qsize_list, + dtype=torch.int, + ) + if self.tp_world_size > 1: + torch.distributed.all_reduce( + qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + ) + qsize_list = list(map(int, qsizes.tolist())) + n_revoke, n_backup, n_release = qsize_list[:3] + extra_release_counts = { + pool_name: count + for pool_name, count in zip(extra_pool_names, qsize_list[3:]) + } + self._drain_storage_control_queues_impl( + n_revoke=n_revoke, + n_backup=n_backup, + n_release=n_release, + extra_release_counts=extra_release_counts, + log_metrics=True, + ) + + def _apply_storage_runtime_config( + self, + *, + storage_backend: Optional[str], + prefetch_threshold: int, + prefetch_timeout_base: float, + prefetch_timeout_per_ki_token: float, + hicache_storage_pass_prefix_keys: bool, + enable_storage: bool, + enable_storage_metrics: bool, + extra_metric_labels: Optional[dict[str, str]], + ) -> None: + self.enable_storage = enable_storage + self.prefetch_threshold = prefetch_threshold + self.prefetch_timeout_base = prefetch_timeout_base + self.prefetch_timeout_per_page = ( + self.page_size / 1024 * prefetch_timeout_per_ki_token + ) + self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys + self.enable_storage_metrics = enable_storage_metrics + + if self.enable_storage_metrics: + attn_cp_rank, attn_cp_size = ( + self.cache_controller.get_attn_cp_rank_and_size() + ) + labels = { + "storage_backend": storage_backend, + "tp_rank": self.cache_controller.tp_rank, + "dp_rank": self.cache_controller.dp_rank, + "pp_rank": self.cache_controller.pp_rank, + "pp_size": self.cache_controller.pp_size, + "attn_cp_rank": attn_cp_rank, + "attn_cp_size": attn_cp_size, + } + if extra_metric_labels: + labels.update(extra_metric_labels) + existing_collector = self.storage_metrics_collector + if existing_collector is None: + self.storage_metrics_collector = StorageMetricsCollector(labels=labels) + elif set(existing_collector.labels.keys()) == set(labels.keys()): + existing_collector.labels = labels + else: + logger.warning( + "Storage metrics labels changed (%s -> %s). Keep existing labels to avoid duplicate metric registration.", + sorted(existing_collector.labels.keys()), + sorted(labels.keys()), + ) + else: + self.storage_metrics_collector = None + + def attach_storage_backend( + self, + storage_backend: str, + storage_backend_extra_config_json: Optional[str] = None, + served_model_name: Optional[str] = None, + hicache_storage_prefetch_policy: Optional[str] = None, + hicache_write_policy: Optional[str] = None, + ) -> tuple[bool, str]: + return ( + False, + "UnifiedRadixCache does not support runtime HiCache storage attach yet. " + "Configure hicache_storage_backend at startup instead.", + ) + + def detach_storage_backend(self) -> tuple[bool, str]: + return ( + False, + "UnifiedRadixCache does not support runtime HiCache storage detach yet. " + "Restart without hicache_storage_backend to disable it.", + ) + + def clear_storage_backend(self) -> bool: + try: + ok = self.cache_controller.clear_storage_backend() + except Exception as e: + logger.error("Failed to clear hierarchical cache storage backend: %s", e) + return False + if ok: + logger.info("Hierarchical cache storage backend cleared successfully!") + return ok + # ---- HiCache: Async Event Management ---- def writing_check(self, write_back: bool = False) -> None: @@ -1387,6 +2033,8 @@ class UnifiedRadixCache(BasePrefixCache): node, params = entry if params is not None: self.dec_lock_ref(node, params) + if self.enable_storage: + self.write_backup_storage(node) cc.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -1415,6 +2063,8 @@ class UnifiedRadixCache(BasePrefixCache): for ack_id in ack_list: node, params = self.ongoing_write_through.pop(ack_id) self.dec_lock_ref(node, params) + if self.enable_storage: + self.write_backup_storage(node) finish_count -= 1 def loading_check(self) -> None: @@ -1484,6 +2134,12 @@ class UnifiedRadixCache(BasePrefixCache): """Called per scheduler step to poll async HiCache events.""" self.writing_check() self.loading_check() + if self.enable_storage: + self.drain_storage_control_queues() + if self.enable_storage_metrics and self.storage_metrics_collector is not None: + self.storage_metrics_collector.log_storage_metrics( + self.cache_controller.storage_backend.get_stats() + ) def flush_write_through_acks(self) -> None: """Flush pending write-through acknowledgements.""" @@ -1883,10 +2539,6 @@ class UnifiedRadixCache(BasePrefixCache): logger.error(msg) self.pretty_print() raise AssertionError(msg) - logger.debug( - f"Sanity check PASSED: {len(all_nodes)} nodes, " - f"{len(self.tree_components)} components" - ) def _check_lru_linked_list( self, diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py index 9f30b9b33..1e0344b72 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py @@ -24,7 +24,7 @@ MAMBA_TRACK_INTERVAL = 128 DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8" DSV4_FLASH_LAUNCH_TIMEOUT = 3600 -register_cuda_ci(est_time=745, stage="base-c", runner_config="8-gpu-h200") +register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200") class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py index 799e94a31..16efd0220 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py @@ -28,8 +28,8 @@ GLM5_LAUNCH_TIMEOUT = 3600 register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True) -class GSM8KTwoPassMixin: - """Mixin: run GSM8K twice with flush in between, verify accuracy diff. +class AccuracyTwoPassMixin: + """Mixin: run an eval twice with flush in between, verify accuracy diff. Subclass must provide: - self.base_url @@ -38,9 +38,14 @@ class GSM8KTwoPassMixin: gsm8k_threshold: float = 0.90 num_gsm8k_questions: int = 200 - max_accuracy_diff: float = 0.02 gsm8k_parallel: int = 40 + mmlu_threshold: float = 0.75 + num_mmlu_examples: int = 200 + mmlu_num_threads: int = 32 + + max_accuracy_diff: float = 0.02 + def _run_gsm8k(self): from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k @@ -57,6 +62,19 @@ class GSM8KTwoPassMixin: metrics = run_few_shot_gsm8k(args) return metrics["accuracy"] + def _run_mmlu(self): + from sglang.test.run_eval import run_eval as run_simple_eval + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mmlu", + num_examples=self.num_mmlu_examples, + num_threads=self.mmlu_num_threads, + ) + metrics = run_simple_eval(args) + return metrics["score"] + def _flush_cache(self): response = requests.post( self.base_url + "/flush_cache", @@ -65,45 +83,52 @@ class GSM8KTwoPassMixin: ) response.raise_for_status() - def test_gsm8k_two_passes(self): - """Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff.""" + def _two_pass(self, name: str, run_fn, threshold: float): # First pass - acc1 = self._run_gsm8k() - print(f"[{self.__class__.__name__}] GSM8K pass 1 accuracy: {acc1:.3f}") + acc1 = run_fn() + print(f"[{self.__class__.__name__}] {name} pass 1 accuracy: {acc1:.3f}") self.assertGreaterEqual( acc1, - self.gsm8k_threshold, - f"Pass 1 accuracy {acc1:.3f} < threshold {self.gsm8k_threshold}", + threshold, + f"{name} pass 1 accuracy {acc1:.3f} < threshold {threshold}", ) # Flush cache self._flush_cache() # Second pass - acc2 = self._run_gsm8k() - print(f"[{self.__class__.__name__}] GSM8K pass 2 accuracy: {acc2:.3f}") + acc2 = run_fn() + print(f"[{self.__class__.__name__}] {name} pass 2 accuracy: {acc2:.3f}") self.assertGreaterEqual( acc2, - self.gsm8k_threshold, - f"Pass 2 accuracy {acc2:.3f} < threshold {self.gsm8k_threshold}", + threshold, + f"{name} pass 2 accuracy {acc2:.3f} < threshold {threshold}", ) - # Verify diff + # Verify diff (only fail when 2nd pass regressed) if acc1 > acc2: diff = abs(acc1 - acc2) print( - f"[{self.__class__.__name__}] Accuracy diff: {diff:.3f} " + f"[{self.__class__.__name__}] {name} accuracy diff: {diff:.3f} " f"(max allowed: {self.max_accuracy_diff})" ) self.assertLessEqual( diff, self.max_accuracy_diff, - f"Accuracy diff {diff:.3f} exceeds max {self.max_accuracy_diff} " + f"{name} accuracy diff {diff:.3f} exceeds max {self.max_accuracy_diff} " f"(pass1={acc1:.3f}, pass2={acc2:.3f})", ) + def test_gsm8k_two_passes(self): + """Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff.""" + self._two_pass("GSM8K", self._run_gsm8k, self.gsm8k_threshold) -class TestGLM5HiCacheL3GSM8K(GSM8KTwoPassMixin, CustomTestCase): + def test_mmlu_two_passes(self): + """Run MMLU twice with flush in between, verify accuracy diff <= max_accuracy_diff.""" + self._two_pass("MMLU", self._run_mmlu, self.mmlu_threshold) + + +class TestGLM5HiCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase): """GLM-5.1-FP8 + HiCache L3 (file backend), with HiRadixTree.""" @classmethod diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py new file mode 100644 index 000000000..99331421b --- /dev/null +++ b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py @@ -0,0 +1,77 @@ +import os +import shutil +import tempfile +import unittest + +from test_unified_radix_cache_kl_hicache_nightly import AccuracyTwoPassMixin + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" +MAMBA_TRACK_INTERVAL = 128 + +register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200") + + +class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): + """Mamba hybrid + HiCache L3 (file backend) + UnifiedRadixCache.""" + + @classmethod + def setUpClass(cls): + cls.model = MAMBA_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.hicache_dir = tempfile.mkdtemp(prefix="hicache_l3_mamba_") + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mem-fraction-static", + "0.85", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + str(MAMBA_TRACK_INTERVAL), + "--enable-hierarchical-cache", + "--hicache-ratio", + "2", + "--hicache-write-policy", + "write_through", + "--hicache-storage-prefetch-policy", + "wait_complete", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--hicache-storage-backend", + "file", + "--max-mamba-cache-size", + "500", + "--weight-loader-prefetch-checkpoints", + ], + env={ + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + if os.path.isdir(cls.hicache_dir): + shutil.rmtree(cls.hicache_dir, ignore_errors=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 819f2c9e6..a5e262a91 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -2274,13 +2274,6 @@ class UnifiedRadixCacheSuite: tokens = self._swa_anchor_chain_tokens(len(chain)) return tree, chain, n, y, x, tokens - def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self): - tree, _, n, y, x, tokens = self._swa_anchor_setup() - result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIs(result.best_match_node, x) - self.assertIs(result.last_device_node, n.parent) - self.assertIs(result.last_host_node, y) - def test_hicache_swa_load_back_anchored_on_best_match_node(self): tree, _, _, y, x, _ = self._swa_anchor_setup() ps = self.cfg.page_size