From 41c018a9ec9f896f73809eb6ecbb10a7644a7229 Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:49:08 +0800 Subject: [PATCH] [UnifiedTree] feat: support runtime attach/detach (#35269) Co-authored-by: hzh0425 --- .../sglang/srt/managers/cache_controller.py | 11 +- .../unified_cache/storage_attachment.py | 400 ++++++++++++++++++ .../unified_cache/unified_tree_core.py | 21 + .../unified_tree_core_interface.py | 10 + .../srt/mem_cache/unified_radix_cache.py | 140 +++--- ...t_hicache_storage_runtime_attach_detach.py | 23 +- .../test_unified_radix_cache_unittest.py | 88 ++++ 7 files changed, 611 insertions(+), 82 deletions(-) create mode 100644 python/sglang/srt/mem_cache/unified_cache/storage_attachment.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 31ee7c5ac..240ee4d5f 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -321,6 +321,11 @@ class HiCacheController: # Dedicated stop event for storage background threads (prefetch/backup). self.storage_stop_event = threading.Event() + # Storage control queues, (re)created whenever the storage threads start. + self.prefetch_hit_queue: Optional[Queue[StorageOperation]] = None + self.ack_backup_queue: Optional[Queue[StorageOperation]] = None + self.host_mem_release_queue: Optional[Queue[torch.Tensor]] = None + self.device = self.mem_pool_device.device self.layer_num = self.mem_pool_device.layer_num self.layer_done_counter = LayerDoneCounter(self.layer_num) @@ -417,9 +422,9 @@ class HiCacheController: self.prefetch_queue = Queue() self.backup_queue = Queue() - self.prefetch_hit_queue: Queue[StorageOperation] = Queue() - self.ack_backup_queue: Queue[StorageOperation] = Queue() - self.host_mem_release_queue: Queue[torch.Tensor] = Queue() + self.prefetch_hit_queue = Queue() + self.ack_backup_queue = Queue() + self.host_mem_release_queue = Queue() self.prefetch_thread.start() self.backup_thread.start() diff --git a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py new file mode 100644 index 000000000..88ac41d49 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py @@ -0,0 +1,400 @@ +"""Runtime attach / detach of the HiCache (L3) storage backend. + +``UnifiedRadixCache`` owns the tree; this component owns the lifecycle of the +storage backend behind it: validating the requested policies, applying the storage +runtime config, starting and stopping the controller's storage threads, and +cleaning up the bookkeeping a half-finished prefetch / backup leaves behind. + +Keeping it beside the tree rather than inside it gives the three entry points -- +startup, the admin API, and the atexit hook -- one implementation. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) +from sglang.srt.observability.metrics_collector import ( + STAT_LOGGER_ROLE_STORAGE, + StorageMetricsCollector, + resolve_collector_class, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache + +logger = logging.getLogger(__name__) + +# Kept in sync with the server-args choices; validated here so an admin request +# carrying a bad policy is rejected before it can touch the controller. +_PREFETCH_POLICIES = ("best_effort", "wait_complete", "timeout") +_WRITE_POLICIES = ("write_back", "write_through", "write_through_selective") + + +class StorageAttachment: + """Attach / detach the storage backend of one ``UnifiedRadixCache``.""" + + def __init__(self, cache: UnifiedRadixCache): + self._cache = cache + + # ---- Lifecycle entry points ---- + + def attach( + 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]: + """Enable the storage backend at runtime. + + Starts the storage threads inside the cache controller and turns on the + prefetch / backup paths. The caller must ensure there are no running or + queued requests, so this cannot race the scheduler thread. + """ + cache = self._cache + + # Validate first: a rejected request must have no side effects. + invalid = self._validate_policies( + hicache_storage_prefetch_policy, hicache_write_policy + ) + if invalid is not None: + return False, invalid + + controller = cache.cache_controller + if controller is None: + return ( + False, + "HiCache is not initialized (no cache controller); " + "launch with --enable-hierarchical-cache to attach a backend.", + ) + + if cache.enable_storage: + current_backend = controller.storage_backend_type + if current_backend != storage_backend: + return ( + False, + f"HiCache storage backend is already enabled with backend " + f"'{current_backend}'. Cannot attach different backend " + f"'{storage_backend}'. Detach first.", + ) + # Same backend: the request degenerates to a policy update. + self._apply_policies(hicache_storage_prefetch_policy, hicache_write_policy) + return ( + True, + "HiCache storage backend already enabled with same backend; " + "policies updated.", + ) + + # Apply policies before the controller attach, so the storage threads + # observe the new values as soon as they start. + self._apply_policies(hicache_storage_prefetch_policy, hicache_write_policy) + + logger.info(f"Attaching HiCache storage backend: {storage_backend}") + try: + ( + extra_config, + prefetch_threshold, + prefetch_timeout_base, + prefetch_timeout_per_ki_token, + hicache_storage_pass_prefix_keys, + ) = HybridCacheController.parse_storage_backend_extra_config( + storage_backend_extra_config_json + ) + except Exception as e: + logger.exception(f"Failed to parse storage_backend_extra_config_json: {e}") + return ( + False, + f"Failed to parse storage_backend_extra_config_json " + f"'{storage_backend_extra_config_json}': {e}", + ) + + try: + controller.attach_storage_backend( + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=served_model_name, + storage_backend_extra_config=extra_config, + host_pools=controller.mem_pool_host.entries, + ) + except Exception as e: + logger.exception( + f"Failed to attach storage backend '{storage_backend}': {e}" + ) + return False, f"Failed to attach storage backend '{storage_backend}': {e}" + + self.apply_runtime_config( + storage_backend=storage_backend, + prefetch_threshold=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=True, + enable_storage_metrics=cache._enable_metrics_flag, + extra_metric_labels=cache.extra_metric_labels, + ) + return True, "Attached HiCache storage backend successfully." + + def detach(self) -> tuple[bool, str]: + """Disable the storage backend at runtime. + + The caller must ensure there are no running or queued requests. Ordering + matters and is the reason this is not just ``controller.detach()``: + + 1. drain the control queues while the bookkeeping is still intact, so + in-flight acks / releases can still be matched to their nodes -- + otherwise host pages and host locks leak; + 2. stop the storage threads; + 3. only then release whatever prefetch / backup is still tracked, since + nothing can race the bookkeeping once the threads are gone; + 4. drain once more, because step 3 only *queues* the host pages -- this + is what hands them back to the pool and its sidecars. + """ + cache = self._cache + controller = cache.cache_controller + if controller is None: + return False, "HiCache storage backend is not initialized." + + try: + cache.drain_storage_control_queues_local() + # Idempotent: ask the controller to clean up even when `enable_storage` + # is already False, since that may be leftover state from an earlier + # partial detach. + controller.detach_storage_backend() + except Exception as e: + logger.exception("Failed to detach storage backend.") + # Never crash the server for an admin operation. The controller raises + # while its threads are still alive, so leave `ongoing_*` untouched -- + # a retry must still be able to match their acks. + return False, f"Failed to detach HiCache storage backend: {e}" + + try: + self._release_pending_storage_ops() + cache.drain_storage_control_queues_local() + except Exception: + logger.exception("Failed post-detach cleanup of storage bookkeeping.") + + cache.enable_storage = False + cache.enable_storage_metrics = False + return True, "Detached HiCache storage backend successfully." + + def shutdown(self) -> None: + """Best-effort auto-detach on process shutdown. + + Keeps startup and runtime behavior consistent: a backend attached either + via CLI args or via the admin API is detached on exit. + """ + try: + if self._cache.enable_storage: + self.detach() + except Exception: + logger.exception("Failed to detach storage backend on process shutdown.") + + def clear(self) -> bool: + """Drop everything the backend has stored, keeping it attached.""" + try: + ok = self._cache.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 + + # ---- Runtime config ---- + + def apply_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: + """Publish the storage knobs onto the tree; the single storage-enable point. + + Both startup and runtime attach funnel through here, so anything that must + happen exactly when storage turns on belongs in this method. + """ + cache = self._cache + + # Nodes already in the tree were built with hashing off. Fill them in as + # storage turns on: a node hashed against an unhashed parent restarts the + # page hash chain mid-sequence, so its L3 keys would cover only a suffix of + # the prefix they claim to represent. + if enable_storage and not cache.enable_storage: + filled = cache.tree_core.backfill_missing_hash_values() + if filled: + logger.info( + "Hashed %d radix nodes that predate the storage backend.", filled + ) + + cache.enable_storage = enable_storage + cache.prefetch_threshold = prefetch_threshold + cache.prefetch_timeout_base = prefetch_timeout_base + cache.prefetch_timeout_per_page = ( + cache.page_size / 1024 * prefetch_timeout_per_ki_token + ) + cache.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys + cache.enable_storage_metrics = enable_storage_metrics + + if enable_storage_metrics: + cache.storage_metrics_collector = self._resolve_metrics_collector( + storage_backend, extra_metric_labels + ) + else: + cache.storage_metrics_collector = None + + def _resolve_metrics_collector( + self, + storage_backend: Optional[str], + extra_metric_labels: Optional[dict[str, str]], + ) -> Optional[StorageMetricsCollector]: + """Build the storage metrics collector, or relabel the existing one. + + A collector is created once and kept across detach / re-attach: building a + second one with the same labels would register duplicate metrics. + """ + cache = self._cache + controller = cache.cache_controller + attn_cp_rank, attn_cp_size = controller.get_attn_cp_rank_and_size() + labels = { + "storage_backend": storage_backend, + "tp_rank": controller.tp_rank, + "dp_rank": controller.dp_rank, + "pp_rank": controller.pp_rank, + "pp_size": 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 = cache.storage_metrics_collector + if existing_collector is None: + from sglang.srt.runtime_context import get_server_args + + storage_cls = resolve_collector_class( + get_server_args(), + STAT_LOGGER_ROLE_STORAGE, + StorageMetricsCollector, + ) + return storage_cls(labels=labels) + + if 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()), + ) + return existing_collector + + # ---- Internals ---- + + @staticmethod + def _validate_policies( + hicache_storage_prefetch_policy: Optional[str], + hicache_write_policy: Optional[str], + ) -> Optional[str]: + """The rejection reason for an invalid policy, or None when both are ok.""" + if ( + hicache_storage_prefetch_policy is not None + and hicache_storage_prefetch_policy not in _PREFETCH_POLICIES + ): + return ( + f"Invalid hicache_storage_prefetch_policy: " + f"{hicache_storage_prefetch_policy!r}. " + f"Expected one of {list(_PREFETCH_POLICIES)}." + ) + if ( + hicache_write_policy is not None + and hicache_write_policy not in _WRITE_POLICIES + ): + return ( + f"Invalid hicache_write_policy: {hicache_write_policy!r}. " + f"Expected one of {list(_WRITE_POLICIES)}." + ) + return None + + def _apply_policies( + self, + hicache_storage_prefetch_policy: Optional[str], + hicache_write_policy: Optional[str], + ) -> None: + cache = self._cache + if hicache_storage_prefetch_policy is not None: + cache.prefetch_stop_policy = hicache_storage_prefetch_policy + logger.info( + f"Set hicache_storage_prefetch_policy to " + f"{hicache_storage_prefetch_policy}" + ) + if hicache_write_policy is not None: + cache.cache_controller.write_policy = hicache_write_policy + cache.write_through_threshold = ( + 1 if hicache_write_policy == "write_through" else 2 + ) + cache.is_write_back = hicache_write_policy == "write_back" + logger.info(f"Set hicache_write_policy to {hicache_write_policy}") + + def _release_pending_storage_ops(self) -> None: + """Release the host pages and locks that still-tracked ops hold. + + A safety net only: the scheduler refuses a detach unless `ongoing_prefetch` + and `ongoing_backup` are already empty, so this normally finds nothing. It + earns its keep on the atexit path, which has no such guard. + + Two deliberate departures from `release_aborted_request`, which does the + same bookkeeping while the storage threads are still running: + + * no cross-rank barrier -- at process exit the other ranks may already be + gone, and a barrier there would hang; + * only the completed prefix of a prefetch is freed. The prefetch IO thread + owns ``host_indices[completed_tokens:]`` and frees it as it drains, so + freeing the whole range here would double-free whatever it already + returned. Leaking the tail of an interrupted prefetch is the safer + failure of the two. + """ + cache = self._cache + controller = cache.cache_controller + + for req_id in list(cache.ongoing_prefetch): + info = cache.ongoing_prefetch[req_id] + try: + if info.host_indices is None: + # Host pages were never allocated for this operation. + cache.revoke_pending_prefetch(req_id) + continue + completed_tokens, _ = controller.terminate_prefetch(info.operation) + del cache.ongoing_prefetch[req_id] + cache.dec_host_lock_ref(info.anchor_node_id, info.anchor_lock_params) + controller.append_host_mem_release( + host_indices=info.host_indices[:completed_tokens], + extra_pools=[ + x for xfers in info.comp_xfers.values() for x in xfers + ], + ) + controller.prefetch_tokens_occupied = max( + 0, controller.prefetch_tokens_occupied - len(info.prefetch_key) + ) + except Exception: + logger.exception("Failed to release pending prefetch %s", req_id) + cache.ongoing_prefetch.pop(req_id, None) + + for ack_id in list(cache.ongoing_backup): + node_id, lock_params = cache.ongoing_backup.pop(ack_id) + try: + cache.dec_host_lock_ref(node_id, lock_params) + except Exception: + logger.exception("Failed to release host lock for backup op %s", ack_id) + + cache.prefetch_loaded_tokens_by_reqid.clear() diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 16180016e..faa14da3e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -515,6 +515,27 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): """The hash values owned by this node, excluding its ancestors.""" return self.node_by_id(node_id).hash_value or [] + def backfill_missing_hash_values(self) -> int: + """Hash every node that was built while storage was disabled. + + Page hashes chain from the parent's last hash, so a node whose parent has + none restarts the chain mid-sequence: its keys then encode only a suffix + of the prefix they claim to represent, which would alias unrelated + requests that happen to start with those tokens. Walks parent-before-child + so each node hashes against an already-filled parent. Returns the number + of nodes filled. + """ + filled = 0 + stack = [self.root_node] + while stack: + node = stack.pop() + # The root anchors every chain and is seeded with an empty hash list. + if node is not self.root_node and node.hash_value is None: + node.hash_value = compute_node_hash_values(node, self.page_size) + filled += 1 + stack.extend(node.children.values()) + return filled + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: """The NodeId anchoring matches; the single root serves every namespace.""" return self.root_node.id diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index b7c63fb55..03c0c65fb 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -173,6 +173,16 @@ class UnifiedTreeCoreInterface(ABC): """The hash values owned by this node, excluding its ancestors.""" ... + @abstractmethod + def backfill_missing_hash_values(self) -> int: + """Hash every node built while storage was disabled; return how many. + + Called when a storage backend is attached at runtime: nodes already in + the tree carry no hash, and hashing their descendants against them would + restart the page hash chain mid-sequence. + """ + ... + @abstractmethod def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: """The NodeId anchoring matches for the namespace.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 9babb6f2f..5477edf2c 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1,5 +1,6 @@ from __future__ import annotations +import atexit import logging import threading import time @@ -69,6 +70,7 @@ from sglang.srt.mem_cache.unified_cache.components import ( from sglang.srt.mem_cache.unified_cache.session_ref_tracker import ( UnifiedSessionRefTracker, ) +from sglang.srt.mem_cache.unified_cache.storage_attachment import StorageAttachment from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401 NodeId, @@ -77,10 +79,8 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401 UnifiedTreeNode, ) from sglang.srt.observability.metrics_collector import ( - STAT_LOGGER_ROLE_STORAGE, StorageMetrics, StorageMetricsCollector, - resolve_collector_class, ) from sglang.srt.session.streaming_session import StreamingSession from sglang.srt.utils.common import ceil_align @@ -230,6 +230,8 @@ class UnifiedRadixCache(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 + # Owns the storage backend lifecycle; built by init_hicache. + self._storage_attachment: Optional[StorageAttachment] = None self.prefetch_stop_policy = "best_effort" self.prefetch_threshold = 256 self.prefetch_timeout_base = 1.0 @@ -457,8 +459,12 @@ class UnifiedRadixCache(BasePrefixCache): self.load_back_threshold = 10 self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy + # Runtime attach/detach of the L3 backend (startup, admin API, atexit). + self._storage_attachment = StorageAttachment(self) + atexit.register(self.shutdown) + if storage_backend is not None: - self._apply_storage_runtime_config( + self._storage_attachment.apply_runtime_config( storage_backend=storage_backend, prefetch_threshold=storage_prefetch_threshold, prefetch_timeout_base=prefetch_timeout_base, @@ -1734,7 +1740,7 @@ class UnifiedRadixCache(BasePrefixCache): return False if operation.host_indices is None: self.cache_controller.terminate_prefetch(operation) - self._revoke_pending_prefetch(req_id) + self.revoke_pending_prefetch(req_id) return True completed_tokens, hash_value = self.cache_controller.terminate_prefetch( @@ -1955,7 +1961,7 @@ class UnifiedRadixCache(BasePrefixCache): ) = self.ongoing_prefetch[rid] if operation.host_indices is None: self.cache_controller.terminate_prefetch(operation) - self._revoke_pending_prefetch(rid) + self.revoke_pending_prefetch(rid) return completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) @@ -2028,7 +2034,7 @@ class UnifiedRadixCache(BasePrefixCache): return len(host_indices) if host_indices is not None else 0 return len(prefetch_key) - def _revoke_pending_prefetch(self, req_id: str) -> None: + def revoke_pending_prefetch(self, req_id: str) -> None: info = self.ongoing_prefetch.pop(req_id, None) if info is None: return @@ -2090,7 +2096,7 @@ class UnifiedRadixCache(BasePrefixCache): if info is None: return True # aborted/cleaned; nothing to retry if operation.is_terminated(): - self._revoke_pending_prefetch(req_id) + self.revoke_pending_prefetch(req_id) return True if buffer_mode and cc.prefetch_rate_limited(): @@ -2116,7 +2122,7 @@ class UnifiedRadixCache(BasePrefixCache): if host_indices is None: if buffer_mode: return False - self._revoke_pending_prefetch(req_id) + self.revoke_pending_prefetch(req_id) return True operation.storage_hit_count = alloc_len @@ -2146,13 +2152,13 @@ class UnifiedRadixCache(BasePrefixCache): continue if operation.is_terminated(): # Aborted while the storage query was in flight. - self._revoke_pending_prefetch(req_id) + self.revoke_pending_prefetch(req_id) continue if operation.storage_hit_count < self.prefetch_threshold: # Below-threshold hit: classify + feed the L3 miss # accounting, then revoke (not enough benefit). self._account_prefetch_outcome(operation, revoked=True) - self._revoke_pending_prefetch(req_id) + self.revoke_pending_prefetch(req_id) continue self._invalidate_absent_from_hit_query(operation) self._account_prefetch_outcome(operation, revoked=False) @@ -2250,62 +2256,30 @@ class UnifiedRadixCache(BasePrefixCache): 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 + def drain_storage_control_queues_local(self) -> None: + """Drain the storage control queues without cross-rank synchronization. + + For the detach / shutdown path, where best-effort cleanup matters more than + keeping the drained counts identical across ranks. The prefetch-hit queue is + deliberately skipped: servicing it would allocate host pages for a prefetch + that can no longer complete. + """ + cc = self.cache_controller + # The storage queues are created by the controller when the storage threads + # start, so they are still None when a backend was never attached. + if cc is None or cc.prefetch_hit_queue is None: + return + self._drain_storage_control_queues_impl( + n_storage_hit=0, + n_backup=None, + n_release=None, + extra_release_counts={ + name: None for name in cc.extra_host_mem_release_queues + }, + log_metrics=False, ) - 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: - from sglang.srt.runtime_context import get_server_args - - storage_cls = resolve_collector_class( - get_server_args(), - STAT_LOGGER_ROLE_STORAGE, - StorageMetricsCollector, - ) - self.storage_metrics_collector = storage_cls(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 + # ---- HiCache: Storage backend lifecycle (delegated) ---- def attach_storage_backend( self, @@ -2315,30 +2289,40 @@ class UnifiedRadixCache(BasePrefixCache): 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.", + """Attach (enable) the HiCache storage backend at runtime.""" + if self._storage_attachment is None: + return ( + False, + "HiCache is not initialized; launch with " + "--enable-hierarchical-cache to attach a storage backend.", + ) + return self._storage_attachment.attach( + storage_backend=storage_backend, + storage_backend_extra_config_json=storage_backend_extra_config_json, + served_model_name=served_model_name, + hicache_storage_prefetch_policy=hicache_storage_prefetch_policy, + hicache_write_policy=hicache_write_policy, ) 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.", - ) + """Detach (disable) the HiCache storage backend at runtime.""" + if self._storage_attachment is None: + return False, "HiCache storage backend is not initialized." + return self._storage_attachment.detach() + + def shutdown(self) -> None: + """Best-effort auto-detach of the storage backend on process shutdown.""" + if self._storage_attachment is not None: + self._storage_attachment.shutdown() 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) + if self._storage_attachment is None: return False + ok = self._storage_attachment.clear() if ok: # L3 is empty now: every storage-presence belief is stale, and a # retained positive would skip that page's backup forever. self.storage_existence_cache.clear() - logger.info("Hierarchical cache storage backend cleared successfully!") return ok # ---- HiCache: Async Event Management ---- diff --git a/test/registered/hicache/test_hicache_storage_runtime_attach_detach.py b/test/registered/hicache/test_hicache_storage_runtime_attach_detach.py index 8dbb452c3..2794ba091 100644 --- a/test/registered/hicache/test_hicache_storage_runtime_attach_detach.py +++ b/test/registered/hicache/test_hicache_storage_runtime_attach_detach.py @@ -28,10 +28,13 @@ from sglang.test.test_utils import ( ) from sglang.utils import wait_for_http_ready -register_cuda_ci(est_time=139, stage="base-b", runner_config="2-gpu-large") +register_cuda_ci(est_time=210, stage="base-b", runner_config="2-gpu-large") class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase): + # Extra server env; subclasses use it to select a tree_cache implementation. + extra_env: dict = {} + @classmethod def setUpClass(cls): cls.temp_dir = tempfile.mkdtemp() @@ -60,6 +63,7 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase): "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir, # Make runs less flaky for CI/dev. "SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1", + **cls.extra_env, } @classmethod @@ -221,6 +225,10 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase): kill_process_tree(process1.pid) time.sleep(2) + self._check_attach_detach_lifecycle() + + def _check_attach_detach_lifecycle(self): + """Attach/detach lifecycle against a server that requires an admin key.""" # Phase B: WITH --admin-api-key, must provide Authorization: Bearer . admin_key = "sglang-test-admin-key" base_url2 = f"http://127.0.0.1:{find_available_port(int(self.base_url.rsplit(':', 1)[1]) + 1)}" @@ -363,5 +371,18 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase): time.sleep(2) +class TestUnifiedRadixCacheStorageRuntimeAttachDetach( + TestHiCacheStorageRuntimeAttachDetach +): + """Same runtime attach/detach lifecycle, backed by UnifiedRadixCache.""" + + extra_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"} + + def test_runtime_attach_detach(self): + # Admin-key gating (phase A of the base test) lives in the HTTP layer and is + # independent of the tree cache implementation, so only the lifecycle is run. + self._check_attach_detach_lifecycle() + + if __name__ == "__main__": unittest.main(verbosity=2) 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 5155fc87b..b05c1d110 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 @@ -65,6 +65,7 @@ from sglang.srt.mem_cache.unified_cache.components.tree_component import ( EvictLayer, TreeComponent, ) +from sglang.srt.mem_cache.unified_cache.storage_attachment import StorageAttachment from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( DecSwaLockOnlyResult, @@ -7136,5 +7137,92 @@ class TestSWAWindowUnderBigramKey(CustomTestCase): cache.sanity_check() +class TestUnifiedRadixCacheStorageAttachBackfill(CustomTestCase): + """Enabling a storage backend must hash nodes that predate it. + + Page hashes chain off the parent's last hash, so a node hashed against an + unhashed parent restarts the chain mid-sequence: its L3 keys then cover only a + suffix of the prefix they claim to represent, which aliases any unrelated + request that happens to start with those tokens. + """ + + cfg = CacheConfig( + page_size=4, + components=(ComponentType.FULL,), + kv_size=64, + max_context_len=64, + ) + + prefix_tokens = array("q", [1, 2, 3, 4, 5, 6, 7, 8]) + suffix_tokens = array("q", [9, 10, 11, 12]) + + def _build_two_level_tree(self, *, storage_on_from_the_start: bool): + """A parent node plus a child extending it, built with storage on or off.""" + cache, allocator, _ = build_fixture(self.cfg) + cache.enable_storage = storage_on_from_the_start + for tokens in (self.prefix_tokens, self.prefix_tokens + self.suffix_tokens): + value = allocator.alloc(len(tokens)) + self.assertIsNotNone(value) + cache.insert(InsertParams(key=RadixKey(tokens), value=value)) + return cache + + @staticmethod + def _hashes_by_token_ids(cache): + """Every non-root node's token ids mapped to its hash chain.""" + root = cache.tree_core.root_node + hashes = {} + stack = [root] + while stack: + node = stack.pop() + if node is not root: + hashes[tuple(node.key.token_ids)] = node.hash_value + stack.extend(node.children.values()) + return hashes + + def test_backfill_reproduces_hashing_from_the_start(self): + expected = self._hashes_by_token_ids( + self._build_two_level_tree(storage_on_from_the_start=True) + ) + self.assertEqual(len(expected), 2, "expected a parent and a child node") + + late = self._build_two_level_tree(storage_on_from_the_start=False) + self.assertTrue( + all(h is None for h in self._hashes_by_token_ids(late).values()), + "nodes built while storage was disabled must start unhashed", + ) + + self.assertEqual(late.tree_core.backfill_missing_hash_values(), len(expected)) + self.assertEqual( + self._hashes_by_token_ids(late), + expected, + "a backfilled chain must be identical to one hashed from the start", + ) + + def test_backfill_is_idempotent(self): + cache = self._build_two_level_tree(storage_on_from_the_start=True) + before = self._hashes_by_token_ids(cache) + self.assertEqual(cache.tree_core.backfill_missing_hash_values(), 0) + self.assertEqual(self._hashes_by_token_ids(cache), before) + + def test_enabling_storage_backfills_the_tree(self): + """The tree is hashed by the time `enable_storage` flips on.""" + cache = self._build_two_level_tree(storage_on_from_the_start=False) + StorageAttachment(cache).apply_runtime_config( + storage_backend="file", + prefetch_threshold=64, + prefetch_timeout_base=1.0, + prefetch_timeout_per_ki_token=0.25, + hicache_storage_pass_prefix_keys=False, + enable_storage=True, + enable_storage_metrics=False, + extra_metric_labels=None, + ) + self.assertTrue(cache.enable_storage) + self.assertTrue( + all(h for h in self._hashes_by_token_ids(cache).values()), + "every node must carry a hash chain once storage is enabled", + ) + + if __name__ == "__main__": unittest.main()