diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 51030ae1d..a0e5bcc1b 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -264,7 +264,12 @@ class HiCacheController: ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator - self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache() + mem_pool_device = token_to_kv_pool_allocator.get_kvcache() + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + if isinstance(mem_pool_device, HybridLinearKVPool): + mem_pool_device = mem_pool_device.full_kv_pool + self.mem_pool_device = mem_pool_device self.mem_pool_host = mem_pool_host self.write_policy = write_policy self.page_size = page_size diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 23d6dfc35..e6634ff8c 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -639,6 +639,7 @@ class Req(ReqDllmMixin): self.extend_logprob_start_len = 0 self.last_node: Any = None self.last_host_node: Any = None + self.last_host_backup_node: Any = None self.host_hit_length = 0 # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 @@ -857,7 +858,11 @@ class Req(ReqDllmMixin): # Whether request reached finished condition return self.finished_reason is not None - def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None): + def init_next_round_input( + self, + tree_cache: Optional[BasePrefixCache] = None, + cow_mamba: Optional[bool] = None, + ): if self.is_dllm(): self._init_fill_ids_for_dllm() self.determine_dllm_phase() @@ -873,23 +878,27 @@ class Req(ReqDllmMixin): token_ids = self.fill_ids[:max_prefix_len] if tree_cache is not None: + if cow_mamba is None: + cow_mamba = tree_cache.supports_mamba() match_result = tree_cache.match_prefix( MatchPrefixParams( key=RadixKey(token_ids=token_ids, extra_key=self.extra_key), req=self, - cow_mamba=tree_cache.supports_mamba(), + cow_mamba=cow_mamba, ) ) ( self.prefix_indices, self.last_node, self.last_host_node, + self.last_host_backup_node, self.host_hit_length, self.mamba_branching_seqlen, ) = ( match_result.device_indices, match_result.last_device_node, match_result.last_host_node, + match_result.last_host_backup_node, match_result.host_hit_length, match_result.mamba_branching_seqlen, ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a7d3c4550..f7641cc34 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -694,9 +694,20 @@ class Scheduler( logger.info("Using experimental C++ radix tree implementation.") self.tree_cache = RadixCacheCpp(params=params, server_args=server_args) elif self.enable_hierarchical_cache: - from sglang.srt.mem_cache.hiradix_cache import HiRadixCache + if self.is_hybrid_ssm: + from sglang.srt.mem_cache.hi_mamba_radix_cache import ( + HiMambaRadixCache, + ) - self.tree_cache = HiRadixCache(params=params, server_args=server_args) + self.tree_cache = HiMambaRadixCache( + params=params, server_args=server_args + ) + else: + from sglang.srt.mem_cache.hiradix_cache import HiRadixCache + + self.tree_cache = HiRadixCache( + params=params, server_args=server_args + ) self.tp_worker.register_hicache_layer_transfer_counter( self.tree_cache.cache_controller.layer_done_counter ) @@ -1707,22 +1718,25 @@ class Scheduler( def _prefetch_kvcache(self, req: Req): if self.enable_hicache_storage: - req.init_next_round_input(self.tree_cache) - if req.last_node.backuped: - # only to initiate the prefetch if the last node is backuped - # otherwise, the allocated GPU memory must be locked for integrity - last_hash = req.last_host_node.get_last_hash_value() + req.init_next_round_input(self.tree_cache, cow_mamba=False) + last_host_node = ( + req.last_host_backup_node + if req.last_host_backup_node is not None + else req.last_host_node + ) + if last_host_node.backuped or last_host_node is self.tree_cache.root_node: + last_hash = last_host_node.get_last_hash_value() matched_len = len(req.prefix_indices) + req.host_hit_length new_input_tokens = req.fill_ids[matched_len:] prefix_keys = ( - req.last_node.get_prefix_hash_values(req.last_node.parent) + last_host_node.get_prefix_hash_values(last_host_node.parent) if self.tree_cache.hicache_storage_pass_prefix_keys else None ) self.tree_cache.prefetch_from_storage( req.rid, - req.last_host_node, + last_host_node, new_input_tokens, last_hash, prefix_keys, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index a383c5aa2..0d6d5313c 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -97,6 +97,7 @@ class MatchResult(NamedTuple): last_host_node : The last TreeNode on the host that was matched. Note that if HiCache is not enabled, this **must** be the same as `last_device_node`. + last_host_backup_node: The deepest backuped node for prefetch from storage. host_hit_length : Length of the KV cache hit on the host, if applicable. 0 if HiCache is not enabled. mamba_branching_seqlen: The mamba radix cache branching point, which is the longest @@ -107,6 +108,7 @@ class MatchResult(NamedTuple): device_indices: torch.Tensor last_device_node: Any last_host_node: Any + last_host_backup_node: Any = None host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py new file mode 100644 index 000000000..235dc8c90 --- /dev/null +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -0,0 +1,1668 @@ +from __future__ import annotations + +import atexit +import heapq +import json +import logging +import os +import threading +import time +from queue import Empty +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import torch + +from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + EvictResult, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.mamba_radix_cache import ( + MambaRadixCache, + TreeNode, + get_last_access_time, +) +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool +from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost +from sglang.srt.mem_cache.radix_cache import ( + RadixKey, + compute_node_hash_values, + split_node_hash_value, +) +from sglang.srt.observability.metrics_collector import StorageMetricsCollector +from sglang.srt.utils import bind_to_closest_numa_node_cuda + +if TYPE_CHECKING: + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class HiMambaRadixCache(MambaRadixCache): + """Hierarchical cache for hybrid Mamba models. + + Only the Full (attention) KV cache is backed up to L2 (host) / L3 (storage). + Mamba states remain on device and are managed by LRU eviction. + """ + + def __init__(self, params: CacheInitParams, server_args: ServerArgs): + self._enable_metrics_flag = params.enable_metrics + if server_args.hicache_io_backend == "direct": + if server_args.hicache_mem_layout == "page_first": + server_args.hicache_mem_layout = "page_first_direct" + logger.warning( + "Page first layout is not supported with direct IO backend, " + "switching to page first direct layout" + ) + + if not server_args.disable_hicache_numa_detect: + bind_to_closest_numa_node_cuda() + + self.page_size = params.page_size + kvcache = params.token_to_kv_pool_allocator.get_kvcache() + + if isinstance(kvcache, HybridLinearKVPool): + kvcache = kvcache.full_kv_pool + self.kvcache = kvcache + + self.full_kv_pool_host = MHATokenToKVPoolHost( + self.kvcache, + server_args.hicache_ratio, + server_args.hicache_size, + params.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + + self.tp_group = params.tp_cache_group + self.tp_world_size = ( + 1 + if self.tp_group is None + else torch.distributed.get_world_size(group=self.tp_group) + ) + + self.enable_storage = server_args.hicache_storage_backend is not None + self.enable_storage_metrics = self.enable_storage and params.enable_metrics + self.extra_metric_labels = server_args.extra_metric_labels + + ( + extra_config, + prefetch_threshold, + prefetch_timeout_base, + prefetch_timeout_per_ki_token, + hicache_storage_pass_prefix_keys, + ) = self._parse_storage_backend_extra_config( + server_args.hicache_storage_backend_extra_config + ) + self.is_prefetch_timeout = self._prefetch_timeout_check_linear_func + self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy + + self.load_cache_event = threading.Event() + self.cache_controller = HiCacheController( + params.token_to_kv_pool_allocator, + self.full_kv_pool_host, + params.page_size, + self.tp_group, + load_cache_event=self.load_cache_event, + write_policy=server_args.hicache_write_policy, + io_backend=server_args.hicache_io_backend, + storage_backend=server_args.hicache_storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=server_args.served_model_name, + storage_backend_extra_config=extra_config, + pp_rank=params.pp_rank, + pp_size=params.pp_size, + ) + self._apply_storage_runtime_config( + storage_backend=server_args.hicache_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=self.enable_storage, + enable_storage_metrics=self.enable_storage_metrics, + extra_metric_labels=self.extra_metric_labels, + ) + + self.ongoing_write_through = {} + self.ongoing_load_back = {} + self.ongoing_prefetch = {} + self.ongoing_backup = {} + # track per-request tokens loaded from storage (L3 hits) + # key: request_id, value: number of tokens actually loaded from storage + self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {} + + self.write_through_threshold = ( + 1 if server_args.hicache_write_policy == "write_through" else 2 + ) + self.load_back_threshold = 10 + + self.evictable_full_device_leaves: set[TreeNode] = set() + self.evictable_full_host_leaves: set[TreeNode] = set() + + # Detach storage backend automatically on process shutdown + atexit.register(self.shutdown) + + super().__init__(params=params) + + def reset(self) -> None: + TreeNode.counter = 0 + self.cache_controller.reset() + self.full_kv_pool_host.clear() + self.ongoing_write_through = {} + self.ongoing_load_back = {} + self.ongoing_prefetch = {} + self.ongoing_backup = {} + self.prefetch_loaded_tokens_by_reqid.clear() + self.evictable_full_device_leaves.clear() + self.evictable_full_host_leaves.clear() + super().reset() + + def write_backup(self, node: TreeNode, write_back=False): + host_indices = self.cache_controller.write( + device_indices=node.value, + node_id=node.id, + ) + if host_indices is None: + self.evict_host(len(node.value)) + host_indices = self.cache_controller.write( + device_indices=node.value, + node_id=node.id, + ) + if host_indices is not None: + node.host_value = host_indices + assert len(node.host_value) > 0 + self.ongoing_write_through[node.id] = node + if not write_back: + # no need to lock nodes if write back + self.inc_lock_ref(node) + else: + return 0 + + return len(host_indices) + + def load_back( + self, node: TreeNode, mem_quota: Optional[int] = None + ) -> Optional[torch.Tensor]: + """Load full KV back from host.""" + last_hit_node = node + nodes_to_load = [] + + while node.evicted: + assert node.backuped, f"No backup on evicted node {node.id}" + nodes_to_load.insert(0, node) + node = node.parent + else: + ancestor_node = node + + delta = self.inc_lock_ref(ancestor_node) + + full_host_indices = torch.cat([n.host_value for n in nodes_to_load]) + if (len(full_host_indices) < self.load_back_threshold) or ( + len(full_host_indices) > mem_quota + delta + if mem_quota is not None + else False + ): + # skip loading back if the total size is too small or exceeding the memory quota + self.dec_lock_ref(ancestor_node) + return None + + full_device_indices = self.cache_controller.load( + host_indices=full_host_indices, + node_id=last_hit_node.id, + ) + if full_device_indices is None: + self.evict(EvictParams(num_tokens=len(full_host_indices))) + full_device_indices = self.cache_controller.load( + host_indices=full_host_indices, + node_id=last_hit_node.id, + ) + self.dec_lock_ref(ancestor_node) + if full_device_indices is None: + # no sufficient GPU memory to load back KV caches + return None + + offset = 0 + for n in nodes_to_load: + n_len = len(n.host_value) + n.value = full_device_indices[offset : offset + n_len].clone() + offset += n_len + + self.full_lru_list.insert_mru(n) + self.full_evictable_size_ += n_len + + if n.mamba_value is not None: + if self.mamba_lru_list.in_list(n): + self.mamba_lru_list.reset_node_mru(n) + else: + self.mamba_lru_list.insert_mru(n) + self.mamba_evictable_size_ += len(n.mamba_value) + + self._update_leaf_status(n) + + self._update_leaf_status(ancestor_node) + + self.inc_lock_ref(last_hit_node) + self.ongoing_load_back[last_hit_node.id] = last_hit_node + + return full_device_indices + + def init_load_back( + self, + last_node: TreeNode, + host_hit_length: int, + mem_quota: Optional[int] = None, + ): + if last_node.evicted: + loading_values = self.load_back(last_node, mem_quota) + if loading_values is not None: + logger.debug( + f"loading back {len(loading_values)} tokens for node {last_node.id}" + ) + return loading_values, last_node + + while last_node.evicted: + last_node = last_node.parent + + return ( + torch.empty((0,), dtype=torch.int64, device=self.device), + last_node, + ) + + def _inc_hit_count(self, node: TreeNode, chunked=False): + if self.cache_controller.write_policy == "write_back" or chunked: + return + node.hit_count += 1 + + if not node.backuped and node.hit_count >= self.write_through_threshold: + # write to host if the node is not backuped + self.write_backup(node) + + def writing_check(self, write_back=False): + if write_back: + # blocking till all write back complete + while len(self.ongoing_write_through) > 0: + for _, finish_event, ack_list in self.cache_controller.ack_write_queue: + finish_event.synchronize() + for ack_id in ack_list: + backuped_node = self.ongoing_write_through.pop(ack_id) + if self.enable_storage: + self.write_backup_storage(backuped_node) + self.cache_controller.ack_write_queue.clear() + assert len(self.ongoing_write_through) == 0 + return + + if len(self.ongoing_write_through) == 0: + return + + finish_count = 0 + for _, finish_event, ack_list in self.cache_controller.ack_write_queue: + if not finish_event.query(): + break + finish_count += 1 + + queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") + if self.tp_world_size > 1: + torch.distributed.all_reduce( + queue_size, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, + ) + finish_count = int(queue_size.item()) + + while finish_count > 0: + _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0) + finish_event.synchronize() + for ack_id in ack_list: + backuped_node = self.ongoing_write_through.pop(ack_id) + self.dec_lock_ref(backuped_node) + if self.enable_storage: + self.write_backup_storage(backuped_node) + finish_count -= 1 + + def loading_check(self): + finish_count = 0 + for _, finish_event, ack_list in self.cache_controller.ack_load_queue: + if not finish_event.query(): + # the KV cache loading is still ongoing + break + finish_count += 1 + for ack_id in ack_list: + end_node = self.ongoing_load_back.pop(ack_id) + self.dec_lock_ref(end_node) + + del self.cache_controller.ack_load_queue[:finish_count] + + def ready_to_load_host_cache(self) -> int: + return self.cache_controller.start_loading() + + def check_hicache_events(self): + self.writing_check() + self.loading_check() + + if self.enable_storage: + self.drain_storage_control_queues() + if self.enable_storage_metrics: + self.storage_metrics_collector.log_storage_metrics( + self.cache_controller.storage_backend.get_stats() + ) + + def _protect_host_node(self, node: TreeNode): + node.protect_host() + self.evictable_full_host_leaves.discard(node) + + def _release_host_node(self, node: TreeNode): + node.release_host() + if node.host_ref_counter == 0: + self._update_full_host_leaf_status(node) + + def _discard_from_leaf_sets(self, node: TreeNode): + self.evictable_full_device_leaves.discard(node) + self.evictable_full_host_leaves.discard(node) + + def _update_leaf_status(self, node: TreeNode): + self._update_full_device_leaf_status(node) + self._update_full_host_leaf_status(node) + + def _update_full_device_leaf_status(self, node: TreeNode): + if node == self.root_node or node.evicted or node.full_lock_ref > 0: + self.evictable_full_device_leaves.discard(node) + return + for child in node.children.values(): + if not child.evicted: + self.evictable_full_device_leaves.discard(node) + return + self.evictable_full_device_leaves.add(node) + + def _update_full_host_leaf_status(self, node: TreeNode): + if ( + not node.evicted + or not node.backuped + or node == self.root_node + or node.host_ref_counter > 0 + ): + self.evictable_full_host_leaves.discard(node) + return + for child in node.children.values(): + if child.evicted and child.backuped: + self.evictable_full_host_leaves.discard(node) + return + self.evictable_full_host_leaves.add(node) + + def _evict_to_host(self, node: TreeNode) -> int: + """Evict full KV to host. Mamba stays on device. Returns num evicted.""" + num_full = len(node.value) + + if not node.backuped: + if self.cache_controller.write_policy == "write_back": + self.write_backup(node, write_back=True) + self.writing_check(write_back=True) + + self.cache_controller.evict_device(node.value) + self.full_evictable_size_ -= num_full + if self.full_lru_list.in_list(node): + self.full_lru_list.remove_node(node) + + node.value = None + self._update_leaf_status(node) + self._update_full_device_leaf_status(node.parent) + return num_full + + def _delete_tombstone_leaf(self, node: TreeNode) -> None: + """Remove a tombstone leaf from the tree and free HiCache resources.""" + assert node.mamba_value is None, f"node has mamba value, {node.id=}" + assert len(node.children) == 0, f"leaf node has children, {node.id=}" + parent = node.parent + key = self.get_child_key_fn(node.key) + v = parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + + self._discard_from_leaf_sets(node) + + if node.backuped and node.host_ref_counter == 0: + self.cache_controller.evict_host(node.host_value) + node.host_value = None + + self._update_leaf_status(parent) + + def _iteratively_delete_tombstone_leaf( + self, node: TreeNode + ) -> Tuple[TreeNode, int, int]: + """Cascade-delete tombstone parents that became leaves. + Returns (node, full_num_evicted, mamba_num_evicted). + """ + full_num_evicted = 0 + mamba_num_evicted = 0 + + while len(node.parent.children) == 0: + if node.parent == self.root_node: + break + if node.parent.mamba_value is not None: + break + if node.parent.full_lock_ref > 0 or node.parent.mamba_lock_ref > 0: + break + + parent = node.parent + + if not parent.evicted: + full_num_evicted += len(parent.value) + self.full_evictable_size_ -= len(parent.value) + self.cache_controller.evict_device(parent.value) + if self.full_lru_list.in_list(parent): + self.full_lru_list.remove_node(parent) + + self._discard_from_leaf_sets(parent) + self._delete_tombstone_leaf(parent) + node = parent + + return node, full_num_evicted, mamba_num_evicted + + def evict(self, params: EvictParams) -> EvictResult: + if self.disable: + return EvictResult() + + full_num_tokens = params.num_tokens + full_num_evicted = 0 + mamba_num_evicted = 0 + + if full_num_tokens > 0: + leaves = list(self.evictable_full_device_leaves) + eviction_heap = [(n.last_access_time, n) for n in leaves] + heapq.heapify(eviction_heap) + + while full_num_evicted < full_num_tokens and eviction_heap: + _, x = heapq.heappop(eviction_heap) + if x not in self.evictable_full_device_leaves: + continue + + evicted_full = self._evict_to_host(x) + full_num_evicted += evicted_full + + parent = x.parent + if parent in self.evictable_full_device_leaves: + heapq.heappush(eviction_heap, (parent.last_access_time, parent)) + + if params.mamba_num > 0: + mamba_num_evicted = self.evict_mamba(params.mamba_num) + + return EvictResult( + num_tokens_evicted=full_num_evicted, + mamba_num_evicted=mamba_num_evicted, + ) + + def evict_host(self, num_tokens: int): + if self.enable_storage: + host_leaves = list(self.evictable_full_host_leaves) + heap = [(n.last_access_time, n) for n in host_leaves] + heapq.heapify(heap) + + num_evicted = 0 + while num_evicted < num_tokens and heap: + _, x = heapq.heappop(heap) + if x not in self.evictable_full_host_leaves: + continue + + num_evicted += self.cache_controller.evict_host(x.host_value) + x.host_value = None + + self.evictable_full_host_leaves.discard(x) + self._update_full_host_leaf_status(x.parent) + if x.parent in self.evictable_full_host_leaves: + heapq.heappush(heap, (x.parent.last_access_time, x.parent)) + return + + # Non-L3 path: evict host leaves and clean up tree + heap = [(n.last_access_time, n) for n in self.evictable_full_host_leaves] + heapq.heapify(heap) + + num_evicted = 0 + while num_evicted < num_tokens and heap: + _, x = heapq.heappop(heap) + if x not in self.evictable_full_host_leaves: + continue + + num_evicted += self.cache_controller.evict_host(x.host_value) + x.host_value = None + + if x.mamba_value is not None: + if self.mamba_lru_list.in_list(x): + self.mamba_lru_list.remove_node(x) + self.mamba_evictable_size_ -= len(x.mamba_value) + self.req_to_token_pool.mamba_pool.free(x.mamba_value) + x.mamba_value = None + + self.evictable_full_host_leaves.discard(x) + + parent = x.parent + child_key = self.get_child_key_fn(x.key) + v = parent.children.pop(child_key, None) + assert v == x, f"parent does not have child key, {x.id=}" + + self._update_leaf_status(parent) + if parent in self.evictable_full_host_leaves: + heapq.heappush(heap, (parent.last_access_time, parent)) + + def evict_mamba(self, mamba_num: int) -> int: + if self.disable or mamba_num <= 0: + return 0 + + x = self.mamba_lru_list.get_lru_no_lock() + mamba_num_evicted = 0 + while mamba_num_evicted < mamba_num and self.mamba_lru_list.in_list(x): + assert x.mamba_value is not None, f"node has no mamba value, {x.id=}" + assert x != self.root_node, f"root node is not evictable, {x.id=}" + assert x.mamba_lock_ref == 0, f"node is in use, {x.id=}" + + if x.evicted: + self.req_to_token_pool.mamba_pool.free(x.mamba_value) + mamba_num_evicted += len(x.mamba_value) + x_next = self.mamba_lru_list.get_prev_no_lock(x) + self.mamba_lru_list.remove_node(x) + self.mamba_evictable_size_ -= len(x.mamba_value) + x.mamba_value = None + + if len(x.children) == 0: + self._delete_tombstone_leaf(x) + _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) + mamba_num_evicted += cascade_mamba + elif len(x.children) > 0: + self.req_to_token_pool.mamba_pool.free(x.mamba_value) + mamba_num_evicted += len(x.mamba_value) + x_next = self.mamba_lru_list.get_prev_no_lock(x) + self.mamba_lru_list.remove_node(x) + self._tombstone_internal_node(x) + else: + assert ( + x.full_lock_ref == 0 + ), f"evict leaf node invalid with {x.id=} {x.full_lock_ref=}" + + if not x.backuped: + if self.cache_controller.write_policy == "write_back": + self.write_backup(x, write_back=True) + self.writing_check(write_back=True) + + self.cache_controller.evict_device(x.value) + self.full_evictable_size_ -= len(x.value) + + self.req_to_token_pool.mamba_pool.free(x.mamba_value) + mamba_num_evicted += len(x.mamba_value) + + x_next = self.mamba_lru_list.get_prev_no_lock(x) + if self.full_lru_list.in_list(x): + self.full_lru_list.remove_node(x) + self.mamba_lru_list.remove_node(x) + self.mamba_evictable_size_ -= len(x.mamba_value) + + if x.backuped: + self.cache_controller.evict_host(x.host_value) + x.host_value = None + + x.value = None + x.mamba_value = None + + self._discard_from_leaf_sets(x) + + parent = x.parent + child_key = self.get_child_key_fn(x.key) + v = parent.children.pop(child_key, None) + assert v == x, f"parent does not have child key, {x.id=}" + + self._update_leaf_status(parent) + _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) + mamba_num_evicted += cascade_mamba + + if not self.mamba_lru_list.in_list(x_next): + x_next = self.mamba_lru_list.get_lru_no_lock() + + x = x_next + + return mamba_num_evicted + + def _unevict_node(self, node: TreeNode, fresh_value: torch.Tensor): + """Restore an evicted node with fresh device KV from the request.""" + n = len(fresh_value) + + node.value = fresh_value.clone() + self.full_lru_list.insert_mru(node) + self.full_evictable_size_ += n + + if node.mamba_value is not None: + if self.mamba_lru_list.in_list(node): + self.mamba_lru_list.reset_node_mru(node) + else: + self.mamba_lru_list.insert_mru(node) + self.mamba_evictable_size_ += len(node.mamba_value) + + self._update_leaf_status(node) + if node.parent is not None: + self._update_leaf_status(node.parent) + + def _insert_helper( + self, + node: TreeNode, + key: RadixKey, + value, + mamba_value, + chunked: bool = False, + ) -> Tuple[int, bool]: + assert mamba_value is not None, "Mamba value should not be None here." + node.last_access_time = get_last_access_time() + if node != self.root_node: + if not node.evicted: + self.full_lru_list.reset_node_mru(node) + if node.mamba_value is not None: + self.mamba_lru_list.reset_node_mru(node) + if len(key) == 0: + return 0, True + + child_key = self.get_child_key_fn(key) + + total_prefix_length = 0 + while len(key) > 0 and child_key in node.children.keys(): + node = node.children[child_key] + node.last_access_time = get_last_access_time() + + if not node.evicted: + self.full_lru_list.reset_node_mru(node) + if node.mamba_value is not None: + self.mamba_lru_list.reset_node_mru(node) + + prefix_len = self.key_match_fn(node.key, key) + + if prefix_len < len(node.key): + new_node = self._split_node(node.key, node, prefix_len) + node = new_node + + if node.evicted: + # Unevicted nodes take ownership of the request's KV pages. + # Do NOT count them in total_prefix_length, otherwise + # cache_finished_req / cache_unfinished_req will free those + # pages even though the tree now references them. + self._unevict_node(node, value[:prefix_len]) + else: + total_prefix_length += prefix_len + self._inc_hit_count(node, chunked) + + key = key[prefix_len:] + value = value[prefix_len:] + + if len(key): + child_key = self.get_child_key_fn(key) + + mamba_value_exist = False + if len(key): + new_node = self._add_new_node(node, key, value, mamba_value) + self._inc_hit_count(new_node, chunked) + elif node.mamba_value is None: + node.mamba_value = mamba_value + if not node.evicted: + self.full_lru_list.reset_node_mru(node) + self.mamba_lru_list.insert_mru(node) + self.mamba_evictable_size_ += len(mamba_value) + node.last_access_time = get_last_access_time() + else: + mamba_value_exist = True + if not node.evicted: + self.full_lru_list.reset_node_mru(node) + self.mamba_lru_list.reset_node_mru(node) + node.last_access_time = get_last_access_time() + + return total_prefix_length, mamba_value_exist + + def _add_new_node( + self, + parent: TreeNode, + key: RadixKey, + value: torch.Tensor, + mamba_value: torch.Tensor, + ) -> TreeNode: + child_key = self.get_child_key_fn(key) + new_node = TreeNode() + new_node.parent = parent + new_node.key = key + new_node.value = value.clone() + new_node.mamba_value = mamba_value + self.full_lru_list.insert_mru(new_node) + self.mamba_lru_list.insert_mru(new_node) + parent.children[child_key] = new_node + self.full_evictable_size_ += len(value) + self.mamba_evictable_size_ += len(mamba_value) + if self.enable_storage: + new_node.hash_value = compute_node_hash_values(new_node, self.page_size) + self._update_full_device_leaf_status(new_node) + self._update_full_device_leaf_status(parent) + return new_node + + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: + key = params.key + + if self.disable or len(key) == 0: + return MatchResult( + device_indices=torch.empty((0,), dtype=torch.int64, device=self.device), + last_device_node=self.root_node, + last_host_node=self.root_node, + last_host_backup_node=self.root_node, + host_hit_length=0, + ) + + if self.page_size != 1: + page_aligned_len = len(key) // self.page_size * self.page_size + key = key[:page_aligned_len] + + value, best_last_node, best_value_len, deepest_node = self._match_prefix_helper( + key + ) + return self._match_post_processor( + params, value, best_last_node, best_value_len, deepest_node + ) + + def _match_prefix_helper( + self, key: RadixKey + ) -> Tuple[List[torch.Tensor], TreeNode, int, TreeNode]: + """Walk tree to find best_last_node (mamba boundary) and deepest node.""" + node = self.root_node + child_key = self.get_child_key_fn(key) + + value: List[torch.Tensor] = [] + best_value_len = 0 + best_last_node = node + + while len(key) > 0 and child_key in node.children.keys(): + child = node.children[child_key] + + if child.evicted and not child.backuped: + break + + if node.mamba_value is not None: + best_value_len = len(value) + best_last_node = node + + prefix_len = self.key_match_fn(child.key, key) + if prefix_len < len(child.key): + new_node = self._split_node(child.key, child, prefix_len) + if not new_node.evicted: + value.append(new_node.value) + node = new_node + break + else: + if not child.evicted: + value.append(child.value) + node = child + key = key[prefix_len:] + if len(key): + child_key = self.get_child_key_fn(key) + + if node.mamba_value is not None: + best_value_len = len(value) + best_last_node = node + + deepest_node = node + return value, best_last_node, best_value_len, deepest_node + + def _match_post_processor( + self, + params: MatchPrefixParams, + value: List[torch.Tensor], + best_last_node: TreeNode, + best_value_len: int, + deepest_node: TreeNode, + ) -> MatchResult: + cow_mamba = params.cow_mamba + req = params.req + + # Full LRU: skip evicted nodes for full_lru_list + lru_node = best_last_node + while lru_node != self.root_node and lru_node.evicted: + lru_node = lru_node.parent + self.full_lru_list.reset_node_and_parents_mru(lru_node, self.root_node) + self.mamba_lru_list.reset_node_and_parents_mru(best_last_node, self.root_node) + + cur_time = get_last_access_time() + node_update = best_last_node + while node_update: + node_update.last_access_time = cur_time + cur_time -= 0.00001 + node_update = node_update.parent + + if len(value) > best_value_len: + from sglang.srt.server_args import get_global_server_args + + mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size + mamba_cache_chunk_aligned_seqlen = ( + sum(len(v) for v in value) // mamba_cache_chunk_size + ) * mamba_cache_chunk_size + mamba_branching_seqlen = ( + mamba_cache_chunk_aligned_seqlen + if mamba_cache_chunk_aligned_seqlen > 0 + else None + ) + else: + mamba_branching_seqlen = None + + # last_device_node & host_hit_length: from best_last_node (mamba boundary) + host_hit_length = 0 + last_device_node = best_last_node + while last_device_node.evicted: + host_hit_length += len(last_device_node.host_value) + last_device_node = last_device_node.parent + + last_host_node = best_last_node + + # last_host_backup_node: from deepest_node, find backuped ancestor + last_host_backup_node = deepest_node + while ( + last_host_backup_node is not self.root_node + and not last_host_backup_node.backuped + ): + last_host_backup_node = last_host_backup_node.parent + + mamba_node = best_last_node + if cow_mamba and mamba_node.mamba_value is not None: + if req.mamba_pool_idx is None: + dst_index = self.req_to_token_pool.mamba_pool.alloc(1) + if dst_index is None: + self.inc_lock_ref(mamba_node) + self.evict_mamba(1) + dst_index = self.req_to_token_pool.mamba_pool.alloc(1) + self.dec_lock_ref(mamba_node) + assert dst_index is not None, "Can not alloc mamba cache" + src_index = mamba_node.mamba_value + self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) + req.mamba_pool_idx = dst_index[0] + else: + src_index = mamba_node.mamba_value + dst_index = req.mamba_pool_idx.unsqueeze(0) + self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) + + value = value[:best_value_len] + if value: + value = torch.cat(value) + else: + value = torch.empty((0,), dtype=torch.int64, device=self.device) + + return MatchResult( + device_indices=value, + last_device_node=last_device_node, + last_host_node=last_host_node, + last_host_backup_node=last_host_backup_node, + host_hit_length=host_hit_length, + mamba_branching_seqlen=mamba_branching_seqlen, + ) + + def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: + if child.evicted: + return self._split_evicted_node(key, child, split_len) + + self.evictable_full_device_leaves.discard(child) + + new_node = super()._split_node(key, child, split_len) + + if child.backuped: + new_node.host_value = child.host_value[:split_len].clone() + child.host_value = child.host_value[split_len:].clone() + + new_node.hash_value, child.hash_value = split_node_hash_value( + child.hash_value, split_len, self.page_size + ) + + self._update_leaf_status(new_node) + self._update_leaf_status(child) + + return new_node + + def _split_evicted_node( + self, key: RadixKey, child: TreeNode, split_len: int + ) -> TreeNode: + self.evictable_full_host_leaves.discard(child) + + new_node = TreeNode() + new_node.children = {self.get_child_key_fn(key[split_len:]): child} + new_node.parent = child.parent + new_node.value = None + new_node.mamba_value = None + new_node.full_lock_ref = child.full_lock_ref + new_node.mamba_lock_ref = 0 + new_node.key = child.key[:split_len] + + if child.backuped: + new_node.host_value = child.host_value[:split_len].clone() + child.host_value = child.host_value[split_len:].clone() + + new_node.hash_value, child.hash_value = split_node_hash_value( + child.hash_value, split_len, self.page_size + ) + + child.last_access_time = get_last_access_time() + if child.mamba_value is not None: + self.mamba_lru_list.remove_node(child) + child.parent = new_node + child.key = child.key[split_len:] + new_node.parent.children[self.get_child_key_fn(key)] = new_node + if child.mamba_value is not None: + self.mamba_lru_list.insert_mru(child) + + self._update_full_host_leaf_status(new_node) + self._update_full_host_leaf_status(child) + + return new_node + + def _collect_all_nodes(self) -> list: + ret = [] + stack = [self.root_node] + while stack: + cur = stack.pop() + if not cur.evicted: + ret.append(cur) + stack.extend(cur.children.values()) + return ret + + def _collect_mamba_nontombstone_nodes(self) -> list: + ret = [] + stack = [self.root_node] + while stack: + cur = stack.pop() + if cur.mamba_value is not None: + ret.append(cur) + stack.extend(cur.children.values()) + return ret + + def all_values_flatten(self) -> torch.Tensor: + values = [] + + def _dfs(node: TreeNode): + for child in node.children.values(): + if not child.evicted: + values.append(child.value) + _dfs(child) + + _dfs(self.root_node) + return torch.cat(values) if values else torch.tensor([]) + + def sanity_check(self): + """Skip if async operations are pending (those nodes are still locked).""" + self.loading_check() + if self.ongoing_load_back or self.ongoing_write_through: + return + super().sanity_check() + + def inc_lock_ref(self, node: TreeNode) -> Optional[int]: + if self.disable: + return 0 + + delta = 0 + if node.mamba_value is not None: + if node.mamba_lock_ref == 0: + self.mamba_evictable_size_ -= len(node.mamba_value) + self.mamba_protected_size_ += len(node.mamba_value) + node.mamba_lock_ref += 1 + + while node != self.root_node: + if node.evicted: + node = node.parent + continue + + assert ( + node.full_lock_ref >= 0 + ), f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + if node.full_lock_ref == 0: + self.full_evictable_size_ -= len(node.value) + self.full_protected_size_ += len(node.value) + delta -= len(node.value) + self.evictable_full_device_leaves.discard(node) + node.full_lock_ref += 1 + node = node.parent + return delta + + def dec_lock_ref(self, node: TreeNode): + if self.disable: + return 0 + + delta = 0 + + if node.mamba_value is not None and node.mamba_lock_ref > 0: + if node.mamba_lock_ref == 1: + self.mamba_evictable_size_ += len(node.mamba_value) + self.mamba_protected_size_ -= len(node.mamba_value) + node.mamba_lock_ref -= 1 + + while node != self.root_node: + if node.evicted: + node = node.parent + continue + + assert ( + node.full_lock_ref > 0 + ), f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + if node.full_lock_ref == 1: + self.full_evictable_size_ += len(node.value) + self.full_protected_size_ -= len(node.value) + delta += len(node.value) + node.full_lock_ref -= 1 + if node.full_lock_ref == 0: + self._update_full_device_leaf_status(node) + node = node.parent + return delta + + # ---- L3 Support ---- + + def shutdown(self): + try: + if self.enable_storage: + self.detach_storage_backend() + except Exception: + logger.exception("Failed to detach storage backend on process shutdown.") + + 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: + prefetch_timeout_per_page = ( + self.page_size / 1024 * prefetch_timeout_per_ki_token + ) + + storage_metrics_collector = None + if enable_storage_metrics: + 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, + } + if extra_metric_labels: + labels.update(extra_metric_labels) + storage_metrics_collector = StorageMetricsCollector(labels=labels) + + self.enable_storage = enable_storage + self.prefetch_threshold = prefetch_threshold + self.prefetch_timeout_base = prefetch_timeout_base + self.prefetch_timeout_per_page = prefetch_timeout_per_page + self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys + self.enable_storage_metrics = enable_storage_metrics + if self.enable_storage_metrics: + self.storage_metrics_collector = storage_metrics_collector + 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]: + if hicache_storage_prefetch_policy is not None: + allowed = ["best_effort", "wait_complete", "timeout"] + if hicache_storage_prefetch_policy not in allowed: + return ( + False, + f"Invalid hicache_storage_prefetch_policy: {hicache_storage_prefetch_policy!r}. " + f"Expected one of {allowed}.", + ) + + if hicache_write_policy is not None: + allowed = ["write_back", "write_through", "write_through_selective"] + if hicache_write_policy not in allowed: + return ( + False, + f"Invalid hicache_write_policy: {hicache_write_policy!r}. " + f"Expected one of {allowed}.", + ) + + if self.enable_storage: + current_backend = self.cache_controller.storage_backend_type + if current_backend == storage_backend: + if hicache_storage_prefetch_policy is not None: + self.prefetch_stop_policy = hicache_storage_prefetch_policy + if hicache_write_policy is not None: + self.cache_controller.write_policy = hicache_write_policy + self.write_through_threshold = ( + 1 if hicache_write_policy == "write_through" else 2 + ) + return ( + True, + "HiCache storage backend already enabled with same backend; policies updated.", + ) + return ( + False, + f"HiCache storage backend is already enabled with backend '{current_backend}'. " + f"Cannot attach different backend '{storage_backend}'. Detach first.", + ) + + if hicache_storage_prefetch_policy is not None: + self.prefetch_stop_policy = hicache_storage_prefetch_policy + if hicache_write_policy is not None: + self.cache_controller.write_policy = hicache_write_policy + self.write_through_threshold = ( + 1 if hicache_write_policy == "write_through" else 2 + ) + + 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, + ) = self._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: + self.cache_controller.attach_storage_backend( + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=served_model_name, + storage_backend_extra_config=extra_config, + ) + 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_storage_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=self._enable_metrics_flag, + extra_metric_labels=self.extra_metric_labels, + ) + return True, "Attached HiCache storage backend successfully." + + def detach_storage_backend(self) -> tuple: + try: + self._drain_storage_control_queues_local() + self.cache_controller.detach_storage_backend() + except Exception as e: + logger.exception("Failed to detach storage backend.") + return False, f"Failed to detach HiCache storage backend: {e}" + + self._drain_storage_control_queues_local() + self._force_release_pending_storage_ops() + + self.enable_storage = False + self.enable_storage_metrics = False + if hasattr(self, "storage_metrics_collector"): + self.storage_metrics_collector = None + return True, "Detached HiCache storage backend successfully." + + def _force_release_pending_storage_ops(self): + cc = self.cache_controller + + try: + for req_id, info in list(self.ongoing_prefetch.items()): + try: + last_host_node, token_ids, host_indices, _operation = info + except Exception: + self.ongoing_prefetch.pop(req_id, None) + continue + try: + if host_indices is not None: + cc.mem_pool_host.free(host_indices) + except Exception: + logger.exception( + "Failed to free host indices for prefetch %s", req_id + ) + try: + self._release_host_node(last_host_node) + except Exception: + logger.exception( + "Failed to release host protection for prefetch %s", req_id + ) + try: + cc.prefetch_tokens_occupied -= len(token_ids) + if cc.prefetch_tokens_occupied < 0: + cc.prefetch_tokens_occupied = 0 + except Exception: + pass + self.ongoing_prefetch.pop(req_id, None) + except Exception: + logger.exception("Force release pending prefetch ops failed.") + + try: + for ack_id, node in list(self.ongoing_backup.items()): + try: + self._release_host_node(node) + except Exception: + logger.exception( + "Failed to release host protection for backup op %s", ack_id + ) + self.ongoing_backup.pop(ack_id, None) + except Exception: + logger.exception("Force release pending backup ops failed.") + + def _drain_storage_control_queues_local(self): + self._drain_storage_control_queues_impl( + n_revoke=None, + n_backup=None, + n_release=None, + log_metrics=False, + ) + + def _drain_storage_control_queues_impl( + self, + n_revoke: Optional[int], + n_backup: Optional[int], + n_release: Optional[int], + log_metrics: bool, + ): + 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(): + for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke): + info = self.ongoing_prefetch.pop(req_id, None) + if info is not None: + last_host_node, token_ids, _, _ = info + self._release_host_node(last_host_node) + cc.prefetch_tokens_occupied -= len(token_ids) + if cc.prefetch_tokens_occupied < 0: + cc.prefetch_tokens_occupied = 0 + + def _drain_backup(): + for operation in _drain_queue(cc.ack_backup_queue, n_backup): + ack_id = operation.id + entry = self.ongoing_backup.pop(ack_id, None) + if entry is not None: + self._release_host_node(entry) + if log_metrics and self.enable_storage_metrics: + self.storage_metrics_collector.log_backuped_tokens( + operation.completed_tokens + ) + + def _drain_release(): + host_indices_list = [] + for host_indices in _drain_queue(cc.host_mem_release_queue, n_release): + host_indices_list.append(host_indices) + if host_indices_list: + host_indices = torch.cat(host_indices_list, dim=0) + cc.mem_pool_host.free(host_indices) + + _drain_revoke() + _drain_backup() + _drain_release() + + def _parse_storage_backend_extra_config( + self, storage_backend_extra_config: Optional[str] + ): + extra_config = {} + if storage_backend_extra_config: + try: + 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) + except Exception as e: + logger.error(f"Invalid backend extra config JSON: {e}") + raise e + + 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( + f"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 self.enable_storage: + try: + if hasattr(self.cache_controller.storage_backend, "clear"): + self.cache_controller.storage_backend.clear() + logger.info( + "Hierarchical cache storage backend cleared successfully!" + ) + return True + else: + logger.warning( + f"Storage backend " + f"{type(self.cache_controller.storage_backend).__name__} " + "does not support clear operation." + ) + return False + except Exception as e: + logger.error(f"Failed to clear hierarchical cache storage backend: {e}") + return False + else: + logger.warning("Hierarchical cache storage backend is not enabled.") + return False + + def drain_storage_control_queues(self): + cc = self.cache_controller + + qsizes = torch.tensor( + [ + cc.prefetch_revoke_queue.qsize(), + cc.ack_backup_queue.qsize(), + cc.host_mem_release_queue.qsize(), + ], + dtype=torch.int, + ) + if self.tp_world_size > 1: + torch.distributed.all_reduce( + qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + ) + + n_revoke, n_backup, n_release = map(int, qsizes.tolist()) + self._drain_storage_control_queues_impl( + n_revoke=n_revoke, + n_backup=n_backup, + n_release=n_release, + log_metrics=True, + ) + + def _prefetch_timeout_check_linear_func(self, operation: PrefetchOperation): + 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: PrefetchOperation): + can_terminate = True + + if self.prefetch_stop_policy == "best_effort": + return can_terminate + + 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.is_prefetch_timeout(operation) + else: + return True + + operation_terminated = operation.is_terminated() + if self.tp_world_size > 1: + states = torch.tensor( + [1 - int(can_terminate), int(operation_terminated)], + dtype=torch.int, + ) + 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 + can_terminate = can_terminate or operation_terminated + return can_terminate + + def terminate_prefetch(self, req_id: str): + 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 write_backup_storage(self, node: TreeNode): + prefix_keys = ( + node.get_prefix_hash_values(node.parent) + if self.hicache_storage_pass_prefix_keys + else None + ) + + operation_id = self.cache_controller.write_storage( + node.host_value, + node.key, + node.hash_value, + prefix_keys, + ) + self.ongoing_backup[operation_id] = node + self._protect_host_node(node) + + def prefetch_from_storage( + self, + req_id: str, + last_host_node: TreeNode, + new_input_tokens: List[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[List[str]] = None, + ): + prefetch_length = len(new_input_tokens) - ( + len(new_input_tokens) % self.page_size + ) + new_input_tokens = new_input_tokens[:prefetch_length] + if ( + not self.enable_storage + or prefetch_length < self.prefetch_threshold + or self.cache_controller.prefetch_rate_limited() + ): + return + + self._protect_host_node(last_host_node) + 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: + self._release_host_node(last_host_node) + return + operation = self.cache_controller.prefetch( + req_id, host_indices, new_input_tokens, last_hash, prefix_keys + ) + self.ongoing_prefetch[req_id] = ( + last_host_node, + new_input_tokens, + host_indices, + operation, + ) + self.cache_controller.prefetch_tokens_occupied += len(new_input_tokens) + + def check_prefetch_progress(self, req_id: str) -> bool: + if req_id not in self.ongoing_prefetch: + return True + + last_host_node, token_ids, host_indices, operation = 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 = completed_tokens_tensor.item() + fetched_token_ids = token_ids[:min_completed_tokens] + written_indices = host_indices[:min_completed_tokens] + matched_length = self._insert_helper_host( + last_host_node, + RadixKey( + token_ids=fetched_token_ids, + extra_key=last_host_node.key.extra_key, + ), + written_indices, + hash_value[: min_completed_tokens // self.page_size], + ) + + self.cache_controller.mem_pool_host.free(host_indices[:matched_length]) + self.cache_controller.append_host_mem_release( + host_indices[min_completed_tokens:completed_tokens] + ) + self._release_host_node(last_host_node) + del self.ongoing_prefetch[req_id] + self.cache_controller.prefetch_tokens_occupied -= len(token_ids) + + loaded_from_storage = min_completed_tokens - matched_length + self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage + + if self.enable_storage_metrics: + self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage) + + return True + + def _insert_helper_host( + self, node: TreeNode, key: RadixKey, host_value, hash_value + ): + node.last_access_time = get_last_access_time() + if len(key) == 0: + return 0 + + child_key = self.get_child_key_fn(key) + + matched_length = 0 + host_value_inserted = False + while len(key) > 0 and child_key in node.children.keys(): + node = node.children[child_key] + node.last_access_time = get_last_access_time() + if node != self.root_node and node.mamba_value is not None: + self.mamba_lru_list.reset_node_mru(node) + prefix_len = self.key_match_fn(node.key, key) + + if node.evicted and not node.backuped: + node.host_value = host_value[:prefix_len].clone() + host_value_inserted = True + self._update_full_host_leaf_status(node) + if node.parent is not None: + self._update_full_host_leaf_status(node.parent) + else: + assert not host_value_inserted, ( + f"matched node after host_value insertion would cause incorrect free, " + f"{node.id=} {matched_length=} {prefix_len=}" + ) + matched_length += prefix_len + + key = key[prefix_len:] + host_value = host_value[prefix_len:] + hash_value = hash_value[prefix_len // self.page_size :] + + if prefix_len < len(node.key): + new_node = self._split_node(node.key, node, prefix_len) + node = new_node + + if len(key): + child_key = self.get_child_key_fn(key) + + if len(key): + new_node = TreeNode() + new_node.parent = node + new_node.key = key + new_node.value = None + new_node.mamba_value = None + new_node.host_value = host_value.clone() + new_node.hash_value = hash_value + node.children[child_key] = new_node + self._update_full_host_leaf_status(new_node) + self._update_full_host_leaf_status(node) + return matched_length + + def release_aborted_request(self, rid: str): + self.prefetch_loaded_tokens_by_reqid.pop(rid, None) + + if rid not in self.ongoing_prefetch: + return + + last_host_node, token_ids, host_indices, operation = 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._release_host_node(last_host_node) + del self.ongoing_prefetch[rid] + self.cache_controller.append_host_mem_release(host_indices[:completed_tokens]) + self.cache_controller.prefetch_tokens_occupied -= len(token_ids) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 07f5278d9..93ebd07a1 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -21,7 +21,7 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache. import heapq from collections import defaultdict -from functools import partial +from functools import lru_cache, partial from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -82,8 +82,11 @@ class TreeNode: self.last_access_time = get_last_access_time() self.hit_count = 0 + self.host_ref_counter = 0 # store the host indices of KV cache self.host_value = None + # store hash values of each pages + self.hash_value: Optional[List[str]] = None # for lru list, invariant: # 1. prev has greater last_access_time @@ -104,6 +107,29 @@ class TreeNode: def backuped(self): return self.host_value is not None + def protect_host(self): + """Protect the host value from eviction.""" + self.host_ref_counter += 1 + + def release_host(self): + """Release the host value, allowing it to be evicted.""" + if self.host_ref_counter > 0: + self.host_ref_counter -= 1 + else: + raise RuntimeError("Host reference counter is already zero.") + + def get_last_hash_value(self) -> Optional[str]: + """Returns the hash value of the last page in this node.""" + 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: "TreeNode") -> List[str]: + if node is None or node.hash_value is None: + return [] + return node.get_prefix_hash_values(node.parent) + node.hash_value + def __lt__(self, other: "TreeNode"): return self.last_access_time < other.last_access_time @@ -457,7 +483,7 @@ class MambaRadixCache(BasePrefixCache): if value is None: value = torch.tensor([x for x in key.token_ids], dtype=torch.int64) prefix_len, mamba_exist = self._insert_helper( - self.root_node, key, value, mamba_value + self.root_node, key, value, mamba_value, params.chunked ) return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist) @@ -1063,6 +1089,7 @@ class MambaRadixCache(BasePrefixCache): key: RadixKey, value, mamba_value, + chunked: bool = False, ) -> Tuple[int, bool]: # Update the last access time from root to leaf, so that # mamba will tombstone the node closer to root first