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.disaggregation.kv_events import StorageMedium from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, DecLockRefResult, EvictParams, EvictResult, IncLockRefResult, InitLoadBackParams, MatchPrefixParams, MatchResult, ) from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, PoolName, PoolTransfer, PrefetchTimeoutConfig, ) from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( PrefetchOperation, ) from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( attach_hybrid_pool_to_mamba_cache, ) from sglang.srt.mem_cache.mamba_radix_cache import ( LRUList, MambaRadixCache, TreeNode, get_last_access_time, ) from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import ( RadixKey, ) from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value from sglang.srt.observability.metrics_collector import ( STAT_LOGGER_ROLE_STORAGE, StorageMetricsCollector, resolve_collector_class, ) 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 HostLRUList(LRUList): def __init__(self): super().__init__(mamba=True) self.prv = "host_mamba_prev" self.nxt = "host_mamba_next" setattr(self.head, self.nxt, self.tail) setattr(self.tail, self.prv, self.head) def reset_node_mru(self, node): assert node.id in self.cache, f"Resetting node {node.id=} not in host mamba lru" assert ( node.mamba_host_value is not None ), f"Resetting host mamba tombstone node in lru list: {node.id=}" self._remove_node(node) self._add_node(node) def insert_mru(self, node): assert ( node.mamba_host_value is not None ), f"Inserting host mamba tombstone node in lru list: {node.id=}" assert ( node.id not in self.cache ), f"Inserting node {node.id=} already in host mamba lru list" self.cache[node.id] = node self._add_node(node) def remove_node(self, node: TreeNode): assert node.id in self.cache, f"Removing node {node.id=} not in host mamba lru" assert ( node.mamba_host_value is not None ), f"Removing host mamba tombstone node from lru list: {node.id=}" del self.cache[node.id] self._remove_node(node) class HiMambaRadixCache(MambaRadixCache): """Hierarchical cache for hybrid Mamba models.""" 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.override( "hicache.mem_layout_force", hicache_mem_layout="page_first_direct" ) logger.warning( "Page first layout is not supported with direct IO backend, " "switching to page first direct layout" ) self.page_size = params.page_size self.hybrid_kv_cache = params.token_to_kv_pool_allocator.get_kvcache() if not isinstance(self.hybrid_kv_cache, HybridLinearKVPool): raise ValueError( "HiMambaRadixCache requires HybridLinearKVPool for hybrid SSM models." ) if not isinstance(params.req_to_token_pool, HybridReqToTokenPool): raise ValueError( "HiMambaRadixCache requires HybridReqToTokenPool for hybrid SSM models." ) self.kvcache = self.hybrid_kv_cache.full_kv_pool 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_config, 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() attach_hybrid_pool_to_mamba_cache( self, params, server_args, extra_config=extra_config, prefetch_threshold=prefetch_threshold, load_cache_event=self.load_cache_event, enable_storage_metrics=self.enable_storage_metrics, attn_cp_group=params.attn_cp_cache_group, attn_tp_group=params.attn_tp_cache_group, ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, prefetch_threshold=prefetch_threshold, prefetch_timeout_config=prefetch_timeout_config, 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() self.mamba_host_lru_list = HostLRUList() # Detach storage backend automatically on process shutdown atexit.register(self.shutdown) super().__init__(params=params) def reset(self) -> None: TreeNode.counter = 0 self._flush_pending_storage_backups_before_reset() self.cache_controller.reset() self.full_kv_pool_host.clear() self.mamba_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() self.mamba_host_lru_list = HostLRUList() logger.info( "HiMambaRadixCache reset completed: host_kv_available=%s host_mamba_available=%s", self.full_kv_pool_host.available_size(), self.mamba_pool_host.available_size(), ) super().reset() def write_backup(self, node: TreeNode, write_back=False) -> int: # Backup invariant (for write-through mode): backed-up nodes must form a # contiguous prefix from root — no gaps. Skip if parent isn't backed # up yet; if not write_back and ( node.parent != self.root_node and not node.parent.backuped ): return 0 # If mamba host slot already exists, refresh its LRU position. if node.mamba_value is not None and node.mamba_host_value is not None: if self.mamba_host_lru_list.in_list(node): self.mamba_host_lru_list.reset_node_mru(node) extra_pools = self.mamba_backup_transfers(node) host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, extra_pools=extra_pools, ) 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, extra_pools=extra_pools, ) if host_indices is not None: node.host_value = host_indices.clone() if extra_pools is not None: self.mamba_backup_commit(node, extra_pools) 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, req=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 mamba_restore_nodes = [] if last_hit_node.mamba_backuped and last_hit_node.mamba_evicted: mamba_restore_nodes.append(last_hit_node) result = self.inc_lock_ref(ancestor_node) delta = result.delta if nodes_to_load: full_host_indices = torch.cat([n.host_value for n in nodes_to_load]) else: full_host_indices = torch.empty((0,), dtype=torch.int64, device="cpu") if ( len(full_host_indices) > 0 and ( (len(full_host_indices) < self.load_back_threshold) or ( len(full_host_indices) > mem_quota + delta if mem_quota is not None else False ) ) and len(mamba_restore_nodes) == 0 ): # skip loading back if the total size is too small or exceeding the memory quota self.dec_lock_ref(ancestor_node) return None logger.debug( f"Init load back from cpu -> gpu, kv hit length: {len(full_host_indices)}, mamba host hit length: {len(mamba_restore_nodes)}" ) mamba_pools = self.mamba_restore_transfers( last_hit_node, mamba_restore_nodes, req ) full_device_indices = self.cache_controller.load( host_indices=full_host_indices, node_id=last_hit_node.id, extra_pools=mamba_pools, ) if full_device_indices is None: if len(full_host_indices) > 0: self.evict(EvictParams(num_tokens=len(full_host_indices))) mamba_pools = self.mamba_restore_transfers( last_hit_node, mamba_restore_nodes, req ) full_device_indices = self.cache_controller.load( host_indices=full_host_indices, node_id=last_hit_node.id, extra_pools=mamba_pools, ) self.dec_lock_ref(ancestor_node) if full_device_indices is None: # no sufficient GPU memory to load back KV caches return None self.mamba_restore_commit(mamba_restore_nodes, mamba_pools) 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._record_store_event(n, medium=StorageMedium.GPU) self.full_lru_list.insert_mru(n) self.full_evictable_size_ += n_len self._update_leaf_status(n) for n in mamba_restore_nodes: 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(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, params: InitLoadBackParams, ): last_node = params.best_match_node mem_quota = params.mem_quota req = params.req if last_node.evicted or (last_node.mamba_evicted and last_node.mamba_backuped): loading_values = self.load_back(last_node, mem_quota, req=req) 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 is not self.root_node and ( last_node.evicted or last_node.mamba_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) self._record_store_event( backuped_node, medium=StorageMedium.CPU ) 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 # Every rank must enter the all_reduce below; ongoing_write_through can # diverge across ranks because loading_check processes DMA completions # independently (no cross-rank sync). finish_count = 0 if len(self.ongoing_write_through) > 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._record_store_event(backuped_node, medium=StorageMedium.CPU) self.dec_lock_ref(backuped_node) if self.enable_storage: self.write_backup_storage(backuped_node) finish_count -= 1 def loading_check(self): # Every rank must enter the all_reduce below; ongoing_load_back can # diverge across ranks. finish_count = 0 for _, finish_event, ack_list in self.cache_controller.ack_load_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_load_queue.pop(0) finish_event.synchronize() for ack_id in ack_list: end_node = self.ongoing_load_back.pop(ack_id) self.dec_lock_ref(end_node) finish_count -= 1 def ready_to_load_host_cache(self) -> int: return self.cache_controller.start_loading() def flush_write_through_acks(self) -> None: self.writing_check() 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, protect_mamba: bool = True): node.protect_host() self.evictable_full_host_leaves.discard(node) if protect_mamba: node.protect_host_mamba() if self.mamba_host_lru_list.in_list(node): self.mamba_host_lru_list.remove_node(node) def _release_host_node(self, node: TreeNode, release_mamba: bool = True): node.release_host() if release_mamba: node.release_host_mamba() if node.host_mamba_ref_counter == 0 and node.mamba_host_value is not None: if not self.mamba_host_lru_list.in_list(node): self.mamba_host_lru_list.insert_mru(node) if node.host_ref_counter == 0 and node.host_mamba_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 or node.host_mamba_ref_counter > 0 or len(node.children) > 0 ): self.evictable_full_host_leaves.discard(node) return self.evictable_full_host_leaves.add(node) def _free_device_mamba(self, node: TreeNode) -> int: if node.mamba_value is None: return 0 mamba_num = len(node.mamba_value) self.req_to_token_pool.mamba_allocator.free(node.mamba_value) if node.mamba_lock_ref > 0: self.mamba_protected_size_ -= mamba_num node.mamba_lock_ref = 0 else: self.mamba_evictable_size_ -= mamba_num if self.mamba_lru_list.in_list(node): self.mamba_lru_list.remove_node(node) node.mamba_value = None return mamba_num def _evict_to_host(self, node: TreeNode) -> Tuple[int, int]: # GPU -> CPU demotion: node stays in tree as evicted+backuped assert not node.evicted, f"already evicted, {node.id=}" assert node.backuped, f"not backuped, {node.id=}" num_full = len(node.value) self._record_remove_event(node, medium=StorageMedium.GPU) 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) mamba_num = self._free_device_mamba(node) node.value = None self._update_leaf_status(node) self._update_full_device_leaf_status(node.parent) return num_full, mamba_num def _evict_regular(self, node: TreeNode) -> Tuple[int, int]: # evict a non-backuped device leaf — free GPU KV + mamba, delete from tree assert not node.evicted, f"already evicted, {node.id=}" assert not node.backuped, f"backuped node, {node.id=}" assert len(node.children) == 0, f"non-leaf, {node.id=}" full_num_evicted = len(node.value) self._record_remove_event(node, medium=StorageMedium.GPU) self.cache_controller.evict_device(node.value) self.full_evictable_size_ -= full_num_evicted if self.full_lru_list.in_list(node): self.full_lru_list.remove_node(node) mamba_num_evicted = self._free_device_mamba(node) if node.mamba_host_value is not None: if self.mamba_host_lru_list.in_list(node): self.mamba_host_lru_list.remove_node(node) self.mamba_pool_host.free(node.mamba_host_value) node.mamba_host_value = None node.value = None self._discard_from_leaf_sets(node) parent = node.parent key = node.key.child_key(self.page_size) v = parent.children.pop(key, None) assert v == node, f"parent does not have child key, {key}" self._update_leaf_status(parent) _, cascade_full_num_evicted, cascade_mamba_num_evicted = ( self._iteratively_delete_tombstone_leaf(node) ) return ( full_num_evicted + cascade_full_num_evicted, mamba_num_evicted + cascade_mamba_num_evicted, ) def _evict_host_leaf(self, node: TreeNode) -> int: # evict a host-resident leaf: free host KV + mamba, delete from tree, cascade assert node.evicted, f"not evicted, {node.id=}" assert node.backuped, f"not backuped, {node.id=}" assert node.mamba_value is None, f"has device mamba, {node.id=}" assert ( node.host_ref_counter == 0 ), f"host kv in use, {node.id=} {node.host_ref_counter=}" assert ( node.host_mamba_ref_counter == 0 ), f"host mamba in use, {node.id=} {node.host_mamba_ref_counter=}" self._record_remove_event(node, medium=StorageMedium.CPU) full_num_evicted = self.cache_controller.evict_host(node.host_value) node.host_value = None if node.mamba_host_value is not None: if self.mamba_host_lru_list.in_list(node): self.mamba_host_lru_list.remove_node(node) self.mamba_pool_host.free(node.mamba_host_value) node.mamba_host_value = None self._discard_from_leaf_sets(node) parent = node.parent key = node.key.child_key(self.page_size) v = parent.children.pop(key, None) assert v == node, f"parent does not have child key, {key}" self._update_leaf_status(parent) _, cascade_full_num_evicted, _ = self._iteratively_delete_tombstone_leaf(node) return full_num_evicted + cascade_full_num_evicted def _delete_tombstone_leaf(self, node: TreeNode) -> None: assert node.mamba_value is None, f"has mamba value, {node.id=}" assert node.mamba_host_value is None, f"has mamba host value, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}" parent = node.parent key = node.key.child_key(self.page_size) 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 and node.host_mamba_ref_counter == 0 ): self._record_remove_event(node, medium=StorageMedium.CPU) 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]: 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.mamba_host_value is not None: break if node.parent.full_lock_ref > 0 or node.parent.mamba_lock_ref > 0: break if ( node.parent.host_ref_counter > 0 or node.parent.host_mamba_ref_counter > 0 ): break parent = node.parent if not parent.evicted: self._record_remove_event(parent, medium=StorageMedium.GPU) 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_device_leaf(self, x: TreeNode) -> Tuple[int, int]: """Evict a device leaf node, choosing the right strategy: - backuped: demote to host via _evict_to_host (node stays in tree) - not backuped + write_back: write_backup first, then demote - not backuped + write_through: _evict_regular (delete from tree) """ 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) return self._evict_to_host(x) else: return self._evict_regular(x) return self._evict_to_host(x) 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, evicted_mamba = self._evict_device_leaf(x) full_num_evicted += evicted_full mamba_num_evicted += evicted_mamba 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): """Evict host-resident leaf nodes: free host KV + mamba, delete from tree, cascade.""" 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._evict_host_leaf(x) if x.parent in self.evictable_full_host_leaves: heapq.heappush(heap, (x.parent.last_access_time, x.parent)) def evict_mamba_host(self, num_mamba_hosts: int) -> int: """Evict host mamba states. Internal host node: free host mamba only (tombstone). Host leaf node: same as Full host evict — _evict_host_leaf_node frees host KV + mamba, deletes from tree, cascades. """ if self.disable or num_mamba_hosts <= 0: return 0 x = self.mamba_host_lru_list.get_lru_no_lock() num_evicted = 0 while num_evicted < num_mamba_hosts and self.mamba_host_lru_list.in_list(x): x_next = self.mamba_host_lru_list.get_prev_no_lock(x) if x in self.evictable_full_host_leaves: # Leaf: evictable_full_host_leaves guarantees both counters == 0 assert ( x.host_ref_counter == 0 ), f"evict host leaf: host_ref_counter != 0 with {x.id=} {x.host_ref_counter=}" assert ( x.host_mamba_ref_counter == 0 ), f"evict host leaf: host_mamba_ref_counter != 0 with {x.id=} {x.host_mamba_ref_counter=}" self._evict_host_leaf(x) num_evicted += 1 else: # Internal host node assert ( x.host_mamba_ref_counter == 0 ), f"evict host mamba internal: host_mamba_ref_counter != 0 with {x.id=} {x.host_mamba_ref_counter=}" self.mamba_host_lru_list.remove_node(x) self.mamba_pool_host.free(x.mamba_host_value) x.mamba_host_value = None num_evicted += 1 x = x_next return num_evicted def evict_mamba(self, mamba_num: int) -> int: """Evict mamba states. Internal node: tombstone — free GPU mamba only, KV stays on GPU. Leaf node: same as Full evict — _evict_to_host moves KV+mamba to host, node stays in tree, then cascade tombstone parent device leaves. """ 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=}" assert ( not x.evicted ), f"evicted node should not be in mamba_lru_list, {x.id=}" if len(x.children) > 0: # Internal: free device mamba only, KV stays on device (tombstone) x_next = self.mamba_lru_list.get_prev_no_lock(x) mamba_num_evicted += len(x.mamba_value) self.req_to_token_pool.mamba_allocator.free(x.mamba_value) self.mamba_lru_list.remove_node(x) self._tombstone_internal_node(x) else: # Leaf: evict KV + mamba atomically assert ( x.full_lock_ref == 0 ), f"evict device leaf: full_lock_ref mismatch with {x.id=} {x.full_lock_ref=} {x.mamba_lock_ref=}" x_next = self.mamba_lru_list.get_prev_no_lock(x) _, mamba_evicted = self._evict_device_leaf(x) mamba_num_evicted += mamba_evicted 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): assert node.evicted, f"not evicted, {node.id=}" assert node.mamba_value is None, f"evicted node has device mamba, {node.id=}" n = len(fresh_value) node.value = fresh_value.clone() self.full_lru_list.insert_mru(node) self.full_evictable_size_ += n self._record_store_event(node, medium=StorageMedium.GPU) 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, prev_prefix_len: int = 0, ) -> 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 = key.child_key(self.page_size) 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 = node.key.match(key, page_size=self.page_size) if prefix_len < len(node.key): new_node = self._split_node(node.key, node, prefix_len) node = new_node if node.evicted: self._unevict_node(node, value[:prefix_len]) else: if prev_prefix_len < total_prefix_length + prefix_len: start = max(0, prev_prefix_len - total_prefix_length) self.token_to_kv_pool_allocator.free(value[start:prefix_len]) total_prefix_length += prefix_len self._inc_hit_count(node, chunked) key = key[prefix_len:] value = value[prefix_len:] if len(key): child_key = key.child_key(self.page_size) 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 = key.child_key(self.page_size) 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 or self.enable_kv_cache_events: new_node.hash_value = compute_node_hash_values(new_node, self.page_size) self._record_store_event(new_node, medium=StorageMedium.GPU) 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, best_match_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 = self._match_prefix_helper(key) return self._match_post_processor(params, value, best_last_node, best_value_len) def _match_prefix_helper( self, key: RadixKey ) -> Tuple[List[torch.Tensor], TreeNode, int]: """Walk tree to find best_last_node (mamba boundary).""" node = self.root_node child_key = key.child_key(self.page_size) 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 or node.mamba_backuped: best_value_len = len(value) best_last_node = node prefix_len = child.key.match(key, page_size=self.page_size) 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 = key.child_key(self.page_size) if node.mamba_value is not None or node.mamba_backuped: best_value_len = len(value) best_last_node = node return value, best_last_node, best_value_len def _match_post_processor( self, params: MatchPrefixParams, value: List[torch.Tensor], best_last_node: TreeNode, best_value_len: int, ) -> 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.runtime_context import get_server_args mamba_cache_chunk_size = get_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 kv_host_hit_length = 0 last_device_node = best_last_node while last_device_node is not self.root_node and last_device_node.evicted: kv_host_hit_length += len(last_device_node.host_value) last_device_node = last_device_node.parent last_host_node = best_last_node while last_host_node is not self.root_node and not last_host_node.backuped: last_host_node = last_host_node.parent mamba_host_hit = ( 1 if (last_host_node.mamba_evicted and last_host_node.mamba_backuped) else 0 ) 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._alloc_with_evict( self.req_to_token_pool.mamba_allocator, 1, self.evict_mamba, lock_node=mamba_node, error_message="Can not alloc mamba cache", ) req.mamba_pool_idx = dst_index[0] req.mamba_cow_src_index = mamba_node.mamba_value req.mamba_needs_clear = False 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, # TODO(ispobock): use best_match_node as start node for load_back best_match_node=last_host_node, host_hit_length=kv_host_hit_length, mamba_host_hit_length=mamba_host_hit, 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() 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 = {key[split_len:].child_key(self.page_size): 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[key.child_key(self.page_size)] = 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) -> IncLockRefResult: if self.disable: return IncLockRefResult(delta=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 IncLockRefResult(delta=delta) def dec_lock_ref( self, node: TreeNode, params: Optional[DecLockRefParams] = None ) -> DecLockRefResult: if self.disable: return DecLockRefResult(delta=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 DecLockRefResult(delta=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_config: PrefetchTimeoutConfig, hicache_storage_pass_prefix_keys: bool, enable_storage: bool, enable_storage_metrics: bool, extra_metric_labels: Optional[Dict[str, str]], ) -> None: 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) from sglang.srt.runtime_context import get_server_args storage_cls = resolve_collector_class( get_server_args(), STAT_LOGGER_ROLE_STORAGE, StorageMetricsCollector, ) storage_metrics_collector = storage_cls(labels=labels) self.enable_storage = enable_storage self.prefetch_threshold = prefetch_threshold self.prefetch_timeout_config = prefetch_timeout_config 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_config, 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, host_pools=self.host_pool_group.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_storage_runtime_config( storage_backend=storage_backend, prefetch_threshold=prefetch_threshold, prefetch_timeout_config=prefetch_timeout_config, 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 prefetch_abort(self, pool_transfers: Optional[list[PoolTransfer]]) -> None: """Free any allocated mamba host slots on prefetch abort/revoke.""" for transfer in pool_transfers or []: if transfer.name == PoolName.MAMBA: if transfer.host_indices is not None: self.mamba_pool_host.free(transfer.host_indices) break 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.prefetch_abort(getattr(_operation, "pool_transfers", None)) except Exception: logger.exception( "Failed to release mamba 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, entry in list(self.ongoing_backup.items()): try: node, mamba_host_protected = entry self._release_host_node(node, release_mamba=mamba_host_protected) 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, _, operation = info self.prefetch_abort(operation.pool_transfers) 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: node, mamba_host_protected = entry self._release_host_node(node, release_mamba=mamba_host_protected) 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 defaults = PrefetchTimeoutConfig() prefetch_threshold = extra_config.pop("prefetch_threshold", 256) prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", defaults.base) prefetch_timeout_per_ki_token = extra_config.pop( "prefetch_timeout_per_ki_token", defaults.per_ki_token ) prefetch_timeout_max = extra_config.pop("prefetch_timeout_max", defaults.max) 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(prefetch_timeout_max, (int, float)): raise ValueError( f"prefetch_timeout_max must be number, got " f"{type(prefetch_timeout_max).__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__}" ) prefetch_timeout_config = PrefetchTimeoutConfig( base=float(prefetch_timeout_base), per_ki_token=float(prefetch_timeout_per_ki_token), max=float(prefetch_timeout_max), ) return ( extra_config, prefetch_threshold, prefetch_timeout_config, 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): cfg = self.prefetch_timeout_config num_tokens = len(operation.hash_value) * self.page_size timeout = min(cfg.max, cfg.base + cfg.per_ki_token * num_tokens / 1024) return time.monotonic() - operation.start_time > timeout 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 if completed and operation.pool_transfers and not operation.pool_transfers_done: can_terminate = False 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 ) extra_pools = self.mamba_archive_transfers(node) operation_id = self.cache_controller.write_storage( node.host_value, node.key, node.hash_value, prefix_keys, extra_pools=extra_pools, ) mamba_host_protected = extra_pools is not None self.ongoing_backup[operation_id] = (node, mamba_host_protected) self._protect_host_node(node, protect_mamba=mamba_host_protected) 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, protect_mamba=False) # Allocate host KV memory host_indices = self._alloc_with_evict( self.cache_controller.mem_pool_host, prefetch_length, self.evict_host, ) if host_indices is None: # truncate the prefetch length to the page-aligned available host size available_size = self.cache_controller.mem_pool_host.available_size() prefetch_length = available_size - (available_size % self.page_size) if prefetch_length < self.prefetch_threshold: self._release_host_node(last_host_node, release_mamba=False) return new_input_tokens = new_input_tokens[: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, release_mamba=False) return # Allocate host mamba slot extra_pools = self.mamba_prefetch_alloc(new_input_tokens, last_hash) if extra_pools is None: self.cache_controller.mem_pool_host.free(host_indices) self._release_host_node(last_host_node, release_mamba=False) return # mamba is also being loaded, protect host mamba as well last_host_node.protect_host_mamba() if self.mamba_host_lru_list.in_list(last_host_node): self.mamba_host_lru_list.remove_node(last_host_node) operation = self.cache_controller.prefetch( req_id, host_indices, new_input_tokens, last_hash, prefix_keys, extra_pools=extra_pools, ) 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() mamba_host_indices = None mamba_loaded = False for transfer in operation.pool_transfers or []: if transfer.name == PoolName.MAMBA: mamba_host_indices = transfer.host_indices mamba_loaded = ( operation.pool_storage_result.extra_pool_hit_pages.get( PoolName.MAMBA, 0 ) >= 1 ) break 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], mamba_host_indices, mamba_loaded, ) # Free host KV memory: matched portion is already in tree, tail was unused 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] ) # Free mamba host slot if it wasn't inserted into the tree if mamba_host_indices is not None: inserted_new = matched_length < min_completed_tokens if not inserted_new or not mamba_loaded: self.mamba_pool_host.free(mamba_host_indices) 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) if loaded_from_storage > 0 and operation.pool_transfers: logger.debug( "HiCache mamba prefetch completed for request %s: prefetched_tokens=%s mamba_states=%s", req_id, loaded_from_storage, int(mamba_loaded), ) return True def _insert_helper_host( self, node: TreeNode, key: RadixKey, host_value, hash_value, mamba_host_value: Optional[torch.Tensor] = None, mamba_loaded: bool = False, ): node.last_access_time = get_last_access_time() if len(key) == 0: return 0 child_key = key.child_key(self.page_size) matched_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 node != self.root_node and node.mamba_value is not None: self.mamba_lru_list.reset_node_mru(node) prefix_len = node.key.match(key, page_size=self.page_size) key = key[prefix_len:] host_value = host_value[prefix_len:] hash_value = hash_value[prefix_len // self.page_size :] matched_length += prefix_len if prefix_len < len(node.key): new_node = self._split_node(node.key, node, prefix_len) node = new_node if len(key): child_key = key.child_key(self.page_size) leaf_node: Optional[TreeNode] = None 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 leaf_node = new_node self._update_full_host_leaf_status(new_node) self._update_full_host_leaf_status(node) self._record_store_event(new_node, medium=StorageMedium.CPU) # Attach mamba state to the new leaf if leaf_node is not None and mamba_host_value is not None and mamba_loaded: leaf_node.mamba_host_value = mamba_host_value.clone() if not self.mamba_host_lru_list.in_list(leaf_node): self.mamba_host_lru_list.insert_mru(leaf_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.prefetch_abort(operation.pool_transfers) self.cache_controller.prefetch_tokens_occupied -= len(token_ids) def _flush_pending_storage_backups_before_reset(self) -> None: if not self.enable_storage: return self.writing_check(write_back=True) deadline = time.monotonic() + 30.0 while time.monotonic() < deadline: self.drain_storage_control_queues() backup_qsize = self.cache_controller.backup_queue.qsize() ack_backup_qsize = self.cache_controller.ack_backup_queue.qsize() ongoing_backup = len(self.ongoing_backup) ongoing_write = len(self.ongoing_write_through) if ( backup_qsize == 0 and ack_backup_qsize == 0 and ongoing_backup == 0 and ongoing_write == 0 ): return time.sleep(0.05) logger.warning( "Timed out waiting for HiCache storage backups to drain before reset: " "ongoing_write=%s ongoing_backup=%s backup_queue=%s ack_backup_queue=%s", len(self.ongoing_write_through), len(self.ongoing_backup), self.cache_controller.backup_queue.qsize(), self.cache_controller.ack_backup_queue.qsize(), ) def _alloc_with_evict( self, pool, size: int, evict_fn, lock_node: Optional[TreeNode] = None, error_message: Optional[str] = None, ) -> Optional[torch.Tensor]: indices = pool.alloc(size) if indices is None: if lock_node is not None: self.inc_lock_ref(lock_node) evict_fn(size) indices = pool.alloc(size) if lock_node is not None: self.dec_lock_ref(lock_node) if indices is None and error_message is not None: raise RuntimeError(error_message) return indices # -- mamba PoolTransfer builders (D↔H↔S) ---------------------------------- def mamba_backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: # build D→H transfer descriptor for mamba state if node.mamba_value is None: return None return [ PoolTransfer( name=PoolName.MAMBA, host_indices=node.mamba_host_value, device_indices=node.mamba_value, ) ] def mamba_backup_commit( self, node: TreeNode, transfers: list[PoolTransfer] ) -> None: # store auto-allocated mamba host indices into the node after D→H backup if not transfers: return host_indices = transfers[0].host_indices if node.mamba_host_value is None and host_indices is not None: node.mamba_host_value = host_indices.clone() self.mamba_host_lru_list.insert_mru(node) def mamba_archive_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: # build H→Storage transfer descriptor for mamba state if node.mamba_host_value is None or not node.hash_value: return None return [ PoolTransfer( name=PoolName.MAMBA, host_indices=node.mamba_host_value, keys=[node.hash_value[-1]], hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] def mamba_prefetch_alloc( self, token_ids: List[int], last_hash: Optional[str], ) -> Optional[list[PoolTransfer]]: # allocate a mamba host slot and build Storage→H transfer descriptor if not token_ids: return None host_indices = self._alloc_with_evict( self.mamba_pool_host, 1, self.evict_mamba_host ) if host_indices is None: return None # placeholder key; I/O thread replaces with correct hash after hit query return [ PoolTransfer( name=PoolName.MAMBA, host_indices=host_indices, keys=["__placeholder__"], hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] def mamba_restore_transfers( self, last_hit_node: TreeNode, nodes_to_restore: list[TreeNode], req, ) -> Optional[list[PoolTransfer]]: # build H→D transfer descriptors for mamba state backed_up_host_indices: list[torch.Tensor] = [] for node in nodes_to_restore: if not node.mamba_backuped: continue backed_up_host_indices.append(node.mamba_host_value) transfers: list[PoolTransfer] = [] if backed_up_host_indices: transfers.append( PoolTransfer( name=PoolName.MAMBA, host_indices=torch.cat(backed_up_host_indices), device_indices=None, ) ) if ( req is not None and last_hit_node in nodes_to_restore and last_hit_node.mamba_host_value is not None ): if req.mamba_pool_idx is None: req.mamba_pool_idx = self._alloc_with_evict( self.req_to_token_pool.mamba_allocator, len(last_hit_node.mamba_host_value), self.evict_mamba, lock_node=last_hit_node, error_message="Cannot alloc request mamba cache for host load back", )[0] transfers.append( PoolTransfer( name=PoolName.MAMBA, host_indices=last_hit_node.mamba_host_value, device_indices=req.mamba_pool_idx.unsqueeze(0), ) ) return transfers if transfers else None def mamba_restore_commit( self, restored_nodes: list[TreeNode], transfers: Optional[list[PoolTransfer]], ) -> None: # write back controller-allocated device indices after H→D restore if not restored_nodes or not transfers or transfers[0].device_indices is None: return device_indices = transfers[0].device_indices offset = 0 for node in restored_nodes: count = len(node.mamba_host_value) node.mamba_value = device_indices[offset : offset + count].clone() offset += count