From 0d2651c3d8e04a58ca7f24c5344bb8ade793f5d8 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Thu, 12 Mar 2026 21:31:41 +0800 Subject: [PATCH 01/39] Init support for himamba tree offloading Co-authored-by: pansicheng --- benchmark/hicache/bench_multiturn.py | 11 + python/sglang/srt/managers/schedule_policy.py | 4 +- .../srt/mem_cache/hi_mamba_radix_cache.py | 625 ++++++++++++++++-- .../sglang/srt/mem_cache/hicache_storage.py | 306 ++++++++- .../hybrid_cache/hybrid_cache_controller.py | 490 ++++++++++++++ .../sglang/srt/mem_cache/mamba_radix_cache.py | 3 + python/sglang/srt/mem_cache/memory_pool.py | 36 + .../sglang/srt/mem_cache/memory_pool_host.py | 369 ++++++++++- .../model_runner_kv_cache_mixin.py | 24 +- .../4-gpu-models/test_qwen3_next_models.py | 104 ++- 10 files changed, 1912 insertions(+), 60 deletions(-) create mode 100644 python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py diff --git a/benchmark/hicache/bench_multiturn.py b/benchmark/hicache/bench_multiturn.py index d821bbc7b78c..49483fe849db 100644 --- a/benchmark/hicache/bench_multiturn.py +++ b/benchmark/hicache/bench_multiturn.py @@ -529,6 +529,8 @@ def response_handler(self): f"({expected} clients), releasing {len(next_round_reqs)} " f"requests for round {current_barrier_round + 1}" ) + self._send_heartbeat(input_len=100, output_len=100) + time.sleep(10) for req in next_round_reqs: self.ready_queue.append(req) next_round_reqs = [] @@ -541,6 +543,15 @@ def response_handler(self): print(f"Error processing response for client {client_id}: {e}") continue + def _send_heartbeat(self, input_len=100, output_len=20): + """Send a small heartbeat request to the server.""" + heartbeat_input = [1] * input_len + payload = gen_payload(heartbeat_input, output_len, self.lora_path) + try: + requests.post(self.url, json=payload, timeout=30) + except Exception as e: + print(f"Heartbeat request failed: {e}") + def run(self): request_thread = threading.Thread(target=self.request_sender, daemon=True) response_thread = threading.Thread(target=self.response_handler, daemon=True) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 7ebd5cd7bb38..683dd77438f9 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -767,9 +767,9 @@ def add_one_req( if total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN - if req.host_hit_length > 0: + if req.host_hit_length > 0 or req.last_host_node is not req.last_node: new_indices, req.last_node = self.tree_cache.init_load_back( - req.last_host_node, req.host_hit_length + req.last_host_node, req.host_hit_length, req=req ) req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 235dc8c900fd..fc015742072a 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -12,20 +12,31 @@ 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.hicache_storage import PoolTransfer, PoolTransferResult +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + PrefetchOperation, +) 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 -from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool +from sglang.srt.mem_cache.memory_pool_host import ( + HostPoolGroup, + MambaPoolHost, + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + PoolEntry, +) from sglang.srt.mem_cache.radix_cache import ( RadixKey, compute_node_hash_values, @@ -41,6 +52,41 @@ 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. @@ -62,13 +108,33 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): 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.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.full_kv_pool_host = MHATokenToKVPoolHost( + self.hybrid_model_layer_ids: list[int] = [] + full_layer_ids = sorted( + self.hybrid_kv_cache.full_attention_layer_id_mapping.keys() + ) + mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) + self.hybrid_model_layer_ids = sorted(set(full_layer_ids) | set(mamba_layer_ids)) + self.transfer_layer_num = len(self.hybrid_model_layer_ids) + self.hybrid_kv_cache.set_model_layer_id_mapping(self.hybrid_model_layer_ids) + params.req_to_token_pool.set_model_layer_id_mapping(self.hybrid_model_layer_ids) + + self.kvcache = self.hybrid_kv_cache.full_kv_pool + kv_host_pool_cls = ( + MLATokenToKVPoolHost + if self.hybrid_kv_cache.use_mla + else MHATokenToKVPoolHost + ) + self.full_kv_pool_host = kv_host_pool_cls( self.kvcache, server_args.hicache_ratio, server_args.hicache_size, @@ -76,6 +142,49 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) + self.mamba_pool_host = MambaPoolHost( + params.req_to_token_pool.mamba_pool, + server_args.hicache_ratio, + server_args.hicache_size, + allocator_type=server_args.hicache_storage_backend, + ) + + full_layer_mapping = dict(self.hybrid_kv_cache.full_attention_layer_id_mapping) + mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) + + def kv_layer_mapper(model_layer_local_id: int) -> Optional[int]: + if not 0 <= model_layer_local_id < len(self.hybrid_model_layer_ids): + return None + return full_layer_mapping.get( + self.hybrid_model_layer_ids[model_layer_local_id] + ) + + def mamba_layer_mapper(model_layer_local_id: int) -> Optional[int]: + if not 0 <= model_layer_local_id < len(self.hybrid_model_layer_ids): + return None + return mamba_layer_mapping.get( + self.hybrid_model_layer_ids[model_layer_local_id] + ) + + self.host_pool_group = HostPoolGroup( + [ + PoolEntry( + name="kv", + host_pool=self.full_kv_pool_host, + device_pool=self.kvcache, + layer_mapper=kv_layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name="mamba", + host_pool=self.mamba_pool_host, + device_pool=params.req_to_token_pool.mamba_pool, + layer_mapper=mamba_layer_mapper, + host_evict_fn=self.evict_mamba_host, + device_evict_fn=self.evict_mamba, + ), + ] + ) self.tp_group = params.tp_cache_group self.tp_world_size = ( @@ -101,9 +210,9 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy self.load_cache_event = threading.Event() - self.cache_controller = HiCacheController( + self.cache_controller = HybridCacheController( params.token_to_kv_pool_allocator, - self.full_kv_pool_host, + self.host_pool_group, params.page_size, self.tp_group, load_cache_event=self.load_cache_event, @@ -115,6 +224,10 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): storage_backend_extra_config=extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + transfer_layer_num=self.transfer_layer_num, + ) + params.req_to_token_pool.register_layer_transfer_counter( + self.cache_controller.layer_done_counter ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, @@ -142,6 +255,7 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): 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) @@ -150,8 +264,10 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): 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 = {} @@ -159,33 +275,60 @@ def reset(self) -> None: 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): + # 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.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 + if extra_pools is not None: + self.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) + if extra_pools is not None: + logger.info( + "HiCache mamba offload prepared for node %s: kv_tokens=%s mamba_states=%s", + node.id, + len(node.host_value), + ( + len(node.mamba_host_value) + if node.mamba_host_value is not None + else 0 + ), + ) else: return 0 return len(host_indices) def load_back( - self, node: TreeNode, mem_quota: Optional[int] = None + self, node: TreeNode, mem_quota: Optional[int] = None, req=None ) -> Optional[torch.Tensor]: """Load full KV back from host.""" last_hit_node = node @@ -198,33 +341,58 @@ def load_back( else: ancestor_node = node + mamba_restore_nodes = list(nodes_to_load) + if ( + last_hit_node not in mamba_restore_nodes + and last_hit_node.mamba_host_value is not None + and last_hit_node.mamba_value is None + ): + mamba_restore_nodes.append(last_hit_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 + 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 + ) ): # skip loading back if the total size is too small or exceeding the memory quota self.dec_lock_ref(ancestor_node) return None + mamba_pools, mamba_nodes = self.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: self.evict(EvictParams(num_tokens=len(full_host_indices))) + mamba_pools, mamba_nodes = self.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.restore_commit(mamba_nodes, mamba_pools) + offset = 0 for n in nodes_to_load: n_len = len(n.host_value) @@ -255,16 +423,25 @@ def init_load_back( last_node: TreeNode, host_hit_length: int, mem_quota: Optional[int] = None, + req=None, ): - if last_node.evicted: - loading_values = self.load_back(last_node, mem_quota) + if last_node.evicted or ( + last_node.mamba_value is None and last_node.mamba_host_value is not None + ): + 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.evicted: + while last_node is not self.root_node and ( + last_node.evicted + or ( + last_node.mamba_value is None + and last_node.mamba_host_value is not None + ) + ): last_node = last_node.parent return ( @@ -414,6 +591,7 @@ def _evict_to_host(self, node: TreeNode) -> int: 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 node.mamba_host_value is None, f"node has mamba host 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) @@ -442,6 +620,8 @@ def _iteratively_delete_tombstone_leaf( 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 @@ -507,6 +687,11 @@ def evict_host(self, num_tokens: int): num_evicted += self.cache_controller.evict_host(x.host_value) x.host_value = None + if x.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(x): + self.mamba_host_lru_list.remove_node(x) + self.mamba_pool_host.free(x.mamba_host_value) + x.mamba_host_value = None self.evictable_full_host_leaves.discard(x) self._update_full_host_leaf_status(x.parent) @@ -534,6 +719,12 @@ def evict_host(self, num_tokens: int): self.req_to_token_pool.mamba_pool.free(x.mamba_value) x.mamba_value = None + if x.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(x): + self.mamba_host_lru_list.remove_node(x) + self.mamba_pool_host.free(x.mamba_host_value) + x.mamba_host_value = None + self.evictable_full_host_leaves.discard(x) parent = x.parent @@ -545,6 +736,22 @@ def evict_host(self, num_tokens: int): if parent in self.evictable_full_host_leaves: heapq.heappush(heap, (parent.last_access_time, parent)) + def evict_mamba_host(self, num_mamba_hosts: int) -> int: + 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.host_ref_counter == 0: + 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: if self.disable or mamba_num <= 0: return 0 @@ -557,6 +764,8 @@ def evict_mamba(self, mamba_num: int) -> int: assert x.mamba_lock_ref == 0, f"node is in use, {x.id=}" if x.evicted: + if x.mamba_host_value is None: + break 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) @@ -564,11 +773,13 @@ def evict_mamba(self, mamba_num: int) -> int: self.mamba_evictable_size_ -= len(x.mamba_value) x.mamba_value = None - if len(x.children) == 0: + if len(x.children) == 0 and x.mamba_host_value is None: self._delete_tombstone_leaf(x) _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) mamba_num_evicted += cascade_mamba elif len(x.children) > 0: + if x.mamba_host_value is None: + break 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) @@ -597,22 +808,27 @@ def evict_mamba(self, mamba_num: int) -> int: self.mamba_evictable_size_ -= len(x.mamba_value) if x.backuped: - self.cache_controller.evict_host(x.host_value) - x.host_value = None + pass + elif x.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(x): + self.mamba_host_lru_list.remove_node(x) + self.mamba_pool_host.free(x.mamba_host_value) + x.mamba_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=}" + if not x.backuped: + 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 + 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() @@ -769,6 +985,10 @@ def _match_prefix_helper( value: List[torch.Tensor] = [] best_value_len = 0 best_last_node = node + mamba_available_size = ( + self.req_to_token_pool.mamba_pool.available_size() + + self.mamba_evictable_size() + ) while len(key) > 0 and child_key in node.children.keys(): child = node.children[child_key] @@ -776,7 +996,9 @@ def _match_prefix_helper( if child.evicted and not child.backuped: break - if node.mamba_value is not None: + if node.mamba_value is not None or ( + node.mamba_host_value is not None and mamba_available_size > 0 + ): best_value_len = len(value) best_last_node = node @@ -795,7 +1017,9 @@ def _match_prefix_helper( if len(key): child_key = self.get_child_key_fn(key) - if node.mamba_value is not None: + if node.mamba_value is not None or ( + node.mamba_host_value is not None and mamba_available_size > 0 + ): best_value_len = len(value) best_last_node = node @@ -842,11 +1066,20 @@ def _match_post_processor( else: mamba_branching_seqlen = None - # last_device_node & host_hit_length: from best_last_node (mamba boundary) + # last_device_node & host_hit_length: + # last_device_node is the last node whose full and mamba states can both + # be directly used on device. host_hit_length remains full-KV-only. host_hit_length = 0 last_device_node = best_last_node - while last_device_node.evicted: - host_hit_length += len(last_device_node.host_value) + while last_device_node is not self.root_node and ( + last_device_node.evicted + or ( + last_device_node.mamba_value is None + and last_device_node.mamba_host_value is not None + ) + ): + if 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 @@ -862,13 +1095,13 @@ def _match_post_processor( 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" + dst_index = self._alloc_with_evict( + self.req_to_token_pool.mamba_pool, + 1, + self.evict_mamba, + lock_node=mamba_node, + error_message="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] @@ -1174,6 +1407,7 @@ def attach_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( @@ -1210,6 +1444,14 @@ def detach_storage_backend(self) -> tuple: 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 == "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 @@ -1227,6 +1469,12 @@ def _force_release_pending_storage_ops(self): 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: @@ -1286,7 +1534,8 @@ 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 + 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: @@ -1489,12 +1738,13 @@ def write_backup_storage(self, node: TreeNode): if self.hicache_storage_pass_prefix_keys else None ) - + extra_pools = self.archive_transfers(node) operation_id = self.cache_controller.write_storage( node.host_value, node.key, node.hash_value, prefix_keys, + extra_pools=extra_pools, ) self.ongoing_backup[operation_id] = node self._protect_host_node(node) @@ -1519,15 +1769,36 @@ def prefetch_from_storage( 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) + host_indices = self._alloc_with_evict( + self.cache_controller.mem_pool_host, + prefetch_length, + self.evict_host, + ) if host_indices is None: self._release_host_node(last_host_node) return + extra_pools = None + prepared = self.prefetch_prepare(new_input_tokens, last_hash) + if prefetch_length > 0 and prepared is None: + self.cache_controller.mem_pool_host.free(host_indices) + self._release_host_node(last_host_node) + return + if prepared is not None: + _, extra_pools = prepared + if extra_pools is not None: + logger.info( + "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", + req_id, + prefetch_length // self.page_size, + 1, + ) operation = self.cache_controller.prefetch( - req_id, host_indices, new_input_tokens, last_hash, prefix_keys + req_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys, + extra_pools=extra_pools, ) self.ongoing_prefetch[req_id] = ( last_host_node, @@ -1568,6 +1839,11 @@ def check_prefetch_progress(self, req_id: str) -> bool: min_completed_tokens = completed_tokens_tensor.item() fetched_token_ids = token_ids[:min_completed_tokens] written_indices = host_indices[:min_completed_tokens] + mamba_host_indices = None + for transfer in operation.pool_transfers or []: + if transfer.name == "mamba": + mamba_host_indices = transfer.host_indices + break matched_length = self._insert_helper_host( last_host_node, RadixKey( @@ -1576,9 +1852,17 @@ def check_prefetch_progress(self, req_id: str) -> bool: ), written_indices, hash_value[: min_completed_tokens // self.page_size], + mamba_host_indices, + operation.pool_storage_result, ) self.cache_controller.mem_pool_host.free(host_indices[:matched_length]) + mamba_loaded = self.prefetch_commit( + operation.pool_transfers, + matched_length, + min_completed_tokens, + operation.pool_storage_result, + ) self.cache_controller.append_host_mem_release( host_indices[min_completed_tokens:completed_tokens] ) @@ -1591,11 +1875,24 @@ def check_prefetch_progress(self, req_id: str) -> bool: 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.info( + "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 + self, + node: TreeNode, + key: RadixKey, + host_value, + hash_value, + mamba_host_value: Optional[torch.Tensor] = None, + pool_storage_result: Optional[PoolTransferResult] = None, ): node.last_access_time = get_last_access_time() if len(key) == 0: @@ -1605,6 +1902,10 @@ def _insert_helper_host( matched_length = 0 host_value_inserted = False + final_mamba_node: Optional[TreeNode] = None + has_mamba = pool_storage_result is None or ( + pool_storage_result.extra_pool_hit_pages.get("mamba", 0) >= 1 + ) while len(key) > 0 and child_key in node.children.keys(): node = node.children[child_key] node.last_access_time = get_last_access_time() @@ -1615,6 +1916,8 @@ def _insert_helper_host( if node.evicted and not node.backuped: node.host_value = host_value[:prefix_len].clone() host_value_inserted = True + if prefix_len == len(key): + final_mamba_node = node self._update_full_host_leaf_status(node) if node.parent is not None: self._update_full_host_leaf_status(node.parent) @@ -1645,8 +1948,13 @@ def _insert_helper_host( new_node.host_value = host_value.clone() new_node.hash_value = hash_value node.children[child_key] = new_node + final_mamba_node = new_node self._update_full_host_leaf_status(new_node) self._update_full_host_leaf_status(node) + if final_mamba_node is not None and mamba_host_value is not None and has_mamba: + final_mamba_node.mamba_host_value = mamba_host_value.clone() + if not self.mamba_host_lru_list.in_list(final_mamba_node): + self.mamba_host_lru_list.insert_mru(final_mamba_node) return matched_length def release_aborted_request(self, rid: str): @@ -1665,4 +1973,221 @@ def release_aborted_request(self, rid: str): 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 + last_log_time = 0.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 + + def _last_page_index(self, token_count: int) -> Optional[int]: + if token_count <= 0: + return None + return token_count // self.page_size - 1 + + def backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: + """PoolTransfers for D→H backup: mamba device → mamba host.""" + if node.mamba_value is None: + return None + return [ + PoolTransfer( + name="mamba", + host_indices=node.mamba_host_value, + device_indices=node.mamba_value, + ) + ] + + def backup_commit(self, node: TreeNode, transfers: list[PoolTransfer]) -> None: + """After D→H backup succeeds: store auto-allocated mamba host indices into node.""" + if not transfers: + return + mamba_host = transfers[0].host_indices + if node.mamba_host_value is None and mamba_host is not None: + node.mamba_host_value = mamba_host + self.mamba_host_lru_list.insert_mru(node) + + def archive_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: + """PoolTransfers for H→Storage archive (write_backup_storage).""" + mamba_host_value = getattr(node, "mamba_host_value", None) + if mamba_host_value is None or not node.hash_value: + return None + return [ + PoolTransfer( + name="mamba", + host_indices=mamba_host_value, + keys=[node.hash_value[-1]], + hit_policy="trailing_pages", + ) + ] + + def prefetch_prepare( + self, + token_ids: List[int], + last_hash: Optional[str], + ) -> Optional[tuple[torch.Tensor, list[PoolTransfer]]]: + """Alloc mamba host slot and build PoolTransfers for Storage→H prefetch. + + Returns (host_handle, transfers) so the caller can track the allocation, + or None if allocation fails (caller should cancel the prefetch). + """ + if not token_ids: + return None + mamba_host_index = self._alloc_with_evict( + self.mamba_pool_host, 1, self.evict_mamba_host + ) + if mamba_host_index is None: + return None + last_page_hash = last_hash + for start in range(0, len(token_ids), self.page_size): + last_page_hash = self.cache_controller.get_hash_str( + token_ids[start : start + self.page_size], last_page_hash + ) + transfers = [ + PoolTransfer( + name="mamba", + host_indices=mamba_host_index, + keys=[last_page_hash], + hit_policy="trailing_pages", + ) + ] + return mamba_host_index, transfers + + def restore_transfers( + self, + last_hit_node: TreeNode, + nodes_to_restore: list[TreeNode], + req, + ) -> tuple[Optional[list[PoolTransfer]], list[TreeNode]]: + """PoolTransfers for H→D restore (load_back). + + Node slots use device_indices=None so the controller auto-allocates them, + mirroring KV. The req slot is pre-allocated and must land in req.mamba_pool_idx. + Returns (transfers, extra_nodes) where extra_nodes need device write-back via + restore_commit() after load succeeds. + """ + mamba_host_list: list[torch.Tensor] = [] + extra_nodes: list[TreeNode] = [] + + for node in nodes_to_restore: + if node.mamba_host_value is None or node.mamba_value is not None: + continue + extra_nodes.append(node) + mamba_host_list.append(node.mamba_host_value) + + transfers: list[PoolTransfer] = [] + if mamba_host_list: + transfers.append( + PoolTransfer( + name="mamba", + host_indices=torch.cat(mamba_host_list), + ) + ) + + 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_pool, + len(last_hit_node.mamba_host_value), + self.evict_mamba, + lock_node=last_hit_node, + error_message="Can not alloc request mamba cache for host load back", + )[0] + transfers.append( + PoolTransfer( + name="mamba", + host_indices=last_hit_node.mamba_host_value, + device_indices=req.mamba_pool_idx.unsqueeze(0), + ) + ) + + return (transfers if transfers else None), extra_nodes + + def restore_commit( + self, + extra_nodes: list[TreeNode], + transfers: Optional[list[PoolTransfer]], + ) -> None: + """After H→D restore succeeds: write back controller-allocated mamba device indices.""" + if not extra_nodes or not transfers or transfers[0].device_indices is None: + return + mamba_device = transfers[0].device_indices + offset = 0 + for n in extra_nodes: + n_len = len(n.mamba_host_value) + n.mamba_value = mamba_device[offset : offset + n_len].clone() + offset += n_len + + def prefetch_commit( + self, + pool_transfers: Optional[list[PoolTransfer]], + matched_length: int, + min_completed_tokens: int, + result: PoolTransferResult, + ) -> bool: + """After Storage→H prefetch completes: free mamba host slot if not inserted into tree. + + Returns True if a mamba state was successfully loaded into the radix tree. + """ + mamba_host_indices = None + for transfer in pool_transfers or []: + if transfer.name == "mamba": + mamba_host_indices = transfer.host_indices + break + if mamba_host_indices is None: + return False + # mamba covers the entire prefix as a single page; loaded means >= 1 page hit. + mamba_loaded = result.extra_pool_hit_pages.get("mamba", 0) >= 1 + if matched_length == min_completed_tokens or not mamba_loaded: + self.mamba_pool_host.free(mamba_host_indices) + return mamba_loaded diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 02680db04a04..e04d7938db0c 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -3,7 +3,7 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Literal, Optional, Set import torch @@ -64,6 +64,47 @@ class HiCacheStorageExtraInfo: extra_info: Optional[dict] = None +@dataclass +class PoolTransfer: + """Unified per-pool transfer descriptor for batch v2 interface. + + device<->host path : host_indices + device_indices + host<->storage path: host_indices + keys + + hit_policy controls how batch_exists_v2 checks this pool: + "all_pages" – every page in the KV hit range must exist. + "trailing_pages"– only the last len(keys) pages must exist (e.g. mamba/SWA). + """ + + name: str + host_indices: Optional[torch.Tensor] = None + device_indices: Optional[torch.Tensor] = None + keys: Optional[List[str]] = None + hit_policy: Literal["all_pages", "trailing_pages"] = "all_pages" + + +@dataclass +class PoolTransferResult: + """Tracks how many pages were successfully processed per pool.""" + + kv_hit_pages: int + extra_pool_hit_pages: dict[str, int] + + @classmethod + def empty(cls) -> "PoolTransferResult": + return cls(0, {}) + + def update_kv_hit_pages(self, kv_hit_pages: int) -> None: + """Accumulate kv_hit_pages across batches (max = last successful batch).""" + self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages) + + def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None: + """Record actual load/write success counts per extra pool.""" + self.extra_pool_hit_pages.update( + {name: sum(rs) for name, rs in results.items()} + ) + + class HiCacheStorage(ABC): """ HiCacheStorage is a class that provides a generic key-value interface for storing and retrieving KV cache. @@ -71,10 +112,49 @@ class HiCacheStorage(ABC): """ # todo, the page size of storage backend does not have to be the same as the same as host memory pool - def register_mem_pool_host(self, mem_pool_host: HostKVCache): self.mem_pool_host = mem_pool_host + def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): + if not hasattr(self, "registered_pools"): + self.registered_pools = {} + self.registered_pools[host_pool_name] = host_pool + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + """Check which KV pages (and associated extra pools) exist in storage. + + Returns a PoolTransferResult with the longest contiguous KV prefix hit and + per-pool hit counts, constrained by each pool's auxiliary_constraints policy. + """ + raise NotImplementedError() + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + """Read data from storage into host memory for each PoolTransfer. + + Returns a dict mapping pool name to a per-entry success list. + """ + raise NotImplementedError() + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + """Write data from host memory to storage for each PoolTransfer. + + Returns a dict mapping pool name to a per-entry success list. + """ + raise NotImplementedError() + def batch_get_v1( self, keys: List[str], @@ -202,7 +282,6 @@ def __init__( self.config_suffix = f"_{model_name}" else: self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}" - if not os.path.exists(self.file_path) and tp_rank == 0: os.makedirs(self.file_path) logger.info(f"Created HiCacheFile storage directory at {self.file_path}") @@ -210,6 +289,18 @@ def __init__( def _get_suffixed_key(self, key: str) -> str: return key + self.config_suffix + def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str: + if component_name is None or component_name in ("__default__", "kv"): + return self._get_suffixed_key(key) + return self._get_suffixed_key(f"{key}.{component_name}") + + def _get_component_path( + self, key: str, component_name: Optional[str] = None + ) -> str: + return os.path.join( + self.file_path, f"{self._get_component_key(key, component_name)}.bin" + ) + def get( self, key: str, @@ -279,6 +370,215 @@ def exists(self, key: str) -> bool: tensor_path = os.path.join(self.file_path, f"{key}.bin") return os.path.exists(tensor_path) + def _collect_existing_component_keys( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + ) -> Set[str]: + target_files = {f"{self._get_component_key(key)}.bin" for key in keys} + for transfer in pool_transfers or []: + for key in keys: + target_files.add(f"{self._get_component_key(key, transfer.name)}.bin") + + existing_files = set() + with os.scandir(self.file_path) as entries: + for entry in entries: + if entry.is_file() and entry.name in target_files: + existing_files.add(entry.name) + return existing_files + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + """Check which KV pages (and associated extra pools) exist in storage. + + Returns a PoolTransferResult with the longest contiguous KV prefix hit and + per-pool hit counts, constrained by each pool's auxiliary_constraints policy. + """ + existing_files = self._collect_existing_component_keys(keys, pool_transfers) + + def has_component(page_idx: int, name: str) -> bool: + return ( + f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files + ) + + # Longest contiguous KV prefix present in storage. + kv_pages = next( + ( + i + for i in range(len(keys)) + if f"{self._get_component_key(keys[i])}.bin" not in existing_files + ), + len(keys), + ) + + hit_count: dict[str, int] = {"kv": kv_pages} if kv_pages else {} + final_pages = kv_pages + + for transfer in pool_transfers or []: + if final_pages == 0: + break + name = transfer.name + if transfer.hit_policy == "all_pages": + boundary = next( + (i for i in range(kv_pages) if not has_component(i, name)), kv_pages + ) + else: # trailing_pages + trailing = max(1, len(transfer.keys) if transfer.keys else 1) + boundary = 0 + for prefix_len in range(kv_pages, 0, -1): + if all( + has_component(i, name) + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + if boundary: + hit_count[name] = boundary + final_pages = min(final_pages, boundary) + + if pool_transfers: + logger.info( + "HiCacheFile batch_exists_v2: kv_pages=%s final_pages=%s hit_count=%s first_key=%s last_key=%s", + kv_pages, + final_pages, + hit_count, + keys[0] if keys else None, + keys[final_pages - 1] if final_pages > 0 else None, + ) + + return PoolTransferResult(final_pages, hit_count) + + def _log_key(self, pool_name: str, key: str) -> str: + return key if pool_name == "kv" else f"{key}.{pool_name}" + + def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool: + """Read one page from storage into host_pool at page_offset.""" + tensor_path = self._get_component_path(key, pool_name) + if not os.path.exists(tensor_path): + return False + try: + with open(tensor_path, "rb", buffering=0) as f: + file_data = f.read() + dummy_page = host_pool.get_dummy_flat_data_page() + expected_bytes = dummy_page.numel() * dummy_page.element_size() + if len(file_data) != expected_bytes: + logger.warning( + "File size mismatch for %s: expected %s, got %s", + self._log_key(pool_name, key), + expected_bytes, + len(file_data), + ) + return False + data_page = ( + torch.frombuffer( + file_data, dtype=dummy_page.dtype, count=dummy_page.numel() + ) + .clone() + .reshape(dummy_page.shape) + ) + host_pool.set_from_flat_data_page(page_offset, data_page) + return True + except Exception as e: + logger.error("Failed to read key %s: %s", self._log_key(pool_name, key), e) + return False + + def _write_page( + self, pool_name: str, key: str, host_pool, page_offset: int + ) -> bool: + """Write one page from host_pool at page_offset to storage as raw bytes.""" + try: + data_page = host_pool.get_data_page(page_offset, flat=True) + data_bytes = ( + data_page.contiguous().view(torch.uint8).reshape(-1).numpy().tobytes() + ) + component_path = self._get_component_path(key, pool_name) + if os.path.exists(component_path): + existing_size = os.path.getsize(component_path) + if existing_size == len(data_bytes): + return True # same-size file → assume identical, skip write + logger.warning( + "Overwriting stale storage entry for %s: expected %s bytes, got %s", + self._log_key(pool_name, key), + len(data_bytes), + existing_size, + ) + with open(component_path, "wb", buffering=0) as f: + f.write(data_bytes) + if pool_name != "kv": + logger.info( + "HiCacheFile wrote auxiliary component %s for key %s", + pool_name, + key, + ) + return True + except Exception as e: + logger.error("Failed to store key %s: %s", self._log_key(pool_name, key), e) + return False + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + del extra_info + results: dict[str, List[bool]] = {} + for transfer in transfers: + host_pool = self.registered_pools[transfer.name] + page_size = getattr(host_pool, "page_size", 1) or 1 + if transfer.host_indices.numel() != len(transfer.keys) * page_size: + logger.error( + "batch_get_v2 indices length mismatch for %s: expected %s, got %s", + transfer.name, + len(transfer.keys) * page_size, + transfer.host_indices.numel(), + ) + results[transfer.name] = [False] * len(transfer.keys) + else: + results[transfer.name] = [ + self._read_page( + transfer.name, + key, + host_pool, + transfer.host_indices[i * page_size].item(), + ) + for i, key in enumerate(transfer.keys) + ] + return results + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + del extra_info + results: dict[str, List[bool]] = {} + for transfer in transfers: + host_pool = self.registered_pools[transfer.name] + page_size = getattr(host_pool, "page_size", 1) or 1 + if transfer.host_indices.numel() != len(transfer.keys) * page_size: + logger.error( + "batch_set_v2 indices length mismatch for %s: expected %s, got %s", + transfer.name, + len(transfer.keys) * page_size, + transfer.host_indices.numel(), + ) + results[transfer.name] = [False] * len(transfer.keys) + else: + results[transfer.name] = [ + self._write_page( + transfer.name, + key, + host_pool, + transfer.host_indices[i * page_size].item(), + ) + for i, key in enumerate(transfer.keys) + ] + return results + def clear(self) -> bool: try: for filename in os.listdir(self.file_path): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py new file mode 100644 index 000000000000..904728b422c9 --- /dev/null +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -0,0 +1,490 @@ +from __future__ import annotations + +import logging +import threading +import time +from typing import TYPE_CHECKING, Any, List, Optional + +import torch + +from sglang.srt.managers.cache_controller import CacheOperation as BaseCacheOperation +from sglang.srt.managers.cache_controller import ( + HiCacheAck, +) +from sglang.srt.managers.cache_controller import ( + HiCacheController as BaseHiCacheController, +) +from sglang.srt.managers.cache_controller import ( + LayerDoneCounter, +) +from sglang.srt.managers.cache_controller import ( + StorageOperation as BaseStorageOperation, +) +from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorageExtraInfo, + PoolTransfer, + PoolTransferResult, +) +from sglang.srt.mem_cache.memory_pool_host import PoolEntry +from sglang.srt.utils import get_device_module + +if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + +logger = logging.getLogger(__name__) +device_module = get_device_module() + + +class CacheOperation(BaseCacheOperation): + def __init__( + self, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + node_id: int, + priority: Optional[int] = None, + pool_transfers: Optional[list[PoolTransfer]] = None, + ): + super().__init__(host_indices, device_indices, node_id, priority) + self.pool_transfers = pool_transfers + + @staticmethod + def merge_pool_transfers( + ops: List["CacheOperation"], + ) -> Optional[list[PoolTransfer]]: + grouped: dict[str, list[PoolTransfer]] = {} + for op in ops: + for t in op.pool_transfers or []: + grouped.setdefault(t.name, []).append(t) + if not grouped: + return None + + def cat_or_none(tensors): + parts = [x for x in tensors if x is not None] + return torch.cat(parts) if parts else None + + return [ + PoolTransfer( + name=name, + host_indices=cat_or_none(t.host_indices for t in ts), + device_indices=cat_or_none(t.device_indices for t in ts), + keys=[k for t in ts if t.keys for k in t.keys] or None, + ) + for name, ts in grouped.items() + ] + + @staticmethod + def merge_ops(ops: List["CacheOperation"]) -> "CacheOperation": + if len(ops) == 1: + return ops[0] + host_indices = torch.cat([op.host_indices for op in ops]) + device_indices = torch.cat([op.device_indices for op in ops]) + node_ids = [] + priority = min(op.priority for op in ops) + for op in ops: + node_ids.extend(op.node_ids) + merged = CacheOperation( + host_indices, + device_indices, + -1, + priority, + pool_transfers=CacheOperation.merge_pool_transfers(ops), + ) + merged.node_ids = node_ids + return merged + + +class StorageOperation(BaseStorageOperation): + def __init__( + self, + host_indices: torch.Tensor, + token_ids: List[int], + last_hash: Optional[str] = None, + hash_value: Optional[List[str]] = None, + prefix_keys: Optional[List[str]] = None, + pool_transfers: Optional[list[PoolTransfer]] = None, + ): + super().__init__(host_indices, token_ids, last_hash, hash_value, prefix_keys) + self.pool_transfers = pool_transfers + self.pool_storage_result = PoolTransferResult.empty() + + +class PrefetchOperation(StorageOperation): + def __init__( + self, + request_id: str, + host_indices: torch.Tensor, + token_ids: List[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[List[str]] = None, + pool_transfers: Optional[list[PoolTransfer]] = None, + ): + self.request_id = request_id + self._lock = threading.Lock() + self._terminated_flag = False + self.start_time = time.monotonic() + super().__init__( + host_indices, + token_ids, + last_hash, + prefix_keys=prefix_keys, + pool_transfers=pool_transfers, + ) + + def increment(self, num_tokens: int): + with self._lock: + if self._terminated_flag: + return False + self.completed_tokens += num_tokens + return True + + def mark_terminate(self): + with self._lock: + self._terminated_flag = True + + def is_terminated(self) -> bool: + return self._terminated_flag + + +class HybridCacheController(BaseHiCacheController): + def __init__( + self, + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, + mem_pool_host: Any, + page_size: int, + tp_group: torch.distributed.ProcessGroup, + load_cache_event: threading.Event, + write_policy: str = "write_through_selective", + io_backend: str = "", + storage_backend: Optional[str] = None, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + pp_rank: int = 0, + pp_size: int = 1, + transfer_layer_num: Optional[int] = None, + ): + self.transfer_layer_num = transfer_layer_num + startup_storage_backend = storage_backend + super().__init__( + token_to_kv_pool_allocator=token_to_kv_pool_allocator, + mem_pool_host=mem_pool_host, + page_size=page_size, + tp_group=tp_group, + load_cache_event=load_cache_event, + write_policy=write_policy, + io_backend=io_backend, + storage_backend=None, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + pp_rank=pp_rank, + pp_size=pp_size, + ) + self.transfer_layer_num = self.transfer_layer_num or self.layer_num + if self.transfer_layer_num != self.layer_done_counter.num_layers: + self.layer_done_counter = LayerDoneCounter(self.transfer_layer_num) + self.mem_pool_device.register_layer_transfer_counter( + self.layer_done_counter + ) + if startup_storage_backend is not None: + self.attach_storage_backend( + storage_backend=startup_storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + host_pools=getattr(mem_pool_host, "entries", None), + ) + + def attach_storage_backend( + self, + storage_backend: str, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + host_pools: Optional[list[PoolEntry]] = None, + ): + super().attach_storage_backend( + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + ) + + for entry in host_pools or []: + self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name) + logger.info("Using pool-based interface for hybrid storage operations") + + def reset(self): + super().reset() + if self.enable_storage: + self.host_mem_release_queue.queue.clear() + self.prefetch_tokens_occupied = 0 + + def write( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> Optional[torch.Tensor]: + host_indices = self.mem_pool_host.alloc(len(device_indices)) + if host_indices is None: + return None + pool_transfers = self._resolve_pool_transfers_allocation( + extra_pools, alloc_host=True + ) + if pool_transfers is None and extra_pools: + self.mem_pool_host.free(host_indices) + return None + + self.write_queue.append( + CacheOperation( + host_indices, + device_indices, + node_id, + priority, + pool_transfers=pool_transfers or None, + ) + ) + self.start_writing() + return host_indices + + def start_writing(self) -> None: + if not self.write_queue: + return + op = CacheOperation.merge_ops(self.write_queue) + host_indices, device_indices = self.move_indices(op) + self.write_queue.clear() + start_event = device_module.Event() + finish_event = device_module.Event() + start_event.record() + with device_module.stream(self.write_stream): + start_event.wait(self.write_stream) + self.mem_pool_host.backup_from_device_all_layer( + self.mem_pool_device, + host_indices, + device_indices, + self.io_backend, + pool_transfers=op.pool_transfers, + ) + finish_event.record() + if host_indices.is_cuda: + host_indices.record_stream(self.write_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.write_stream) + self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + + def load( + self, + host_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> Optional[torch.Tensor]: + device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) + if device_indices is None: + return None + pool_transfers = self._resolve_pool_transfers_allocation( + extra_pools, alloc_host=False + ) + if pool_transfers is None and extra_pools: + self.mem_pool_device_allocator.free(device_indices) + return None + + self.load_queue.append( + CacheOperation( + host_indices, + device_indices, + node_id, + priority, + pool_transfers=pool_transfers or None, + ) + ) + return device_indices + + def start_loading(self) -> int: + if not self.load_queue: + return -1 + producer_id = self.layer_done_counter.update_producer() + op = CacheOperation.merge_ops(self.load_queue) + host_indices, device_indices = self.move_indices(op) + self.load_queue.clear() + producer_event = self.layer_done_counter.events[producer_id] + producer_event.start_event.record() + with device_module.stream(self.load_stream): + producer_event.start_event.wait(self.load_stream) + for i in range(self.transfer_layer_num): + self.mem_pool_host.load_to_device_per_layer( + self.mem_pool_device, + host_indices, + device_indices, + i, + self.io_backend, + pool_transfers=op.pool_transfers, + ) + producer_event.complete(i) + if host_indices.is_cuda: + host_indices.record_stream(self.load_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.load_stream) + self.ack_load_queue.append( + HiCacheAck( + producer_event.start_event, + producer_event.finish_event, + op.node_ids, + ) + ) + return producer_id + + def prefetch( + self, + request_id: str, + host_indices: torch.Tensor, + new_input_tokens: List[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[List[str]] = None, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> PrefetchOperation: + operation = PrefetchOperation( + request_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys=prefix_keys, + pool_transfers=extra_pools, + ) + self.prefetch_queue.put(operation) + return operation + + def write_storage( + self, + host_indices: torch.Tensor, + token_ids: List[int], + hash_value: Optional[List[str]] = None, + prefix_keys: Optional[List[str]] = None, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> int: + operation = StorageOperation( + host_indices, + token_ids, + hash_value=hash_value, + prefix_keys=prefix_keys, + pool_transfers=extra_pools, + ) + self.backup_queue.put(operation) + return operation.id + + def _storage_hit_query(self, operation) -> tuple[list[str], int]: + last_hash = operation.last_hash + hash_value = [] + for start in range(0, len(operation.token_ids), self.page_size): + last_hash = self.get_hash_str( + operation.token_ids[start : start + self.page_size], last_hash + ) + hash_value.append(last_hash) + + extra_info = HiCacheStorageExtraInfo( + prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None + ) + if operation.pool_transfers: + hit_result = self.storage_backend.batch_exists_v2( + hash_value, operation.pool_transfers, extra_info + ) + else: + kv_hit_count = self.storage_backend.batch_exists(hash_value, extra_info) + hit_result = PoolTransferResult( + kv_hit_pages=kv_hit_count, extra_pool_hit_pages={} + ) + + kv_hit_pages = hit_result.kv_hit_pages + operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages) + + if kv_hit_pages > 0 and operation.pool_transfers: + self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages) + + return ( + hash_value[:kv_hit_pages], + kv_hit_pages * self.page_size, + ) + + def _page_transfer(self, operation): + super()._page_transfer(operation) + if operation.pool_transfers and not operation.is_terminated(): + results = self.storage_backend.batch_get_v2(operation.pool_transfers) + operation.pool_storage_result.update_extra_pool_hit_pages(results) + + def _page_backup(self, operation): + super()._page_backup(operation) + + if operation.pool_transfers: + results = self.storage_backend.batch_set_v2(operation.pool_transfers) + operation.pool_storage_result.update_extra_pool_hit_pages(results) + + def _sync_trailing_keys( + self, + pool_transfers: list[PoolTransfer], + all_hashes: list[str], + kv_hit_pages: int, + ) -> None: + """Re-align trailing-page sidecar keys after KV hit truncation. + + When the storage hit is shorter than the original target prefix, each + pool transfer's keys must be updated to the last N hashes of the actual + hit range instead of the last N hashes of the original target range. + For mamba (N=1) this is just the last hit page hash; for SWA (N>1) it + is a sliding window of the last N hit pages. + """ + for transfer in pool_transfers: + if transfer.hit_policy != "trailing_pages": + continue + trailing_n = len(transfer.keys) if transfer.keys else 1 + transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages] + + def _resolve_pool_transfers_allocation( + self, + extra_pools: Optional[list[PoolTransfer]], + alloc_host: bool, + ) -> Optional[list[PoolTransfer]]: + """Auto-alloc host or device indices for PoolTransfers where they are None.""" + if not extra_pools: + return None + newly_allocated: list[tuple[PoolTransfer, Any, torch.Tensor]] = [] + for pool in extra_pools: + entry = self.mem_pool_host.entry_map.get(pool.name) + if entry is None: + continue + if alloc_host: + if pool.host_indices is not None or pool.device_indices is None: + continue + entry_pool, evict_fn, size = ( + entry.host_pool, + entry.host_evict_fn, + len(pool.device_indices), + ) + else: + if pool.device_indices is not None or pool.host_indices is None: + continue + entry_pool, evict_fn, size = ( + entry.device_pool, + entry.device_evict_fn, + len(pool.host_indices), + ) + indices = entry_pool.alloc(size) + if indices is None and evict_fn: + evict_fn(size) + indices = entry_pool.alloc(size) + if indices is None: + # Roll back all previous allocations using each pool's own entry_pool. + for prev_pool, prev_entry_pool, prev_indices in newly_allocated: + prev_entry_pool.free(prev_indices) + if alloc_host: + prev_pool.host_indices = None + else: + prev_pool.device_indices = None + return None + if alloc_host: + pool.host_indices = indices + else: + pool.device_indices = indices + newly_allocated.append((pool, entry_pool, indices)) + return extra_pools diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 93ebd07a1ae4..f0e7f185f70d 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -71,6 +71,7 @@ def __init__(self, id: Optional[int] = None): self.key: RadixKey = None self.value: Optional[torch.Tensor] = None self.mamba_value: Optional[torch.Tensor] = None + self.mamba_host_value: Optional[torch.Tensor] = None # invariant: for any node, if mamba_lock_ref is locked, full_lock_ref must be locked; # if full_lock_ref is locked, mamba_lock_ref doesn't need to be locked. So, # full_lock_ref is always >= mamba_lock_ref. @@ -95,6 +96,8 @@ def __init__(self, id: Optional[int] = None): self.next = None self.mamba_prev = None self.mamba_next = None + self.host_mamba_prev = None + self.host_mamba_next = None self.id = TreeNode.counter if id is None else id TreeNode.counter += 1 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 16b1410c3090..4298cca4a074 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -469,6 +469,8 @@ def __init__( self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver + self.layer_transfer_counter = None + self.model_layer_id_to_local: dict[int, int] = {} self._init_mamba_pool( size=mamba_size, mamba_spec_state_size=mamba_spec_state_size, @@ -510,6 +512,16 @@ def _init_mamba_pool( ) ) + def register_layer_transfer_counter( + self, layer_transfer_counter: "LayerDoneCounter" + ): + self.layer_transfer_counter = layer_transfer_counter + + def set_model_layer_id_mapping(self, model_layer_ids: List[int]): + self.model_layer_id_to_local = { + layer_id: i for i, layer_id in enumerate(model_layer_ids) + } + # For chunk prefill req, we do not need to allocate mamba cache, # We could use allocated mamba cache instead. def alloc(self, reqs: List["Req"]) -> Optional[List[int]]: @@ -564,6 +576,9 @@ def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor: def mamba2_layer_cache(self, layer_id: int): assert layer_id in self.mamba_map + if self.layer_transfer_counter is not None: + wait_id = self.model_layer_id_to_local.get(layer_id, layer_id) + self.layer_transfer_counter.wait_until(wait_id) return self.mamba_pool.mamba2_layer_cache(self.mamba_map[layer_id]) def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState: @@ -1234,6 +1249,8 @@ def __init__( self.page_size = page_size # TODO support pp? self.start_layer = 0 + self.layer_transfer_counter = None + self.model_layer_id_to_local: dict[int, int] = {} self.head_num = head_num self.head_dim = head_dim self.mamba_pool = mamba_pool @@ -1317,15 +1334,34 @@ def _transfer_full_attention_id(self, layer_id: int): ) return self.full_attention_layer_id_mapping[layer_id] + def register_layer_transfer_counter( + self, layer_transfer_counter: "LayerDoneCounter" + ): + self.layer_transfer_counter = layer_transfer_counter + + def set_model_layer_id_mapping(self, model_layer_ids: List[int]): + self.model_layer_id_to_local = { + layer_id: i for i, layer_id in enumerate(model_layer_ids) + } + + def _wait_for_model_layer(self, layer_id: int): + if self.layer_transfer_counter is None: + return + wait_id = self.model_layer_id_to_local.get(layer_id, layer_id) + self.layer_transfer_counter.wait_until(wait_id) + def get_key_buffer(self, layer_id: int): + self._wait_for_model_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_key_buffer(layer_id) def get_value_buffer(self, layer_id: int): + self._wait_for_model_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_value_buffer(layer_id) def get_kv_buffer(self, layer_id: int): + self._wait_for_model_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_kv_buffer(layer_id) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 0b8199d9f7e9..4f291445abc5 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -2,9 +2,11 @@ import logging import threading from collections import defaultdict +from dataclasses import dataclass from functools import wraps -from typing import Optional +from typing import Any, Callable, Optional +import numpy as np import psutil import torch @@ -19,6 +21,7 @@ ) from sglang.srt.mem_cache.memory_pool import ( KVCache, + MambaPool, MHATokenToKVPool, MLATokenToKVPool, NSATokenToKVPool, @@ -1073,6 +1076,370 @@ def get_page_buffer_meta(self, indices): return ptr_list, element_size_list +class MambaPoolHost(HostKVCache): + def __init__( + self, + device_pool: MambaPool, + host_to_device_ratio: float, + host_size: int, + pin_memory: bool = True, + device: str = "cpu", + allocator_type: str = "default", + ): + self.device_pool = device_pool + self.page_size = 1 + self.layout = "layer_first" + self.pin_memory = pin_memory + self.device = device + self.allocator = get_allocator_from_storage(allocator_type) + self.num_mamba_layers = device_pool.num_mamba_layers + self.conv_state_shapes = [ + conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv + ] + self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:] + self.conv_dtype = device_pool.mamba_cache.conv[0].dtype + self.temporal_dtype = device_pool.mamba_cache.temporal.dtype + self.dtype = self.conv_dtype + self.size_per_token = self.get_size_per_token() + if host_size > 0: + self.size = int(host_size * 1e9 // self.size_per_token) + else: + self.size = int(device_pool.size * host_to_device_ratio) + self.page_num = self.size // self.page_size + 1 + self.size = self.page_num * self.page_size + + assert ( + self.size > device_pool.size + ), "The host memory should be larger than the device memory with the current protocol" + + host_mem = psutil.virtual_memory() + requested_bytes = self.size * self.size_per_token + ten_gb = 10 * (1024**3) + available_bytes = host_mem.available - ten_gb + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory available. Requesting " + f"{requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free. Please reduce the " + f"size of the hierarchical cache." + ) + logger.info( + "Allocating %.2f GB host memory for hierarchical Mamba cache.", + requested_bytes / 1e9, + ) + + self.init_kv_buffer() + self.lock = threading.RLock() + self.clear() + + def _iter_serialized_page_tensors(self, index: int): + yield self.temporal_buffer[:, index : index + self.page_size] + for conv_buf in self.conv_buffer: + yield conv_buf[:, index : index + self.page_size] + + @staticmethod + def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor: + return tensor.contiguous().view(torch.uint8).reshape(-1) + + def init_kv_buffer(self): + alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] + temporal_dims = (self.num_mamba_layers, self.size) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.num_mamba_layers, self.size) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + ) + + @synchronized + def clear(self): + self.mem_state = torch.zeros( + (self.size,), dtype=torch.uint8, device=self.device + ) + self.free_slots = torch.arange(self.size, dtype=torch.int64) + + def available_size(self): + return len(self.free_slots) + + @synchronized + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + assert ( + need_size % self.page_size == 0 + ), "The requested size should be a multiple of the page size." + if need_size > self.available_size(): + return None + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + return select_index + + @synchronized + def free(self, indices: torch.Tensor) -> int: + self.free_slots = torch.cat([self.free_slots, indices]) + return len(indices) + + def get_size_per_token(self): + conv_total_size = 0 + for conv_shape in self.conv_state_shapes: + conv_total_size += int(np.prod(conv_shape)) * self.conv_dtype.itemsize + temporal_size = ( + int(np.prod(self.temporal_state_shape)) * self.temporal_dtype.itemsize + ) + return (conv_total_size + temporal_size) * self.num_mamba_layers + + @staticmethod + def _item_size_per_index(tensor: torch.Tensor) -> int: + if tensor.shape[0] == 0: + return 0 + return int(tensor[0].numel() * tensor.element_size()) + + @staticmethod + def _copy_tensor( + src: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + io_backend: Optional[str], + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel" and not (_is_npu or _is_xpu): + transfer_kv_per_layer_mla( + src=src, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=MambaPoolHost._item_size_per_index(src), + ) + return + if io_backend == "direct" and not (_is_npu or _is_xpu): + transfer_kv_direct( + src_layers=[src], + dst_layers=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + page_size=1, + ) + return + + src_take = src.index_select(0, src_indices.to(src.device)) + dst.index_copy_(0, dst_indices.to(dst.device), src_take.to(dst.device)) + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + ): + for conv_idx, host_conv_state in enumerate(self.conv_buffer): + self._copy_tensor( + host_conv_state[layer_id], + device_pool.mamba_cache.conv[conv_idx][layer_id], + host_indices, + device_indices, + io_backend, + ) + self._copy_tensor( + self.temporal_buffer[layer_id], + device_pool.mamba_cache.temporal[layer_id], + host_indices, + device_indices, + io_backend, + ) + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend="kernel" + ): + for layer_id in range(self.num_mamba_layers): + for conv_idx, host_conv_state in enumerate(self.conv_buffer): + self._copy_tensor( + device_pool.mamba_cache.conv[conv_idx][layer_id], + host_conv_state[layer_id], + device_indices, + host_indices, + io_backend, + ) + self._copy_tensor( + device_pool.mamba_cache.temporal[layer_id], + self.temporal_buffer[layer_id], + device_indices, + host_indices, + io_backend, + ) + + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + data_page = torch.cat( + [ + self._flatten_tensor_bytes(tensor) + for tensor in self._iter_serialized_page_tensors(index) + ] + ) + return data_page.flatten() if flat else data_page + + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + self.page_size * self.size_per_token, + dtype=torch.uint8, + device=self.device, + pin_memory=self.pin_memory, + ) + + def set_from_flat_data_page( + self, + index: int, + data_page: torch.Tensor, + ) -> None: + flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1) + expected_num_bytes = self.page_size * self.size_per_token + if flat_bytes.numel() != expected_num_bytes: + raise ValueError( + f"Invalid Mamba page size: expected {expected_num_bytes} bytes, " + f"got {flat_bytes.numel()} bytes." + ) + + start = 0 + for tensor in self._iter_serialized_page_tensors(index): + num_bytes = tensor.numel() * tensor.element_size() + tensor_bytes = flat_bytes[start : start + num_bytes] + start += num_bytes + restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape) + tensor.copy_(restored) + + +@dataclass +class PoolEntry: + name: str + host_pool: Any + device_pool: Any + layer_mapper: Callable[[int], Optional[int]] + is_primary_index_anchor: bool = False + # Optional eviction callbacks for auto-alloc in HybridCacheController. + # host_evict_fn(n): evict n slots from the host pool (used by write()). + # device_evict_fn(n): evict n slots from the device pool (used by load()). + host_evict_fn: Optional[Callable] = None + device_evict_fn: Optional[Callable] = None + + +class HostPoolGroup: + def __init__(self, entries: list[PoolEntry]): + if not entries: + raise ValueError("HostPoolGroup requires at least one pool entry.") + self.entries = entries + self.entry_map = {entry.name: entry for entry in entries} + self.anchor_entry = next( + (entry for entry in entries if entry.is_primary_index_anchor), + entries[0], + ) + + self.layout = self.anchor_entry.host_pool.layout + self.page_size = self.anchor_entry.host_pool.page_size + self.device = self.anchor_entry.host_pool.device + self.size = self.anchor_entry.host_pool.size + + def clear(self) -> None: + for entry in self.entries: + entry.host_pool.clear() + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + return self.anchor_entry.host_pool.alloc(need_size) + + def free(self, indices: torch.Tensor) -> int: + return self.anchor_entry.host_pool.free(indices) + + def get_data_page(self, index, flat: bool = True): + return self.anchor_entry.host_pool.get_data_page(index, flat) + + def get_dummy_flat_data_page(self): + return self.anchor_entry.host_pool.get_dummy_flat_data_page() + + def set_from_flat_data_page(self, index: int, data_page) -> None: + return self.anchor_entry.host_pool.set_from_flat_data_page(index, data_page) + + def _resolve_entry_indices( + self, + entry: PoolEntry, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + pool_transfers: Optional[list] = None, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if entry is self.anchor_entry: + return host_indices, device_indices + for transfer in pool_transfers or []: + if transfer.name != entry.name or transfer.host_indices is None: + continue + return transfer.host_indices, transfer.device_indices + return None + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + pool_transfers: Optional[list] = None, + ) -> None: + for entry in self.entries: + local_layer_id = entry.layer_mapper(layer_id) + if local_layer_id is None: + continue + resolved = self._resolve_entry_indices( + entry, + host_indices, + device_indices, + pool_transfers=pool_transfers, + ) + if resolved is None: + continue + entry_host_indices, entry_device_indices = resolved + entry.host_pool.load_to_device_per_layer( + entry.device_pool, + entry_host_indices, + entry_device_indices, + local_layer_id, + io_backend, + ) + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + pool_transfers: Optional[list] = None, + ) -> None: + for entry in self.entries: + resolved = self._resolve_entry_indices( + entry, + host_indices, + device_indices, + pool_transfers=pool_transfers, + ) + if resolved is None: + continue + entry_host_indices, entry_device_indices = resolved + entry.host_pool.backup_from_device_all_layer( + entry.device_pool, + entry_host_indices, + entry_device_indices, + io_backend, + ) + + class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device_pool: NSATokenToKVPool diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index c46ff0a388c7..d47af942f796 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -1,5 +1,6 @@ from __future__ import annotations +import dataclasses import logging from typing import TYPE_CHECKING @@ -43,6 +44,19 @@ _is_npu = is_npu() +def _filter_linear_cache_params_for_pp(model_runner: "ModelRunner", cache_params): + if cache_params is None: + return None + if model_runner.pp_size <= 1: + return cache_params + local_layers = [ + layer_id + for layer_id in cache_params.layers + if model_runner.start_layer <= layer_id < model_runner.end_layer + ] + return dataclasses.replace(cache_params, layers=local_layers) + + class ModelRunnerKVCacheMixin: def get_cell_size_per_token(self: ModelRunner, num_layers: int) -> int: kv_size = torch._utils._element_size(self.kv_cache_dtype) @@ -367,13 +381,16 @@ def _init_pools(self: ModelRunner, max_num_reqs: int): # if max_num_reqs <= 32, we pre-allocate 2x requests pre_alloc_size = max_num_reqs * 2 if max_num_reqs <= 32 else 0 if config := self.mambaish_config: + cache_params = _filter_linear_cache_params_for_pp( + self, config.mamba2_cache_params + ) self.req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, - cache_params=config.mamba2_cache_params, + cache_params=cache_params, speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), pre_alloc_size=pre_alloc_size, @@ -390,6 +407,9 @@ def _init_pools(self: ModelRunner, max_num_reqs: int): pre_alloc_size=pre_alloc_size, ) elif config := self.mambaish_config: + cache_params = _filter_linear_cache_params_for_pp( + self, config.mamba2_cache_params + ) self.req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=self.server_args.max_mamba_cache_size, @@ -398,7 +418,7 @@ def _init_pools(self: ModelRunner, max_num_reqs: int): + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, - cache_params=config.mamba2_cache_params, + cache_params=cache_params, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, enable_overlap_schedule=not self.server_args.disable_overlap_schedule, diff --git a/test/registered/4-gpu-models/test_qwen3_next_models.py b/test/registered/4-gpu-models/test_qwen3_next_models.py index 982cb6dcef85..ac18aebe865e 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models.py @@ -1,3 +1,5 @@ +import shutil +import tempfile import unittest from types import SimpleNamespace @@ -20,12 +22,28 @@ register_cuda_ci(est_time=350, suite="stage-c-test-4-gpu-h100") QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" - +QWEN3_NEXT_8B_MODEL = "/sgl-workspace/mnt/Qwen3.5-9B" ACC_THRESHOLDS = { - QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93}, + QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.7}, + QWEN3_NEXT_8B_MODEL: {"kl_div": 0.0025, "gsm8k": 0.7}, } +def flush_cache(base_url, retries: int = 60, retry_interval: float = 1.0) -> bool: + """Flush device cache to force remote storage access.""" + import time + + for _ in range(retries): + try: + response = requests.post(f"{base_url}/flush_cache", timeout=10) + if response.status_code == 200: + return True + except requests.RequestException: + return False + time.sleep(retry_interval) + return False + + def send_request_helper(base_url: str, text: str): response = requests.post( base_url + "/generate", @@ -126,5 +144,87 @@ def test_prefix_cache_branching(self): print("test_prefix_cache_branching passed") +class TestQwen3NextWithHiCache(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN3_NEXT_8B_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.storage_dir = tempfile.mkdtemp(prefix="qwen3-next-hicache-") + env = { + "FLASHINFER_DISABLE_VERSION_CHECK": "1", + "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.storage_dir, + } + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + env=env, + other_args=[ + "--tp", + "4", + "--mamba-scheduler-strategy", + "extra_buffer", + "--page-size", + "64", + "--max-total-tokens", + "131072", + "--max-mamba-cache-size", + "600", + "--hicache-mem-layout", + "page_first_direct", + "--enable-hierarchical-cache", + "--hicache-ratio", + "2", + "--hicache-size", + "0", + "--hicache-write-policy", + "write_through", + "--hicache-storage-backend", + "file", + "--hicache-storage-prefetch-policy", + "wait_complete", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + shutil.rmtree(cls.storage_dir, ignore_errors=True) + + def _run_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=100, + max_new_tokens=512, + parallel=10, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + return run_eval(args) + + def test_gsm8k(self): + first_metrics = self._run_gsm8k() + print(f"first_metrics={first_metrics}") + self.assertGreaterEqual( + first_metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + print(f"flush cache") + assert flush_cache(self.base_url) + + second_metrics = self._run_gsm8k() + print(f"second_metrics={second_metrics}") + self.assertGreaterEqual( + second_metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + self.assertLessEqual( + abs(second_metrics["accuracy"] - first_metrics["accuracy"]), + 0.05, + f"HiCache prefetch accuracy drift too large: " + f"first={first_metrics['accuracy']}, second={second_metrics['accuracy']}", + ) + + if __name__ == "__main__": unittest.main() From 638ff30a61a59ae337aea464ab9a87ad8cbd6a88 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Fri, 13 Mar 2026 14:13:35 +0800 Subject: [PATCH 02/39] Fix mamba prefix match logic --- .../srt/mem_cache/hi_mamba_radix_cache.py | 140 +++++++----------- .../sglang/srt/mem_cache/mamba_radix_cache.py | 7 + 2 files changed, 59 insertions(+), 88 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index fc015742072a..b031c9d9d4ba 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -341,12 +341,8 @@ def load_back( else: ancestor_node = node - mamba_restore_nodes = list(nodes_to_load) - if ( - last_hit_node not in mamba_restore_nodes - and last_hit_node.mamba_host_value is not None - and last_hit_node.mamba_value is None - ): + mamba_restore_nodes = [] + if (last_hit_node.mamba_backuped and last_hit_node.mamba_evicted): mamba_restore_nodes.append(last_hit_node) delta = self.inc_lock_ref(ancestor_node) @@ -368,7 +364,7 @@ def load_back( self.dec_lock_ref(ancestor_node) return None - mamba_pools, mamba_nodes = self.restore_transfers( + mamba_pools = self.restore_transfers( last_hit_node, mamba_restore_nodes, req ) full_device_indices = self.cache_controller.load( @@ -378,7 +374,7 @@ def load_back( ) if full_device_indices is None: self.evict(EvictParams(num_tokens=len(full_host_indices))) - mamba_pools, mamba_nodes = self.restore_transfers( + mamba_pools = self.restore_transfers( last_hit_node, mamba_restore_nodes, req ) full_device_indices = self.cache_controller.load( @@ -391,7 +387,7 @@ def load_back( # no sufficient GPU memory to load back KV caches return None - self.restore_commit(mamba_nodes, mamba_pools) + self.restore_commit(mamba_restore_nodes, mamba_pools) offset = 0 for n in nodes_to_load: @@ -401,16 +397,15 @@ def load_back( 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) + 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) @@ -764,11 +759,14 @@ def evict_mamba(self, mamba_num: int) -> int: assert x.mamba_lock_ref == 0, f"node is in use, {x.id=}" if x.evicted: + x_next = self.mamba_lru_list.get_prev_no_lock(x) if x.mamba_host_value is None: - break + # TODO: Evicted nodes should always have a host mamba backup? (written by + # write_backup before KV eviction). + x = x_next + continue 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 @@ -778,11 +776,9 @@ def evict_mamba(self, mamba_num: int) -> int: _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) mamba_num_evicted += cascade_mamba elif len(x.children) > 0: - if x.mamba_host_value is None: - break + x_next = self.mamba_lru_list.get_prev_no_lock(x) 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: @@ -985,10 +981,6 @@ def _match_prefix_helper( value: List[torch.Tensor] = [] best_value_len = 0 best_last_node = node - mamba_available_size = ( - self.req_to_token_pool.mamba_pool.available_size() - + self.mamba_evictable_size() - ) while len(key) > 0 and child_key in node.children.keys(): child = node.children[child_key] @@ -996,9 +988,7 @@ def _match_prefix_helper( if child.evicted and not child.backuped: break - if node.mamba_value is not None or ( - node.mamba_host_value is not None and mamba_available_size > 0 - ): + if node.mamba_value is not None or node.mamba_backuped: best_value_len = len(value) best_last_node = node @@ -1017,13 +1007,11 @@ def _match_prefix_helper( if len(key): child_key = self.get_child_key_fn(key) - if node.mamba_value is not None or ( - node.mamba_host_value is not None and mamba_available_size > 0 - ): + if node.mamba_value is not None or node.mamba_backuped: best_value_len = len(value) best_last_node = node - deepest_node = node + deepest_node = best_last_node return value, best_last_node, best_value_len, deepest_node def _match_post_processor( @@ -1066,31 +1054,15 @@ def _match_post_processor( else: mamba_branching_seqlen = None - # last_device_node & host_hit_length: - # last_device_node is the last node whose full and mamba states can both - # be directly used on device. host_hit_length remains full-KV-only. host_hit_length = 0 last_device_node = best_last_node - while last_device_node is not self.root_node and ( - last_device_node.evicted - or ( - last_device_node.mamba_value is None - and last_device_node.mamba_host_value is not None - ) - ): - if last_device_node.evicted: - host_hit_length += len(last_device_node.host_value) + while last_device_node is not self.root_node and 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 + while last_host_node is not self.root_node and not last_host_node.backuped: + last_host_node = last_host_node.parent mamba_node = best_last_node if cow_mamba and mamba_node.mamba_value is not None: @@ -1120,7 +1092,7 @@ def _match_post_processor( device_indices=value, last_device_node=last_device_node, last_host_node=last_host_node, - last_host_backup_node=last_host_backup_node, + last_host_backup_node=last_host_node, host_hit_length=host_hit_length, mamba_branching_seqlen=mamba_branching_seqlen, ) @@ -2104,21 +2076,12 @@ def restore_transfers( last_hit_node: TreeNode, nodes_to_restore: list[TreeNode], req, - ) -> tuple[Optional[list[PoolTransfer]], list[TreeNode]]: - """PoolTransfers for H→D restore (load_back). - - Node slots use device_indices=None so the controller auto-allocates them, - mirroring KV. The req slot is pre-allocated and must land in req.mamba_pool_idx. - Returns (transfers, extra_nodes) where extra_nodes need device write-back via - restore_commit() after load succeeds. - """ + ) -> Optional[list[PoolTransfer]]: + """PoolTransfers for H→D restore (load_back).""" mamba_host_list: list[torch.Tensor] = [] - extra_nodes: list[TreeNode] = [] - for node in nodes_to_restore: - if node.mamba_host_value is None or node.mamba_value is not None: + if not node.mamba_backuped: continue - extra_nodes.append(node) mamba_host_list.append(node.mamba_host_value) transfers: list[PoolTransfer] = [] @@ -2130,28 +2093,29 @@ def restore_transfers( ) ) - 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_pool, - len(last_hit_node.mamba_host_value), - self.evict_mamba, - lock_node=last_hit_node, - error_message="Can not alloc request mamba cache for host load back", - )[0] - transfers.append( - PoolTransfer( - name="mamba", - host_indices=last_hit_node.mamba_host_value, - device_indices=req.mamba_pool_idx.unsqueeze(0), - ) - ) - - return (transfers if transfers else None), extra_nodes + # TODO(hzh): FIX Req's mamba restore logic here + # 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_pool, + # len(last_hit_node.mamba_host_value), + # self.evict_mamba, + # lock_node=last_hit_node, + # error_message="Can not alloc request mamba cache for host load back", + # )[0] + # transfers.append( + # PoolTransfer( + # name="mamba", + # host_indices=last_hit_node.mamba_host_value, + # device_indices=req.mamba_pool_idx.unsqueeze(0), + # ) + # ) + + return (transfers if transfers else None) def restore_commit( self, diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index f0e7f185f70d..4ced651fa722 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -106,10 +106,17 @@ def __init__(self, id: Optional[int] = None): def evicted(self): return self.value is None + @property + def mamba_evicted(self): + return self.mamba_value is None + @property def backuped(self): return self.host_value is not None + @property + def mamba_backuped(self): + return self.mamba_host_value is not None def protect_host(self): """Protect the host value from eviction.""" self.host_ref_counter += 1 From 3556ed6e9ec66f14f29f0ea45f108778eae550fe Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Sat, 14 Mar 2026 15:03:01 +0000 Subject: [PATCH 03/39] Enhance evict logic & Loading host mamba into both req and node's mamba space --- .../srt/mem_cache/hi_mamba_radix_cache.py | 330 ++++++++++-------- .../sglang/srt/mem_cache/mamba_radix_cache.py | 1 + .../sglang/srt/mem_cache/memory_pool_host.py | 68 ++-- 3 files changed, 216 insertions(+), 183 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index b031c9d9d4ba..8a65ab9fd6d5 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -342,7 +342,7 @@ def load_back( ancestor_node = node mamba_restore_nodes = [] - if (last_hit_node.mamba_backuped and last_hit_node.mamba_evicted): + if last_hit_node.mamba_backuped and last_hit_node.mamba_evicted: mamba_restore_nodes.append(last_hit_node) delta = self.inc_lock_ref(ancestor_node) @@ -364,9 +364,7 @@ def load_back( self.dec_lock_ref(ancestor_node) return None - mamba_pools = self.restore_transfers( - last_hit_node, mamba_restore_nodes, req - ) + mamba_pools = self.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, @@ -564,24 +562,114 @@ def _update_full_host_leaf_status(self, node: TreeNode): 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) + def _free_gpu_mamba(self, node: TreeNode) -> int: + """Free GPU mamba on a node. Returns num mamba tokens freed.""" + if node.mamba_value is None: + return 0 + mamba_num = len(node.mamba_value) + self.req_to_token_pool.mamba_pool.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]: + """Evict a backuped device node to host: free GPU KV + GPU mamba. + + Node stays in the tree as evicted+backuped. + Caller must ensure node.backuped is True before calling. + Returns (full_num_evicted, mamba_num_evicted). + """ + assert not node.evicted, f"_evict_to_host on already-evicted node, {node.id=}" + assert node.backuped, f"_evict_to_host on non-backuped node, {node.id=}" - 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) + num_full = len(node.value) 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_gpu_mamba(node) + node.value = None self._update_leaf_status(node) self._update_full_device_leaf_status(node.parent) - return num_full + return num_full, mamba_num + + def _evict_regular_leaf(self, node: TreeNode) -> Tuple[int, int]: + """Evict a non-backuped device leaf: free GPU KV + mamba, delete from tree. + + Used in write_through when a node was never backed up. + Returns (full_num_evicted, mamba_num_evicted). + """ + assert not node.evicted, f"_evict_regular on already-evicted node, {node.id=}" + assert not node.backuped, f"_evict_regular on backuped node, {node.id=}" + assert len(node.children) == 0, f"_evict_regular on non-leaf, {node.id=}" + + num_full = len(node.value) + + 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_gpu_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 = 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._update_leaf_status(parent) + self._iteratively_delete_tombstone_leaf(node) + return num_full, mamba_num + + def _evict_host_leaf_node(self, x: TreeNode) -> int: + """Evict a host-resident leaf: free host KV + host mamba, delete from tree, cascade. + + Returns num host KV tokens evicted. + """ + assert x.evicted, f"host leaf not evicted, {x.id=}" + assert x.backuped, f"host leaf not backuped, {x.id=}" + assert x.mamba_value is None, f"host leaf has GPU mamba, {x.id=}" + assert ( + x.host_ref_counter == 0 + ), f"host leaf in use, {x.id=} {x.host_ref_counter=}" + + full_num_evicted = self.cache_controller.evict_host(x.host_value) + x.host_value = None + + if x.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(x): + self.mamba_host_lru_list.remove_node(x) + self.mamba_pool_host.free(x.mamba_host_value) + x.mamba_host_value = None + + self._discard_from_leaf_sets(x) + parent = x.parent + key = self.get_child_key_fn(x.key) + v = parent.children.pop(key, None) + assert v == x, f"parent does not have child key, {key}" + + self._update_leaf_status(parent) + self._iteratively_delete_tombstone_leaf(x) + + return full_num_evicted def _delete_tombstone_leaf(self, node: TreeNode) -> None: """Remove a tombstone leaf from the tree and free HiCache resources.""" @@ -635,6 +723,22 @@ def _iteratively_delete_tombstone_leaf( 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_leaf (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_leaf(x) + return self._evict_to_host(x) + def evict(self, params: EvictParams) -> EvictResult: if self.disable: return EvictResult() @@ -653,15 +757,16 @@ def evict(self, params: EvictParams) -> EvictResult: if x not in self.evictable_full_device_leaves: continue - evicted_full = self._evict_to_host(x) + 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) + mamba_num_evicted += self.evict_mamba(params.mamba_num) return EvictResult( num_tokens_evicted=full_num_evicted, @@ -669,32 +774,7 @@ def evict(self, params: EvictParams) -> EvictResult: ) 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 - if x.mamba_host_value is not None: - if self.mamba_host_lru_list.in_list(x): - self.mamba_host_lru_list.remove_node(x) - self.mamba_pool_host.free(x.mamba_host_value) - x.mamba_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 + """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) @@ -704,34 +784,18 @@ def evict_host(self, num_tokens: int): if x not in self.evictable_full_host_leaves: continue - num_evicted += self.cache_controller.evict_host(x.host_value) - x.host_value = None + num_evicted += self._evict_host_leaf_node(x) - 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 - - if x.mamba_host_value is not None: - if self.mamba_host_lru_list.in_list(x): - self.mamba_host_lru_list.remove_node(x) - self.mamba_pool_host.free(x.mamba_host_value) - x.mamba_host_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)) + 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 @@ -739,15 +803,31 @@ def evict_mamba_host(self, num_mamba_hosts: int) -> int: 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.host_ref_counter == 0: + if x.host_ref_counter > 0: + x = x_next + continue + + if x in self.evictable_full_host_leaves: + # Host leaf: evict host KV + mamba, delete from tree + self._evict_host_leaf_node(x) + num_evicted += 1 + else: + # Internal host node: free host mamba only (tombstone) 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 @@ -757,74 +837,25 @@ def evict_mamba(self, mamba_num: int) -> int: 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 x.evicted: + if len(x.children) > 0: + # Internal: free GPU mamba only, KV stays on GPU (tombstone) x_next = self.mamba_lru_list.get_prev_no_lock(x) - if x.mamba_host_value is None: - # TODO: Evicted nodes should always have a host mamba backup? (written by - # write_backup before KV eviction). - x = x_next - continue - self.req_to_token_pool.mamba_pool.free(x.mamba_value) - mamba_num_evicted += len(x.mamba_value) - self.mamba_lru_list.remove_node(x) - self.mamba_evictable_size_ -= len(x.mamba_value) - x.mamba_value = None - - if len(x.children) == 0 and x.mamba_host_value is None: - self._delete_tombstone_leaf(x) - _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) - mamba_num_evicted += cascade_mamba - elif len(x.children) > 0: - x_next = self.mamba_lru_list.get_prev_no_lock(x) - self.req_to_token_pool.mamba_pool.free(x.mamba_value) mamba_num_evicted += len(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 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: - pass - elif x.mamba_host_value is not None: - if self.mamba_host_lru_list.in_list(x): - self.mamba_host_lru_list.remove_node(x) - self.mamba_pool_host.free(x.mamba_host_value) - x.mamba_host_value = None - - x.value = None - x.mamba_value = None - - self._discard_from_leaf_sets(x) - - if not x.backuped: - 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 + _, 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() @@ -835,6 +866,11 @@ def evict_mamba(self, mamba_num: int) -> int: def _unevict_node(self, node: TreeNode, fresh_value: torch.Tensor): """Restore an evicted node with fresh device KV from the request.""" + assert node.evicted, f"_unevict_node on non-evicted node, {node.id=}" + # invariant: evicted => no GPU mamba + assert ( + node.mamba_value is None + ), f"evicted node should not have GPU mamba, {node.id=}" n = len(fresh_value) node.value = fresh_value.clone() @@ -2085,37 +2121,39 @@ def restore_transfers( mamba_host_list.append(node.mamba_host_value) transfers: list[PoolTransfer] = [] + # Transfer from host to node's device mamba space if mamba_host_list: transfers.append( PoolTransfer( name="mamba", host_indices=torch.cat(mamba_host_list), + device_indices=None, # will be allocated by cache_controller + ) + ) + + # Transfer from host to request's device mamba space + 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_pool, + len(last_hit_node.mamba_host_value), + self.evict_mamba, + lock_node=last_hit_node, + error_message="Can not alloc request mamba cache for host load back", + )[0] + transfers.append( + PoolTransfer( + name="mamba", + host_indices=last_hit_node.mamba_host_value, + device_indices=req.mamba_pool_idx.unsqueeze(0), ) ) - # TODO(hzh): FIX Req's mamba restore logic here - # 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_pool, - # len(last_hit_node.mamba_host_value), - # self.evict_mamba, - # lock_node=last_hit_node, - # error_message="Can not alloc request mamba cache for host load back", - # )[0] - # transfers.append( - # PoolTransfer( - # name="mamba", - # host_indices=last_hit_node.mamba_host_value, - # device_indices=req.mamba_pool_idx.unsqueeze(0), - # ) - # ) - - return (transfers if transfers else None) + return transfers if transfers else None def restore_commit( self, diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 4ced651fa722..7cd801b991ea 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -117,6 +117,7 @@ def backuped(self): @property def mamba_backuped(self): return self.mamba_host_value is not None + def protect_host(self): """Protect the host value from eviction.""" self.host_ref_counter += 1 diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 4f291445abc5..cc3745e82c3c 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1369,21 +1369,6 @@ def get_dummy_flat_data_page(self): def set_from_flat_data_page(self, index: int, data_page) -> None: return self.anchor_entry.host_pool.set_from_flat_data_page(index, data_page) - def _resolve_entry_indices( - self, - entry: PoolEntry, - host_indices: torch.Tensor, - device_indices: torch.Tensor, - pool_transfers: Optional[list] = None, - ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: - if entry is self.anchor_entry: - return host_indices, device_indices - for transfer in pool_transfers or []: - if transfer.name != entry.name or transfer.host_indices is None: - continue - return transfer.host_indices, transfer.device_indices - return None - def load_to_device_per_layer( self, device_pool, @@ -1393,23 +1378,30 @@ def load_to_device_per_layer( io_backend, pool_transfers: Optional[list] = None, ) -> None: - for entry in self.entries: - local_layer_id = entry.layer_mapper(layer_id) - if local_layer_id is None: - continue - resolved = self._resolve_entry_indices( - entry, + # 1. Anchor (KV) transfer + anchor = self.anchor_entry + local_layer_id = anchor.layer_mapper(layer_id) + if local_layer_id is not None: + anchor.host_pool.load_to_device_per_layer( + anchor.device_pool, host_indices, device_indices, - pool_transfers=pool_transfers, + local_layer_id, + io_backend, ) - if resolved is None: + + # 2. Extra pool transfers + for transfer in pool_transfers or []: + entry = self.entry_map.get(transfer.name) + if entry is None or transfer.host_indices is None: + continue + local_layer_id = entry.layer_mapper(layer_id) + if local_layer_id is None: continue - entry_host_indices, entry_device_indices = resolved entry.host_pool.load_to_device_per_layer( entry.device_pool, - entry_host_indices, - entry_device_indices, + transfer.host_indices, + transfer.device_indices, local_layer_id, io_backend, ) @@ -1422,20 +1414,22 @@ def backup_from_device_all_layer( io_backend, pool_transfers: Optional[list] = None, ) -> None: - for entry in self.entries: - resolved = self._resolve_entry_indices( - entry, - host_indices, - device_indices, - pool_transfers=pool_transfers, - ) - if resolved is None: + # 1. Anchor (KV) backup + self.anchor_entry.host_pool.backup_from_device_all_layer( + self.anchor_entry.device_pool, + host_indices, + device_indices, + io_backend, + ) + # 2. Extra pool transfers + for transfer in pool_transfers or []: + entry = self.entry_map.get(transfer.name) + if entry is None or transfer.host_indices is None: continue - entry_host_indices, entry_device_indices = resolved entry.host_pool.backup_from_device_all_layer( entry.device_pool, - entry_host_indices, - entry_device_indices, + transfer.host_indices, + transfer.device_indices, io_backend, ) From 282737af1c92ba729d9ce21226b926b1f786d2bc Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Sun, 15 Mar 2026 07:36:27 +0000 Subject: [PATCH 04/39] Fix hang issue --- python/sglang/srt/managers/cache_controller.py | 2 +- .../hybrid_cache/hybrid_cache_controller.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ff3376ca9031..d4c31e2d70f4 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -971,7 +971,7 @@ def prefetch_thread_func(self): operation.host_indices[storage_hit_count:] ) operation.host_indices = operation.host_indices[:storage_hit_count] - logger.debug( + logger.info( f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}." ) self.prefetch_buffer.put(operation) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 904728b422c9..9b3a142ba50f 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -408,18 +408,23 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]: ) def _page_transfer(self, operation): - super()._page_transfer(operation) + # Transfer extra pools if operation.pool_transfers and not operation.is_terminated(): results = self.storage_backend.batch_get_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) - def _page_backup(self, operation): - super()._page_backup(operation) + # Transfer kv pools + super()._page_transfer(operation) + def _page_backup(self, operation): + # Backup extra pools if operation.pool_transfers: results = self.storage_backend.batch_set_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) - + + # Backup kv pools + super()._page_backup(operation) + def _sync_trailing_keys( self, pool_transfers: list[PoolTransfer], From f2a062991130132e44dffa49d6efd2f82a7b67b9 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Sun, 15 Mar 2026 09:01:59 +0000 Subject: [PATCH 05/39] Refactor prefetch logic; fix memory leak issue --- .../srt/mem_cache/hi_mamba_radix_cache.py | 139 +++++++----------- 1 file changed, 54 insertions(+), 85 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 8a65ab9fd6d5..899d5cf08755 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -18,7 +18,7 @@ MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult +from sglang.srt.mem_cache.hicache_storage import PoolTransfer from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, PrefetchOperation, @@ -330,6 +330,9 @@ def write_backup(self, node: TreeNode, write_back=False): def load_back( self, node: TreeNode, mem_quota: Optional[int] = None, req=None ) -> Optional[torch.Tensor]: + logger.info( + f"Try load back, node: {node.id}, mem_quota: {mem_quota}, req: {req.rid if req is not None else None}" + ) """Load full KV back from host.""" last_hit_node = node nodes_to_load = [] @@ -365,6 +368,7 @@ def load_back( return None mamba_pools = self.restore_transfers(last_hit_node, mamba_restore_nodes, req) + logger.info(f"Try init load back, mamba_pools: {mamba_pools}") full_device_indices = self.cache_controller.load( host_indices=full_host_indices, node_id=last_hit_node.id, @@ -845,6 +849,7 @@ def evict_mamba(self, mamba_num: int) -> int: # Internal: free GPU mamba only, KV stays on GPU (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_pool.free(x.mamba_value) self.mamba_lru_list.remove_node(x) self._tombstone_internal_node(x) else: @@ -1777,6 +1782,8 @@ def prefetch_from_storage( return self._protect_host_node(last_host_node) + + # Allocate host KV memory host_indices = self._alloc_with_evict( self.cache_controller.mem_pool_host, prefetch_length, @@ -1785,21 +1792,20 @@ def prefetch_from_storage( if host_indices is None: self._release_host_node(last_host_node) return - extra_pools = None - prepared = self.prefetch_prepare(new_input_tokens, last_hash) - if prefetch_length > 0 and prepared is None: + + # Allocate host mamba slot + extra_pools = self.prefetch_prepare(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) return - if prepared is not None: - _, extra_pools = prepared - if extra_pools is not None: - logger.info( - "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", - req_id, - prefetch_length // self.page_size, - 1, - ) + + logger.info( + "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", + req_id, + prefetch_length // self.page_size, + 1, + ) operation = self.cache_controller.prefetch( req_id, host_indices, @@ -1845,13 +1851,21 @@ def check_prefetch_progress(self, req_id: str) -> bool: 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] + + # Resolve mamba transfer state mamba_host_indices = None + mamba_loaded = False for transfer in operation.pool_transfers or []: if transfer.name == "mamba": mamba_host_indices = transfer.host_indices + mamba_loaded = ( + operation.pool_storage_result.extra_pool_hit_pages.get("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( @@ -1861,19 +1875,21 @@ def check_prefetch_progress(self, req_id: str) -> bool: written_indices, hash_value[: min_completed_tokens // self.page_size], mamba_host_indices, - operation.pool_storage_result, + 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]) - mamba_loaded = self.prefetch_commit( - operation.pool_transfers, - matched_length, - min_completed_tokens, - operation.pool_storage_result, - ) 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) @@ -1900,7 +1916,7 @@ def _insert_helper_host( host_value, hash_value, mamba_host_value: Optional[torch.Tensor] = None, - pool_storage_result: Optional[PoolTransferResult] = None, + mamba_loaded: bool = False, ): node.last_access_time = get_last_access_time() if len(key) == 0: @@ -1909,11 +1925,6 @@ def _insert_helper_host( child_key = self.get_child_key_fn(key) matched_length = 0 - host_value_inserted = False - final_mamba_node: Optional[TreeNode] = None - has_mamba = pool_storage_result is None or ( - pool_storage_result.extra_pool_hit_pages.get("mamba", 0) >= 1 - ) while len(key) > 0 and child_key in node.children.keys(): node = node.children[child_key] node.last_access_time = get_last_access_time() @@ -1921,24 +1932,10 @@ def _insert_helper_host( 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 - if prefix_len == len(key): - final_mamba_node = node - 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 :] + matched_length += prefix_len if prefix_len < len(node.key): new_node = self._split_node(node.key, node, prefix_len) @@ -1947,6 +1944,8 @@ def _insert_helper_host( if len(key): child_key = self.get_child_key_fn(key) + # Insert remaining tokens as a new leaf node + leaf_node: Optional[TreeNode] = None if len(key): new_node = TreeNode() new_node.parent = node @@ -1956,13 +1955,15 @@ def _insert_helper_host( new_node.host_value = host_value.clone() new_node.hash_value = hash_value node.children[child_key] = new_node - final_mamba_node = new_node + leaf_node = new_node self._update_full_host_leaf_status(new_node) self._update_full_host_leaf_status(node) - if final_mamba_node is not None and mamba_host_value is not None and has_mamba: - final_mamba_node.mamba_host_value = mamba_host_value.clone() - if not self.mamba_host_lru_list.in_list(final_mamba_node): - self.mamba_host_lru_list.insert_mru(final_mamba_node) + + # 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): @@ -2035,11 +2036,6 @@ def _alloc_with_evict( raise RuntimeError(error_message) return indices - def _last_page_index(self, token_count: int) -> Optional[int]: - if token_count <= 0: - return None - return token_count // self.page_size - 1 - def backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: """PoolTransfers for D→H backup: mamba device → mamba host.""" if node.mamba_value is None: @@ -2079,7 +2075,7 @@ def prefetch_prepare( self, token_ids: List[int], last_hash: Optional[str], - ) -> Optional[tuple[torch.Tensor, list[PoolTransfer]]]: + ) -> Optional[list[PoolTransfer]]: """Alloc mamba host slot and build PoolTransfers for Storage→H prefetch. Returns (host_handle, transfers) so the caller can track the allocation, @@ -2092,20 +2088,17 @@ def prefetch_prepare( ) if mamba_host_index is None: return None - last_page_hash = last_hash - for start in range(0, len(token_ids), self.page_size): - last_page_hash = self.cache_controller.get_hash_str( - token_ids[start : start + self.page_size], last_page_hash - ) + # Use a placeholder key; the I/O thread + # will replace it with the correct hash after query hit determination. transfers = [ PoolTransfer( name="mamba", host_indices=mamba_host_index, - keys=[last_page_hash], + keys=["__placeholder__"], hit_policy="trailing_pages", ) ] - return mamba_host_index, transfers + return transfers def restore_transfers( self, @@ -2169,27 +2162,3 @@ def restore_commit( n_len = len(n.mamba_host_value) n.mamba_value = mamba_device[offset : offset + n_len].clone() offset += n_len - - def prefetch_commit( - self, - pool_transfers: Optional[list[PoolTransfer]], - matched_length: int, - min_completed_tokens: int, - result: PoolTransferResult, - ) -> bool: - """After Storage→H prefetch completes: free mamba host slot if not inserted into tree. - - Returns True if a mamba state was successfully loaded into the radix tree. - """ - mamba_host_indices = None - for transfer in pool_transfers or []: - if transfer.name == "mamba": - mamba_host_indices = transfer.host_indices - break - if mamba_host_indices is None: - return False - # mamba covers the entire prefix as a single page; loaded means >= 1 page hit. - mamba_loaded = result.extra_pool_hit_pages.get("mamba", 0) >= 1 - if matched_length == min_completed_tokens or not mamba_loaded: - self.mamba_pool_host.free(mamba_host_indices) - return mamba_loaded From 7316d0e6db1b91ebcbc5e36e2bbb7dceb940ae64 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Sun, 15 Mar 2026 14:33:39 +0000 Subject: [PATCH 06/39] Support 'page_first_direct' layout for mambaPoolHost --- .../srt/mem_cache/hi_mamba_radix_cache.py | 1 + .../hybrid_cache/hybrid_cache_controller.py | 4 +- .../sglang/srt/mem_cache/memory_pool_host.py | 173 +++++++++++------- 3 files changed, 111 insertions(+), 67 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 899d5cf08755..1be5bd1e3c27 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -147,6 +147,7 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): server_args.hicache_ratio, server_args.hicache_size, allocator_type=server_args.hicache_storage_backend, + layout=server_args.hicache_mem_layout, ) full_layer_mapping = dict(self.hybrid_kv_cache.full_attention_layer_id_mapping) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 9b3a142ba50f..82526085693c 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -421,10 +421,10 @@ def _page_backup(self, operation): if operation.pool_transfers: results = self.storage_backend.batch_set_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) - + # Backup kv pools super()._page_backup(operation) - + def _sync_trailing_keys( self, pool_transfers: list[PoolTransfer], diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index cc3745e82c3c..a0328de58961 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1077,6 +1077,7 @@ def get_page_buffer_meta(self, indices): class MambaPoolHost(HostKVCache): + def __init__( self, device_pool: MambaPool, @@ -1085,10 +1086,12 @@ def __init__( pin_memory: bool = True, device: str = "cpu", allocator_type: str = "default", + layout: str = "layer_first", ): self.device_pool = device_pool self.page_size = 1 - self.layout = "layer_first" + assert layout in ["page_first_direct", "layer_first"], "Unsupported layout: {layout}" + self.layout = layout self.pin_memory = pin_memory self.device = device self.allocator = get_allocator_from_storage(allocator_type) @@ -1124,45 +1127,92 @@ def __init__( f"size of the hierarchical cache." ) logger.info( - "Allocating %.2f GB host memory for hierarchical Mamba cache.", + "Allocating %.2f GB host memory for hierarchical Mamba cache (layout=%s).", requested_bytes / 1e9, + self.layout, ) self.init_kv_buffer() self.lock = threading.RLock() self.clear() - def _iter_serialized_page_tensors(self, index: int): - yield self.temporal_buffer[:, index : index + self.page_size] - for conv_buf in self.conv_buffer: - yield conv_buf[:, index : index + self.page_size] - - @staticmethod - def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor: - return tensor.contiguous().view(torch.uint8).reshape(-1) - def init_kv_buffer(self): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] - temporal_dims = (self.num_mamba_layers, self.size) + self.temporal_state_shape - self.temporal_buffer = alloc_func( - temporal_dims, - dtype=self.temporal_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - self.conv_buffer = [] - for conv_shape in self.conv_state_shapes: - conv_dims = (self.num_mamba_layers, self.size) + conv_shape - self.conv_buffer.append( - alloc_func( - conv_dims, - dtype=self.conv_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, + + if self.layout == "page_first_direct": + # page-first: (size, num_layers, *shape) — per-page data is contiguous + temporal_dims = ( + self.size, + self.num_mamba_layers, + ) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.size, self.num_mamba_layers) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) ) + else: + # layer-first: (num_layers, size, *shape) + temporal_dims = ( + self.num_mamba_layers, + self.size, + ) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.num_mamba_layers, self.size) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + ) + + def _get_temporal_view(self, layer_id: int) -> torch.Tensor: + if self.layout == "page_first_direct": + return self.temporal_buffer[:, layer_id] + return self.temporal_buffer[layer_id] + + def _get_conv_view(self, layer_id: int, conv_idx: int) -> torch.Tensor: + if self.layout == "page_first_direct": + return self.conv_buffer[conv_idx][:, layer_id] + return self.conv_buffer[conv_idx][layer_id] + + def _iter_page_tensors(self, index: int): + if self.layout == "page_first_direct": + yield self.temporal_buffer[index] + for conv_buf in self.conv_buffer: + yield conv_buf[index] + else: + yield self.temporal_buffer[:, index : index + self.page_size] + for conv_buf in self.conv_buffer: + yield conv_buf[:, index : index + self.page_size] + + @staticmethod + def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor: + return tensor.contiguous().view(torch.uint8).reshape(-1) @synchronized def clear(self): @@ -1199,6 +1249,9 @@ def get_size_per_token(self): ) return (conv_total_size + temporal_size) * self.num_mamba_layers + def get_ksize_per_token(self): + return self.get_size_per_token() + @staticmethod def _item_size_per_index(tensor: torch.Tensor) -> int: if tensor.shape[0] == 0: @@ -1211,11 +1264,11 @@ def _copy_tensor( dst: torch.Tensor, src_indices: torch.Tensor, dst_indices: torch.Tensor, - io_backend: Optional[str], + io_backend: str, ) -> None: if src_indices.numel() == 0: return - if io_backend == "kernel" and not (_is_npu or _is_xpu): + if io_backend == "kernel": transfer_kv_per_layer_mla( src=src, dst=dst, @@ -1223,8 +1276,7 @@ def _copy_tensor( dst_indices=dst_indices, item_size=MambaPoolHost._item_size_per_index(src), ) - return - if io_backend == "direct" and not (_is_npu or _is_xpu): + elif io_backend == "direct": transfer_kv_direct( src_layers=[src], dst_layers=[dst], @@ -1232,10 +1284,8 @@ def _copy_tensor( dst_indices=dst_indices, page_size=1, ) - return - - src_take = src.index_select(0, src_indices.to(src.device)) - dst.index_copy_(0, dst_indices.to(dst.device), src_take.to(dst.device)) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") def load_to_device_per_layer( self, @@ -1245,47 +1295,47 @@ def load_to_device_per_layer( layer_id, io_backend="kernel", ): - for conv_idx, host_conv_state in enumerate(self.conv_buffer): - self._copy_tensor( - host_conv_state[layer_id], - device_pool.mamba_cache.conv[conv_idx][layer_id], - host_indices, - device_indices, - io_backend, - ) self._copy_tensor( - self.temporal_buffer[layer_id], + self._get_temporal_view(layer_id), device_pool.mamba_cache.temporal[layer_id], host_indices, device_indices, io_backend, ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + self._get_conv_view(layer_id, conv_idx), + device_pool.mamba_cache.conv[conv_idx][layer_id], + host_indices, + device_indices, + io_backend, + ) def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend="kernel" ): for layer_id in range(self.num_mamba_layers): - for conv_idx, host_conv_state in enumerate(self.conv_buffer): - self._copy_tensor( - device_pool.mamba_cache.conv[conv_idx][layer_id], - host_conv_state[layer_id], - device_indices, - host_indices, - io_backend, - ) self._copy_tensor( device_pool.mamba_cache.temporal[layer_id], - self.temporal_buffer[layer_id], + self._get_temporal_view(layer_id), device_indices, host_indices, io_backend, ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + device_pool.mamba_cache.conv[conv_idx][layer_id], + self._get_conv_view(layer_id, conv_idx), + device_indices, + host_indices, + io_backend, + ) def get_data_page(self, index, flat: bool = True) -> torch.Tensor: data_page = torch.cat( [ self._flatten_tensor_bytes(tensor) - for tensor in self._iter_serialized_page_tensors(index) + for tensor in self._iter_page_tensors(index) ] ) return data_page.flatten() if flat else data_page @@ -1304,15 +1354,8 @@ def set_from_flat_data_page( data_page: torch.Tensor, ) -> None: flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1) - expected_num_bytes = self.page_size * self.size_per_token - if flat_bytes.numel() != expected_num_bytes: - raise ValueError( - f"Invalid Mamba page size: expected {expected_num_bytes} bytes, " - f"got {flat_bytes.numel()} bytes." - ) - start = 0 - for tensor in self._iter_serialized_page_tensors(index): + for tensor in self._iter_page_tensors(index): num_bytes = tensor.numel() * tensor.element_size() tensor_bytes = flat_bytes[start : start + num_bytes] start += num_bytes @@ -1421,7 +1464,7 @@ def backup_from_device_all_layer( device_indices, io_backend, ) - # 2. Extra pool transfers + # 2. Extra pool backup for transfer in pool_transfers or []: entry = self.entry_map.get(transfer.name) if entry is None or transfer.host_indices is None: From 687452000d047f9d31e97d48eccbbec929b711f2 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Sun, 15 Mar 2026 15:30:46 +0000 Subject: [PATCH 07/39] Optimize code style --- .../srt/mem_cache/hi_mamba_radix_cache.py | 216 ++++++++---------- .../sglang/srt/mem_cache/memory_pool_host.py | 5 +- 2 files changed, 94 insertions(+), 127 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 1be5bd1e3c27..634a85cc659d 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -88,11 +88,7 @@ def remove_node(self, node: TreeNode): 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. - """ + """Hierarchical cache for hybrid Mamba models.""" def __init__(self, params: CacheInitParams, server_args: ServerArgs): self._enable_metrics_flag = params.enable_metrics @@ -290,7 +286,7 @@ def write_backup(self, node: TreeNode, write_back=False): if self.mamba_host_lru_list.in_list(node): self.mamba_host_lru_list.reset_node_mru(node) - extra_pools = self.backup_transfers(node) + extra_pools = self.mamba_backup_transfers(node) host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, @@ -306,7 +302,7 @@ def write_backup(self, node: TreeNode, write_back=False): if host_indices is not None: node.host_value = host_indices if extra_pools is not None: - self.backup_commit(node, extra_pools) + self.mamba_backup_commit(node, extra_pools) assert len(node.host_value) > 0 self.ongoing_write_through[node.id] = node if not write_back: @@ -331,9 +327,6 @@ def write_backup(self, node: TreeNode, write_back=False): def load_back( self, node: TreeNode, mem_quota: Optional[int] = None, req=None ) -> Optional[torch.Tensor]: - logger.info( - f"Try load back, node: {node.id}, mem_quota: {mem_quota}, req: {req.rid if req is not None else None}" - ) """Load full KV back from host.""" last_hit_node = node nodes_to_load = [] @@ -368,8 +361,9 @@ def load_back( self.dec_lock_ref(ancestor_node) return None - mamba_pools = self.restore_transfers(last_hit_node, mamba_restore_nodes, req) - logger.info(f"Try init load back, mamba_pools: {mamba_pools}") + 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, @@ -377,7 +371,7 @@ def load_back( ) if full_device_indices is None: self.evict(EvictParams(num_tokens=len(full_host_indices))) - mamba_pools = self.restore_transfers( + mamba_pools = self.mamba_restore_transfers( last_hit_node, mamba_restore_nodes, req ) full_device_indices = self.cache_controller.load( @@ -390,7 +384,7 @@ def load_back( # no sufficient GPU memory to load back KV caches return None - self.restore_commit(mamba_restore_nodes, mamba_pools) + self.mamba_restore_commit(mamba_restore_nodes, mamba_pools) offset = 0 for n in nodes_to_load: @@ -567,8 +561,7 @@ def _update_full_host_leaf_status(self, node: TreeNode): return self.evictable_full_host_leaves.add(node) - def _free_gpu_mamba(self, node: TreeNode) -> int: - """Free GPU mamba on a node. Returns num mamba tokens freed.""" + def _free_device_mamba(self, node: TreeNode) -> int: if node.mamba_value is None: return 0 mamba_num = len(node.mamba_value) @@ -584,14 +577,9 @@ def _free_gpu_mamba(self, node: TreeNode) -> int: return mamba_num def _evict_to_host(self, node: TreeNode) -> Tuple[int, int]: - """Evict a backuped device node to host: free GPU KV + GPU mamba. - - Node stays in the tree as evicted+backuped. - Caller must ensure node.backuped is True before calling. - Returns (full_num_evicted, mamba_num_evicted). - """ - assert not node.evicted, f"_evict_to_host on already-evicted node, {node.id=}" - assert node.backuped, f"_evict_to_host on non-backuped node, {node.id=}" + # 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) @@ -600,22 +588,18 @@ def _evict_to_host(self, node: TreeNode) -> Tuple[int, int]: if self.full_lru_list.in_list(node): self.full_lru_list.remove_node(node) - mamba_num = self._free_gpu_mamba(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_leaf(self, node: TreeNode) -> Tuple[int, int]: - """Evict a non-backuped device leaf: free GPU KV + mamba, delete from tree. - - Used in write_through when a node was never backed up. - Returns (full_num_evicted, mamba_num_evicted). - """ - assert not node.evicted, f"_evict_regular on already-evicted node, {node.id=}" - assert not node.backuped, f"_evict_regular on backuped node, {node.id=}" - assert len(node.children) == 0, f"_evict_regular on non-leaf, {node.id=}" + 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=}" num_full = len(node.value) @@ -624,7 +608,7 @@ def _evict_regular_leaf(self, node: TreeNode) -> Tuple[int, int]: if self.full_lru_list.in_list(node): self.full_lru_list.remove_node(node) - mamba_num = self._free_gpu_mamba(node) + mamba_num = self._free_device_mamba(node) if node.mamba_host_value is not None: if self.mamba_host_lru_list.in_list(node): @@ -644,42 +628,38 @@ def _evict_regular_leaf(self, node: TreeNode) -> Tuple[int, int]: self._iteratively_delete_tombstone_leaf(node) return num_full, mamba_num - def _evict_host_leaf_node(self, x: TreeNode) -> int: - """Evict a host-resident leaf: free host KV + host mamba, delete from tree, cascade. - - Returns num host KV tokens evicted. - """ - assert x.evicted, f"host leaf not evicted, {x.id=}" - assert x.backuped, f"host leaf not backuped, {x.id=}" - assert x.mamba_value is None, f"host leaf has GPU mamba, {x.id=}" + 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 ( - x.host_ref_counter == 0 - ), f"host leaf in use, {x.id=} {x.host_ref_counter=}" + node.host_ref_counter == 0 + ), f"in use, {node.id=} {node.host_ref_counter=}" - full_num_evicted = self.cache_controller.evict_host(x.host_value) - x.host_value = None + num_evicted = self.cache_controller.evict_host(node.host_value) + node.host_value = None - if x.mamba_host_value is not None: - if self.mamba_host_lru_list.in_list(x): - self.mamba_host_lru_list.remove_node(x) - self.mamba_pool_host.free(x.mamba_host_value) - x.mamba_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(x) - parent = x.parent - key = self.get_child_key_fn(x.key) + self._discard_from_leaf_sets(node) + parent = node.parent + key = self.get_child_key_fn(node.key) v = parent.children.pop(key, None) - assert v == x, f"parent does not have child key, {key}" + assert v == node, f"parent does not have child key, {key}" self._update_leaf_status(parent) - self._iteratively_delete_tombstone_leaf(x) + self._iteratively_delete_tombstone_leaf(node) - return full_num_evicted + return num_evicted 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 node.mamba_host_value is None, f"node has mamba host value, {node.id=}" + 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 = self.get_child_key_fn(node.key) @@ -697,9 +677,6 @@ def _delete_tombstone_leaf(self, node: TreeNode) -> None: 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 @@ -733,7 +710,7 @@ def _evict_device_leaf(self, x: TreeNode) -> Tuple[int, int]: - 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_leaf (delete from tree) + - not backuped + write_through: _evict_regular (delete from tree) """ if not x.backuped: if self.cache_controller.write_policy == "write_back": @@ -741,7 +718,7 @@ def _evict_device_leaf(self, x: TreeNode) -> Tuple[int, int]: self.writing_check(write_back=True) return self._evict_to_host(x) else: - return self._evict_regular_leaf(x) + return self._evict_regular(x) return self._evict_to_host(x) def evict(self, params: EvictParams) -> EvictResult: @@ -789,7 +766,7 @@ def evict_host(self, num_tokens: int): if x not in self.evictable_full_host_leaves: continue - num_evicted += self._evict_host_leaf_node(x) + 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)) @@ -813,11 +790,10 @@ def evict_mamba_host(self, num_mamba_hosts: int) -> int: continue if x in self.evictable_full_host_leaves: - # Host leaf: evict host KV + mamba, delete from tree - self._evict_host_leaf_node(x) + self._evict_host_leaf(x) num_evicted += 1 else: - # Internal host node: free host mamba only (tombstone) + # internal host node: free host mamba only (tombstone) self.mamba_host_lru_list.remove_node(x) self.mamba_pool_host.free(x.mamba_host_value) x.mamba_host_value = None @@ -847,7 +823,7 @@ def evict_mamba(self, mamba_num: int) -> int: ), f"evicted node should not be in mamba_lru_list, {x.id=}" if len(x.children) > 0: - # Internal: free GPU mamba only, KV stays on GPU (tombstone) + # 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_pool.free(x.mamba_value) @@ -871,12 +847,8 @@ def evict_mamba(self, mamba_num: int) -> int: 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.""" - assert node.evicted, f"_unevict_node on non-evicted node, {node.id=}" - # invariant: evicted => no GPU mamba - assert ( - node.mamba_value is None - ), f"evicted node should not have GPU mamba, {node.id=}" + 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() @@ -1752,7 +1724,7 @@ def write_backup_storage(self, node: TreeNode): if self.hicache_storage_pass_prefix_keys else None ) - extra_pools = self.archive_transfers(node) + extra_pools = self.mamba_archive_transfers(node) operation_id = self.cache_controller.write_storage( node.host_value, node.key, @@ -1795,7 +1767,7 @@ def prefetch_from_storage( return # Allocate host mamba slot - extra_pools = self.prefetch_prepare(new_input_tokens, last_hash) + 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) @@ -1853,7 +1825,6 @@ def check_prefetch_progress(self, req_id: str) -> bool: ) min_completed_tokens = completed_tokens_tensor.item() - # Resolve mamba transfer state mamba_host_indices = None mamba_loaded = False for transfer in operation.pool_transfers or []: @@ -1945,7 +1916,6 @@ def _insert_helper_host( if len(key): child_key = self.get_child_key_fn(key) - # Insert remaining tokens as a new leaf node leaf_node: Optional[TreeNode] = None if len(key): new_node = TreeNode() @@ -1992,7 +1962,6 @@ def _flush_pending_storage_backups_before_reset(self) -> None: self.writing_check(write_back=True) deadline = time.monotonic() + 30.0 - last_log_time = 0.0 while time.monotonic() < deadline: self.drain_storage_control_queues() backup_qsize = self.cache_controller.backup_queue.qsize() @@ -2037,8 +2006,10 @@ def _alloc_with_evict( raise RuntimeError(error_message) return indices - def backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: - """PoolTransfers for D→H backup: mamba device → mamba host.""" + # -- 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 [ @@ -2049,83 +2020,76 @@ def backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: ) ] - def backup_commit(self, node: TreeNode, transfers: list[PoolTransfer]) -> None: - """After D→H backup succeeds: store auto-allocated mamba host indices into node.""" + 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 - mamba_host = transfers[0].host_indices - if node.mamba_host_value is None and mamba_host is not None: - node.mamba_host_value = mamba_host + 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 self.mamba_host_lru_list.insert_mru(node) - def archive_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: - """PoolTransfers for H→Storage archive (write_backup_storage).""" - mamba_host_value = getattr(node, "mamba_host_value", None) - if mamba_host_value is None or not node.hash_value: + 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="mamba", - host_indices=mamba_host_value, + host_indices=node.mamba_host_value, keys=[node.hash_value[-1]], hit_policy="trailing_pages", ) ] - def prefetch_prepare( + def mamba_prefetch_alloc( self, token_ids: List[int], last_hash: Optional[str], ) -> Optional[list[PoolTransfer]]: - """Alloc mamba host slot and build PoolTransfers for Storage→H prefetch. - - Returns (host_handle, transfers) so the caller can track the allocation, - or None if allocation fails (caller should cancel the prefetch). - """ + # allocate a mamba host slot and build Storage→H transfer descriptor if not token_ids: return None - mamba_host_index = self._alloc_with_evict( + host_indices = self._alloc_with_evict( self.mamba_pool_host, 1, self.evict_mamba_host ) - if mamba_host_index is None: + if host_indices is None: return None - # Use a placeholder key; the I/O thread - # will replace it with the correct hash after query hit determination. - transfers = [ + # placeholder key; I/O thread replaces with correct hash after hit query + return [ PoolTransfer( name="mamba", - host_indices=mamba_host_index, + host_indices=host_indices, keys=["__placeholder__"], hit_policy="trailing_pages", ) ] - return transfers - def restore_transfers( + def mamba_restore_transfers( self, last_hit_node: TreeNode, nodes_to_restore: list[TreeNode], req, ) -> Optional[list[PoolTransfer]]: - """PoolTransfers for H→D restore (load_back).""" - mamba_host_list: list[torch.Tensor] = [] + # 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 - mamba_host_list.append(node.mamba_host_value) + backed_up_host_indices.append(node.mamba_host_value) transfers: list[PoolTransfer] = [] - # Transfer from host to node's device mamba space - if mamba_host_list: + if backed_up_host_indices: transfers.append( PoolTransfer( name="mamba", - host_indices=torch.cat(mamba_host_list), - device_indices=None, # will be allocated by cache_controller + host_indices=torch.cat(backed_up_host_indices), + device_indices=None, ) ) - # Transfer from host to request's device mamba space if ( req is not None and last_hit_node in nodes_to_restore @@ -2137,7 +2101,7 @@ def restore_transfers( len(last_hit_node.mamba_host_value), self.evict_mamba, lock_node=last_hit_node, - error_message="Can not alloc request mamba cache for host load back", + error_message="Cannot alloc request mamba cache for host load back", )[0] transfers.append( PoolTransfer( @@ -2149,17 +2113,17 @@ def restore_transfers( return transfers if transfers else None - def restore_commit( + def mamba_restore_commit( self, - extra_nodes: list[TreeNode], + restored_nodes: list[TreeNode], transfers: Optional[list[PoolTransfer]], ) -> None: - """After H→D restore succeeds: write back controller-allocated mamba device indices.""" - if not extra_nodes or not transfers or transfers[0].device_indices is 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 - mamba_device = transfers[0].device_indices + device_indices = transfers[0].device_indices offset = 0 - for n in extra_nodes: - n_len = len(n.mamba_host_value) - n.mamba_value = mamba_device[offset : offset + n_len].clone() - offset += n_len + for node in restored_nodes: + count = len(node.mamba_host_value) + node.mamba_value = device_indices[offset : offset + count].clone() + offset += count diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index a0328de58961..4b4fbf0ceff2 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1090,7 +1090,10 @@ def __init__( ): self.device_pool = device_pool self.page_size = 1 - assert layout in ["page_first_direct", "layer_first"], "Unsupported layout: {layout}" + assert layout in [ + "page_first_direct", + "layer_first", + ], "Unsupported layout: {layout}" self.layout = layout self.pin_memory = pin_memory self.device = device From a03ae570ab9c42c05a5f7282faacadceff760160 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Sun, 15 Mar 2026 15:32:51 +0000 Subject: [PATCH 08/39] Optimize code style --- .../srt/mem_cache/hi_mamba_radix_cache.py | 1118 ++++++++--------- 1 file changed, 559 insertions(+), 559 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 634a85cc659d..1564390eee5e 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -1265,595 +1265,190 @@ def dec_lock_ref(self, node: TreeNode): 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 - ) + def can_terminate_prefetch(self, operation: PrefetchOperation): + can_terminate = True - 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) + if self.prefetch_stop_policy == "best_effort": + return can_terminate - 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 + if len(operation.hash_value) == 0: + completed = False 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.", + completed = ( + operation.completed_tokens == len(operation.hash_value) * self.page_size ) - 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 - ) + 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 - 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 + operation_terminated = operation.is_terminated() + if self.tp_world_size > 1: + states = torch.tensor( + [1 - int(can_terminate), int(operation_terminated)], + dtype=torch.int, ) - 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}", + 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 - 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}" + def terminate_prefetch(self, req_id: str): + if req_id not in self.ongoing_prefetch: + return - 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." + _, _, _, operation = self.ongoing_prefetch[req_id] + if operation.host_indices is None: + return + operation.mark_terminate() - 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}" + def pop_prefetch_loaded_tokens(self, req_id: str) -> int: + return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) - self._drain_storage_control_queues_local() - self._force_release_pending_storage_ops() + 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, + ) + self.ongoing_backup[operation_id] = node + self._protect_host_node(node) - 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_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 - 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 == "mamba": - if transfer.host_indices is not None: - self.mamba_pool_host.free(transfer.host_indices) - break + self._protect_host_node(last_host_node) - def _force_release_pending_storage_ops(self): - cc = self.cache_controller + # 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: + self._release_host_node(last_host_node) + return - 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, 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.") + # 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) + return - 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, + logger.info( + "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", + req_id, + prefetch_length // self.page_size, + 1, ) + 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 _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: - 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 check_prefetch_progress(self, req_id: str) -> bool: + if req_id not in self.ongoing_prefetch: + return True - 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 + last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[ + req_id + ] - extra_config = tomllib.load(f) - elif ext in (".yaml", ".yml"): - import yaml + if operation.host_indices is None: + return True - 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 + if not self.can_terminate_prefetch(operation): + return False - 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 + completed_tokens, hash_value = self.cache_controller.terminate_prefetch( + operation ) - 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__}" + min_completed_tokens = completed_tokens + if self.tp_world_size > 1: + completed_tokens_tensor = torch.tensor( + min_completed_tokens, dtype=torch.int ) - 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__}" + torch.distributed.all_reduce( + completed_tokens_tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, ) + min_completed_tokens = completed_tokens_tensor.item() - 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 + mamba_host_indices = None + mamba_loaded = False + for transfer in operation.pool_transfers or []: + if transfer.name == "mamba": + mamba_host_indices = transfer.host_indices + mamba_loaded = ( + operation.pool_storage_result.extra_pool_hit_pages.get("mamba", 0) + >= 1 + ) + break - def drain_storage_control_queues(self): - cc = self.cache_controller + 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, + ) - 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 - ) - 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, - ) - 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) - - # 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: - self._release_host_node(last_host_node) - 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) - return - - logger.info( - "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", - req_id, - prefetch_length // self.page_size, - 1, - ) - 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 == "mamba": - mamba_host_indices = transfer.host_indices - mamba_loaded = ( - operation.pool_storage_result.extra_pool_hit_pages.get("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 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 @@ -2127,3 +1722,408 @@ def mamba_restore_commit( count = len(node.mamba_host_value) node.mamba_value = device_indices[offset : offset + count].clone() offset += count + + # ---- 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, + 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_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 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 == "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, 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, _, 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: + 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 + ) From 2836195a8873d2716284bf2f745ffc05e0ddbe3b Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 10:44:28 +0800 Subject: [PATCH 09/39] Optimize code style --- .../sglang/srt/managers/cache_controller.py | 2 +- python/sglang/srt/managers/schedule_policy.py | 3 +- .../4-gpu-models/test_qwen3_next_models.py | 104 +----------------- 3 files changed, 5 insertions(+), 104 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index d4c31e2d70f4..ff3376ca9031 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -971,7 +971,7 @@ def prefetch_thread_func(self): operation.host_indices[storage_hit_count:] ) operation.host_indices = operation.host_indices[:storage_hit_count] - logger.info( + logger.debug( f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}." ) self.prefetch_buffer.put(operation) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 683dd77438f9..57417e50e228 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -767,7 +767,8 @@ def add_one_req( if total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN - if req.host_hit_length > 0 or req.last_host_node is not req.last_node: + if req.host_hit_length > 0: + # TODO: also need to add a check for mamba_hit_length. new_indices, req.last_node = self.tree_cache.init_load_back( req.last_host_node, req.host_hit_length, req=req ) diff --git a/test/registered/4-gpu-models/test_qwen3_next_models.py b/test/registered/4-gpu-models/test_qwen3_next_models.py index ac18aebe865e..982cb6dcef85 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models.py @@ -1,5 +1,3 @@ -import shutil -import tempfile import unittest from types import SimpleNamespace @@ -22,28 +20,12 @@ register_cuda_ci(est_time=350, suite="stage-c-test-4-gpu-h100") QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" -QWEN3_NEXT_8B_MODEL = "/sgl-workspace/mnt/Qwen3.5-9B" + ACC_THRESHOLDS = { - QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.7}, - QWEN3_NEXT_8B_MODEL: {"kl_div": 0.0025, "gsm8k": 0.7}, + QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93}, } -def flush_cache(base_url, retries: int = 60, retry_interval: float = 1.0) -> bool: - """Flush device cache to force remote storage access.""" - import time - - for _ in range(retries): - try: - response = requests.post(f"{base_url}/flush_cache", timeout=10) - if response.status_code == 200: - return True - except requests.RequestException: - return False - time.sleep(retry_interval) - return False - - def send_request_helper(base_url: str, text: str): response = requests.post( base_url + "/generate", @@ -144,87 +126,5 @@ def test_prefix_cache_branching(self): print("test_prefix_cache_branching passed") -class TestQwen3NextWithHiCache(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = QWEN3_NEXT_8B_MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.storage_dir = tempfile.mkdtemp(prefix="qwen3-next-hicache-") - env = { - "FLASHINFER_DISABLE_VERSION_CHECK": "1", - "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.storage_dir, - } - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - env=env, - other_args=[ - "--tp", - "4", - "--mamba-scheduler-strategy", - "extra_buffer", - "--page-size", - "64", - "--max-total-tokens", - "131072", - "--max-mamba-cache-size", - "600", - "--hicache-mem-layout", - "page_first_direct", - "--enable-hierarchical-cache", - "--hicache-ratio", - "2", - "--hicache-size", - "0", - "--hicache-write-policy", - "write_through", - "--hicache-storage-backend", - "file", - "--hicache-storage-prefetch-policy", - "wait_complete", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - shutil.rmtree(cls.storage_dir, ignore_errors=True) - - def _run_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=100, - max_new_tokens=512, - parallel=10, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - return run_eval(args) - - def test_gsm8k(self): - first_metrics = self._run_gsm8k() - print(f"first_metrics={first_metrics}") - self.assertGreaterEqual( - first_metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) - - print(f"flush cache") - assert flush_cache(self.base_url) - - second_metrics = self._run_gsm8k() - print(f"second_metrics={second_metrics}") - self.assertGreaterEqual( - second_metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) - self.assertLessEqual( - abs(second_metrics["accuracy"] - first_metrics["accuracy"]), - 0.05, - f"HiCache prefetch accuracy drift too large: " - f"first={first_metrics['accuracy']}, second={second_metrics['accuracy']}", - ) - - if __name__ == "__main__": unittest.main() From 91f954ffaa6d51f3e441731701a6c065f130b2e1 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 11:01:31 +0800 Subject: [PATCH 10/39] Remove pp filter --- .../model_runner_kv_cache_mixin.py | 24 ++----------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index d47af942f796..c46ff0a388c7 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -1,6 +1,5 @@ from __future__ import annotations -import dataclasses import logging from typing import TYPE_CHECKING @@ -44,19 +43,6 @@ _is_npu = is_npu() -def _filter_linear_cache_params_for_pp(model_runner: "ModelRunner", cache_params): - if cache_params is None: - return None - if model_runner.pp_size <= 1: - return cache_params - local_layers = [ - layer_id - for layer_id in cache_params.layers - if model_runner.start_layer <= layer_id < model_runner.end_layer - ] - return dataclasses.replace(cache_params, layers=local_layers) - - class ModelRunnerKVCacheMixin: def get_cell_size_per_token(self: ModelRunner, num_layers: int) -> int: kv_size = torch._utils._element_size(self.kv_cache_dtype) @@ -381,16 +367,13 @@ def _init_pools(self: ModelRunner, max_num_reqs: int): # if max_num_reqs <= 32, we pre-allocate 2x requests pre_alloc_size = max_num_reqs * 2 if max_num_reqs <= 32 else 0 if config := self.mambaish_config: - cache_params = _filter_linear_cache_params_for_pp( - self, config.mamba2_cache_params - ) self.req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, - cache_params=cache_params, + cache_params=config.mamba2_cache_params, speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), pre_alloc_size=pre_alloc_size, @@ -407,9 +390,6 @@ def _init_pools(self: ModelRunner, max_num_reqs: int): pre_alloc_size=pre_alloc_size, ) elif config := self.mambaish_config: - cache_params = _filter_linear_cache_params_for_pp( - self, config.mamba2_cache_params - ) self.req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=self.server_args.max_mamba_cache_size, @@ -418,7 +398,7 @@ def _init_pools(self: ModelRunner, max_num_reqs: int): + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, - cache_params=cache_params, + cache_params=config.mamba2_cache_params, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, enable_overlap_schedule=not self.server_args.disable_overlap_schedule, From 520ca17c660cab64f5786ff7b932204a2641176f Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 06:48:08 +0000 Subject: [PATCH 11/39] Optimize layerwise waiting logic --- .../srt/mem_cache/hi_mamba_radix_cache.py | 36 ++++++++---------- .../hybrid_cache/hybrid_cache_controller.py | 15 ++++---- python/sglang/srt/mem_cache/memory_pool.py | 37 +++++++------------ 3 files changed, 36 insertions(+), 52 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 1564390eee5e..f5212ed3285b 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -114,16 +114,6 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): "HiMambaRadixCache requires HybridReqToTokenPool for hybrid SSM models." ) - self.hybrid_model_layer_ids: list[int] = [] - full_layer_ids = sorted( - self.hybrid_kv_cache.full_attention_layer_id_mapping.keys() - ) - mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) - self.hybrid_model_layer_ids = sorted(set(full_layer_ids) | set(mamba_layer_ids)) - self.transfer_layer_num = len(self.hybrid_model_layer_ids) - self.hybrid_kv_cache.set_model_layer_id_mapping(self.hybrid_model_layer_ids) - params.req_to_token_pool.set_model_layer_id_mapping(self.hybrid_model_layer_ids) - self.kvcache = self.hybrid_kv_cache.full_kv_pool kv_host_pool_cls = ( MLATokenToKVPoolHost @@ -146,22 +136,25 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): layout=server_args.hicache_mem_layout, ) + full_layer_ids = sorted( + self.hybrid_kv_cache.full_attention_layer_id_mapping.keys() + ) + mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) + self.transfer_layer_num = len(set(full_layer_ids) | set(mamba_layer_ids)) + full_layer_mapping = dict(self.hybrid_kv_cache.full_attention_layer_id_mapping) mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) + transfer_layer_num = self.transfer_layer_num - def kv_layer_mapper(model_layer_local_id: int) -> Optional[int]: - if not 0 <= model_layer_local_id < len(self.hybrid_model_layer_ids): + def kv_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: return None - return full_layer_mapping.get( - self.hybrid_model_layer_ids[model_layer_local_id] - ) + return full_layer_mapping.get(layer_id) - def mamba_layer_mapper(model_layer_local_id: int) -> Optional[int]: - if not 0 <= model_layer_local_id < len(self.hybrid_model_layer_ids): + def mamba_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: return None - return mamba_layer_mapping.get( - self.hybrid_model_layer_ids[model_layer_local_id] - ) + return mamba_layer_mapping.get(layer_id) self.host_pool_group = HostPoolGroup( [ @@ -226,6 +219,9 @@ def mamba_layer_mapper(model_layer_local_id: int) -> Optional[int]: params.req_to_token_pool.register_layer_transfer_counter( self.cache_controller.layer_done_counter ) + self.hybrid_kv_cache.register_layer_transfer_counter( + self.cache_controller.layer_done_counter + ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, prefetch_threshold=prefetch_threshold, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 82526085693c..4725accb1fd5 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -163,7 +163,6 @@ def __init__( pp_size: int = 1, transfer_layer_num: Optional[int] = None, ): - self.transfer_layer_num = transfer_layer_num startup_storage_backend = storage_backend super().__init__( token_to_kv_pool_allocator=token_to_kv_pool_allocator, @@ -180,12 +179,12 @@ def __init__( pp_rank=pp_rank, pp_size=pp_size, ) - self.transfer_layer_num = self.transfer_layer_num or self.layer_num - if self.transfer_layer_num != self.layer_done_counter.num_layers: - self.layer_done_counter = LayerDoneCounter(self.transfer_layer_num) - self.mem_pool_device.register_layer_transfer_counter( - self.layer_done_counter - ) + # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)), + # not just the full attention layers reported by full_kv_pool. + if transfer_layer_num is not None and transfer_layer_num != self.layer_num: + self.layer_num = transfer_layer_num + self.layer_done_counter = LayerDoneCounter(self.layer_num) + if startup_storage_backend is not None: self.attach_storage_backend( storage_backend=startup_storage_backend, @@ -313,7 +312,7 @@ def start_loading(self) -> int: producer_event.start_event.record() with device_module.stream(self.load_stream): producer_event.start_event.wait(self.load_stream) - for i in range(self.transfer_layer_num): + for i in range(self.layer_num): self.mem_pool_host.load_to_device_per_layer( self.mem_pool_device, host_indices, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 4298cca4a074..176389996b36 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -469,8 +469,9 @@ def __init__( self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver + # TODO: Support PP + self.start_layer = 0 self.layer_transfer_counter = None - self.model_layer_id_to_local: dict[int, int] = {} self._init_mamba_pool( size=mamba_size, mamba_spec_state_size=mamba_spec_state_size, @@ -517,11 +518,6 @@ def register_layer_transfer_counter( ): self.layer_transfer_counter = layer_transfer_counter - def set_model_layer_id_mapping(self, model_layer_ids: List[int]): - self.model_layer_id_to_local = { - layer_id: i for i, layer_id in enumerate(model_layer_ids) - } - # For chunk prefill req, we do not need to allocate mamba cache, # We could use allocated mamba cache instead. def alloc(self, reqs: List["Req"]) -> Optional[List[int]]: @@ -577,8 +573,7 @@ def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor: def mamba2_layer_cache(self, layer_id: int): assert layer_id in self.mamba_map if self.layer_transfer_counter is not None: - wait_id = self.model_layer_id_to_local.get(layer_id, layer_id) - self.layer_transfer_counter.wait_until(wait_id) + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) return self.mamba_pool.mamba2_layer_cache(self.mamba_map[layer_id]) def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState: @@ -1247,10 +1242,8 @@ def __init__( self.device = device self.full_layer_nums = len(full_attention_layer_ids) self.page_size = page_size - # TODO support pp? - self.start_layer = 0 + self.start_layer = 0 # TODO: Support PP self.layer_transfer_counter = None - self.model_layer_id_to_local: dict[int, int] = {} self.head_num = head_num self.head_dim = head_dim self.mamba_pool = mamba_pool @@ -1338,30 +1331,26 @@ def register_layer_transfer_counter( self, layer_transfer_counter: "LayerDoneCounter" ): self.layer_transfer_counter = layer_transfer_counter + # The layer-wise wait logic is executed at the Hybrid LinearPool level; + # no additional wait is needed in the full_kv_pool + self.full_kv_pool.register_layer_transfer_counter(None) - def set_model_layer_id_mapping(self, model_layer_ids: List[int]): - self.model_layer_id_to_local = { - layer_id: i for i, layer_id in enumerate(model_layer_ids) - } - - def _wait_for_model_layer(self, layer_id: int): - if self.layer_transfer_counter is None: - return - wait_id = self.model_layer_id_to_local.get(layer_id, layer_id) - self.layer_transfer_counter.wait_until(wait_id) + def _wait_for_layer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) def get_key_buffer(self, layer_id: int): - self._wait_for_model_layer(layer_id) + self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_key_buffer(layer_id) def get_value_buffer(self, layer_id: int): - self._wait_for_model_layer(layer_id) + self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_value_buffer(layer_id) def get_kv_buffer(self, layer_id: int): - self._wait_for_model_layer(layer_id) + self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_kv_buffer(layer_id) From 3cd247da711740ae1c2eb626f034bd4846a83822 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 08:07:58 +0000 Subject: [PATCH 12/39] Support loading mamba state only --- python/sglang/srt/managers/schedule_batch.py | 10 +++- python/sglang/srt/managers/schedule_policy.py | 3 +- python/sglang/srt/managers/scheduler.py | 6 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 4 +- .../srt/mem_cache/hi_mamba_radix_cache.py | 57 ++++++++++--------- .../hybrid_cache/hybrid_cache_controller.py | 14 +++-- .../sglang/srt/mem_cache/memory_pool_host.py | 2 +- 7 files changed, 51 insertions(+), 45 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5597d8baea43..9c35f5c4f76d 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -639,8 +639,8 @@ def __init__( 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 + self.mamba_host_hit_length = 0 # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 # The node to lock until for swa radix tree lock ref @@ -816,6 +816,10 @@ def output_ids_through_stop(self) -> List[int]: return self.output_ids[: self.finished_len] return self.output_ids + @property + def should_init_load_back_host(self) -> bool: + return self.host_hit_length > 0 or self.mamba_host_hit_length > 0 + def pop_committed_kv_cache(self) -> int: """Return the length of committed KV cache and mark them as freed.""" assert ( @@ -891,15 +895,15 @@ def init_next_round_input( self.prefix_indices, self.last_node, self.last_host_node, - self.last_host_backup_node, self.host_hit_length, + self.mamba_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_host_hit_length, match_result.mamba_branching_seqlen, ) if match_result.cache_protected_len is not None: diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 57417e50e228..526eb2836e73 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -767,8 +767,7 @@ def add_one_req( if total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN - if req.host_hit_length > 0: - # TODO: also need to add a check for mamba_hit_length. + if req.should_init_load_back_host: new_indices, req.last_node = self.tree_cache.init_load_back( req.last_host_node, req.host_hit_length, req=req ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ec3a254adf98..4beafa0c27a6 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1715,11 +1715,7 @@ def handle_batch_generate_request( def _prefetch_kvcache(self, req: Req): if self.enable_hicache_storage: 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 - ) + last_host_node = 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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index e5c251a836e0..040162dcb452 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -97,9 +97,9 @@ 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_host_hit_length: Length of the mamba host hit mamba_branching_seqlen: The mamba radix cache branching point, which is the longest page-aligned position that could've been cache hit if there exists a mamba state. @@ -108,8 +108,8 @@ 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_host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: 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 index f5212ed3285b..08e15369d4f2 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -345,18 +345,25 @@ def load_back( 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 + 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.info( + 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 ) @@ -366,7 +373,9 @@ def load_back( extra_pools=mamba_pools, ) if full_device_indices is None: - self.evict(EvictParams(num_tokens=len(full_host_indices))) + 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 ) @@ -413,9 +422,7 @@ def init_load_back( mem_quota: Optional[int] = None, req=None, ): - if last_node.evicted or ( - last_node.mamba_value is None and last_node.mamba_host_value is not None - ): + 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( @@ -424,11 +431,7 @@ def init_load_back( return loading_values, last_node while last_node is not self.root_node and ( - last_node.evicted - or ( - last_node.mamba_value is None - and last_node.mamba_host_value is not None - ) + last_node.evicted or last_node.mamba_evicted ): last_node = last_node.parent @@ -966,7 +969,6 @@ def match_prefix(self, params: MatchPrefixParams) -> 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, ) @@ -974,17 +976,13 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: 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 - ) + 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, TreeNode]: - """Walk tree to find best_last_node (mamba boundary) and deepest node.""" + ) -> Tuple[List[torch.Tensor], TreeNode, int]: + """Walk tree to find best_last_node (mamba boundary).""" node = self.root_node child_key = self.get_child_key_fn(key) @@ -1021,8 +1019,7 @@ def _match_prefix_helper( best_value_len = len(value) best_last_node = node - deepest_node = best_last_node - return value, best_last_node, best_value_len, deepest_node + return value, best_last_node, best_value_len def _match_post_processor( self, @@ -1030,7 +1027,6 @@ def _match_post_processor( value: List[torch.Tensor], best_last_node: TreeNode, best_value_len: int, - deepest_node: TreeNode, ) -> MatchResult: cow_mamba = params.cow_mamba req = params.req @@ -1074,6 +1070,11 @@ def _match_post_processor( while last_host_node is not self.root_node and not last_host_node.backuped: last_host_node = last_host_node.parent + # Calculate mamba host hit length + mamba_host_hit_length = 0 + if last_host_node.mamba_evicted and last_host_node.mamba_backuped: + mamba_host_hit_length = len(last_host_node.mamba_host_value) + mamba_node = best_last_node if cow_mamba and mamba_node.mamba_value is not None: if req.mamba_pool_idx is None: @@ -1102,8 +1103,8 @@ def _match_post_processor( device_indices=value, last_device_node=last_device_node, last_host_node=last_host_node, - last_host_backup_node=last_host_node, host_hit_length=host_hit_length, + mamba_host_hit_length=mamba_host_hit_length, mamba_branching_seqlen=mamba_branching_seqlen, ) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 4725accb1fd5..9f85416ecd65 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -280,14 +280,20 @@ def load( node_id: int = -1, extra_pools: Optional[list[PoolTransfer]] = None, ) -> Optional[torch.Tensor]: - device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) - if device_indices is None: - return None + need_load_kv = host_indices.numel() > 0 + if need_load_kv: + device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) + if device_indices is None: + return None + else: + device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) + pool_transfers = self._resolve_pool_transfers_allocation( extra_pools, alloc_host=False ) if pool_transfers is None and extra_pools: - self.mem_pool_device_allocator.free(device_indices) + if need_load_kv: + self.mem_pool_device_allocator.free(device_indices) return None self.load_queue.append( diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 4b4fbf0ceff2..77fef235fd46 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1427,7 +1427,7 @@ def load_to_device_per_layer( # 1. Anchor (KV) transfer anchor = self.anchor_entry local_layer_id = anchor.layer_mapper(layer_id) - if local_layer_id is not None: + if local_layer_id is not None and host_indices.numel() > 0: anchor.host_pool.load_to_device_per_layer( anchor.device_pool, host_indices, From 77163eb92f28c2e5a40f1f80e8a29e4b220c264a Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 18:02:03 +0800 Subject: [PATCH 13/39] Using enum to replace str --- .../srt/mem_cache/hi_mamba_radix_cache.py | 24 ++++----- .../sglang/srt/mem_cache/hicache_storage.py | 49 ++++++++++++++----- 2 files changed, 49 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 08e15369d4f2..803705d57430 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -18,7 +18,7 @@ MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolTransfer +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, PrefetchOperation, @@ -159,14 +159,14 @@ def mamba_layer_mapper(layer_id: int) -> Optional[int]: self.host_pool_group = HostPoolGroup( [ PoolEntry( - name="kv", + name=PoolName.KV, host_pool=self.full_kv_pool_host, device_pool=self.kvcache, layer_mapper=kv_layer_mapper, is_primary_index_anchor=True, ), PoolEntry( - name="mamba", + name=PoolName.MAMBA, host_pool=self.mamba_pool_host, device_pool=params.req_to_token_pool.mamba_pool, layer_mapper=mamba_layer_mapper, @@ -1420,10 +1420,12 @@ def check_prefetch_progress(self, req_id: str) -> bool: mamba_host_indices = None mamba_loaded = False for transfer in operation.pool_transfers or []: - if transfer.name == "mamba": + if transfer.name == PoolName.MAMBA: mamba_host_indices = transfer.host_indices mamba_loaded = ( - operation.pool_storage_result.extra_pool_hit_pages.get("mamba", 0) + operation.pool_storage_result.extra_pool_hit_pages.get( + PoolName.MAMBA, 0 + ) >= 1 ) break @@ -1606,7 +1608,7 @@ def mamba_backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]] return None return [ PoolTransfer( - name="mamba", + name=PoolName.MAMBA, host_indices=node.mamba_host_value, device_indices=node.mamba_value, ) @@ -1629,7 +1631,7 @@ def mamba_archive_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer] return None return [ PoolTransfer( - name="mamba", + name=PoolName.MAMBA, host_indices=node.mamba_host_value, keys=[node.hash_value[-1]], hit_policy="trailing_pages", @@ -1652,7 +1654,7 @@ def mamba_prefetch_alloc( # placeholder key; I/O thread replaces with correct hash after hit query return [ PoolTransfer( - name="mamba", + name=PoolName.MAMBA, host_indices=host_indices, keys=["__placeholder__"], hit_policy="trailing_pages", @@ -1676,7 +1678,7 @@ def mamba_restore_transfers( if backed_up_host_indices: transfers.append( PoolTransfer( - name="mamba", + name=PoolName.MAMBA, host_indices=torch.cat(backed_up_host_indices), device_indices=None, ) @@ -1697,7 +1699,7 @@ def mamba_restore_transfers( )[0] transfers.append( PoolTransfer( - name="mamba", + name=PoolName.MAMBA, host_indices=last_hit_node.mamba_host_value, device_indices=req.mamba_pool_idx.unsqueeze(0), ) @@ -1888,7 +1890,7 @@ def detach_storage_backend(self) -> tuple: 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 == "mamba": + if transfer.name == PoolName.MAMBA: if transfer.host_indices is not None: self.mamba_pool_host.free(transfer.host_indices) break diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index e04d7938db0c..fcee37bd64c9 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -3,6 +3,7 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass +from enum import Enum from typing import Any, List, Literal, Optional, Set import torch @@ -64,6 +65,13 @@ class HiCacheStorageExtraInfo: extra_info: Optional[dict] = None +class PoolName(str, Enum): + """Well-known pool names used as PoolTransfer/PoolEntry identifiers.""" + + KV = "kv" + MAMBA = "mamba" + + @dataclass class PoolTransfer: """Unified per-pool transfer descriptor for batch v2 interface. @@ -126,10 +134,30 @@ def batch_exists_v2( pool_transfers: Optional[List[PoolTransfer]] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> PoolTransferResult: - """Check which KV pages (and associated extra pools) exist in storage. - - Returns a PoolTransferResult with the longest contiguous KV prefix hit and - per-pool hit counts, constrained by each pool's auxiliary_constraints policy. + """Check which cache pages exist in storage, respecting per-pool hit policies. + + Longest-prefix semantics + Extra-pool hit policies (``PoolTransfer.hit_policy``) + ------------------------------------------------------ + Each ``PoolTransfer`` in ``pool_transfers`` describes a secondary + cache pool (e.g. Mamba SSM states) that must be co-present with the + KV pages. The final ``final_pages`` is the minimum across all pools, + so a missing auxiliary page shrinks the usable prefix. + + - ``"all_pages"`` (default): every page in [0, kv_hit) must exist + for this pool. Used for pools that are required for every token + in the prefix (e.g. DeepSeek DSA pool). + + - ``"trailing_pages"``: only the *last* ``len(transfer.keys)`` pages + of the KV prefix need to exist. Used for pools whose data covers + only the tail of a prefix (e.g. Mamba/SWA Pool). + + Returns + ------- + PoolTransferResult + ``kv_hit_pages`` = length of the usable KV prefix. + ``extra_pool_hit_pages`` maps each pool name to the number of pages + that were found. """ raise NotImplementedError() @@ -290,7 +318,7 @@ def _get_suffixed_key(self, key: str) -> str: return key + self.config_suffix def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str: - if component_name is None or component_name in ("__default__", "kv"): + if component_name is None or component_name in ("__default__", PoolName.KV): return self._get_suffixed_key(key) return self._get_suffixed_key(f"{key}.{component_name}") @@ -393,11 +421,6 @@ def batch_exists_v2( pool_transfers: Optional[List[PoolTransfer]] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> PoolTransferResult: - """Check which KV pages (and associated extra pools) exist in storage. - - Returns a PoolTransferResult with the longest contiguous KV prefix hit and - per-pool hit counts, constrained by each pool's auxiliary_constraints policy. - """ existing_files = self._collect_existing_component_keys(keys, pool_transfers) def has_component(page_idx: int, name: str) -> bool: @@ -415,7 +438,7 @@ def has_component(page_idx: int, name: str) -> bool: len(keys), ) - hit_count: dict[str, int] = {"kv": kv_pages} if kv_pages else {} + hit_count: dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {} final_pages = kv_pages for transfer in pool_transfers or []: @@ -453,7 +476,7 @@ def has_component(page_idx: int, name: str) -> bool: return PoolTransferResult(final_pages, hit_count) def _log_key(self, pool_name: str, key: str) -> str: - return key if pool_name == "kv" else f"{key}.{pool_name}" + return key if pool_name == PoolName.KV else f"{key}.{pool_name}" def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool: """Read one page from storage into host_pool at page_offset.""" @@ -508,7 +531,7 @@ def _write_page( ) with open(component_path, "wb", buffering=0) as f: f.write(data_bytes) - if pool_name != "kv": + if pool_name != PoolName.KV: logger.info( "HiCacheFile wrote auxiliary component %s for key %s", pool_name, From df5867cb83f4bede184279268689d8bbb59c9081 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 16 Mar 2026 18:53:34 +0800 Subject: [PATCH 14/39] support MambaPoolHost layout --- .../sglang/srt/mem_cache/memory_pool_host.py | 179 ++++++++++++++---- 1 file changed, 146 insertions(+), 33 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 77fef235fd46..065f4f274b54 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1091,14 +1091,17 @@ def __init__( self.device_pool = device_pool self.page_size = 1 assert layout in [ + "page_first", "page_first_direct", "layer_first", ], "Unsupported layout: {layout}" + self.layout = layout self.pin_memory = pin_memory self.device = device self.allocator = get_allocator_from_storage(allocator_type) self.num_mamba_layers = device_pool.num_mamba_layers + self.conv_state_shapes = [ conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv ] @@ -1107,10 +1110,12 @@ def __init__( self.temporal_dtype = device_pool.mamba_cache.temporal.dtype self.dtype = self.conv_dtype self.size_per_token = self.get_size_per_token() + if host_size > 0: self.size = int(host_size * 1e9 // self.size_per_token) else: self.size = int(device_pool.size * host_to_device_ratio) + self.page_num = self.size // self.page_size + 1 self.size = self.page_num * self.page_size @@ -1142,11 +1147,12 @@ def __init__( def init_kv_buffer(self): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] - if self.layout == "page_first_direct": - # page-first: (size, num_layers, *shape) — per-page data is contiguous + if self.layout in ["page_first", "page_first_direct"]: + # page-first: (page_num, num_layers, 1, *shape) — per-page data is contiguous temporal_dims = ( self.size, self.num_mamba_layers, + 1, ) + self.temporal_state_shape self.temporal_buffer = alloc_func( temporal_dims, @@ -1157,7 +1163,7 @@ def init_kv_buffer(self): ) self.conv_buffer = [] for conv_shape in self.conv_state_shapes: - conv_dims = (self.size, self.num_mamba_layers) + conv_shape + conv_dims = (self.size, self.num_mamba_layers, 1) + conv_shape self.conv_buffer.append( alloc_func( conv_dims, @@ -1193,18 +1199,8 @@ def init_kv_buffer(self): ) ) - def _get_temporal_view(self, layer_id: int) -> torch.Tensor: - if self.layout == "page_first_direct": - return self.temporal_buffer[:, layer_id] - return self.temporal_buffer[layer_id] - - def _get_conv_view(self, layer_id: int, conv_idx: int) -> torch.Tensor: - if self.layout == "page_first_direct": - return self.conv_buffer[conv_idx][:, layer_id] - return self.conv_buffer[conv_idx][layer_id] - def _iter_page_tensors(self, index: int): - if self.layout == "page_first_direct": + if self.layout in ["page_first", "page_first_direct"]: yield self.temporal_buffer[index] for conv_buf in self.conv_buffer: yield conv_buf[index] @@ -1290,6 +1286,81 @@ def _copy_tensor( else: raise ValueError(f"Unsupported io_backend: {io_backend}") + @staticmethod + def _copy_tensor_pf_lf( + src: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + layer_id: int, + num_layers: int, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + item_size = MambaPoolHost._item_size_per_index(dst) + transfer_kv_per_layer_mla_pf_lf( + src=src, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + layer_id=layer_id, + item_size=item_size, + src_layout_dim=item_size * num_layers, + ) + elif io_backend == "direct": + transfer_kv_per_layer_direct_pf_lf( + src_ptrs=[src], + dst_ptrs=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + layer_id=layer_id, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + @staticmethod + def _copy_tensor_all_layers_lf_pf( + src_layers: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_layers: int, + device: str, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + item_size = MambaPoolHost._item_size_per_index(src_layers[0]) + src_ptrs = torch.tensor( + [src_layers[i].data_ptr() for i in range(num_layers)], + dtype=torch.uint64, + device=device, + ) + transfer_kv_all_layer_mla_lf_pf( + src_layers=src_ptrs, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=item_size, + dst_layout_dim=item_size * num_layers, + num_layers=num_layers, + ) + elif io_backend == "direct": + src_ptrs = [src_layers[i] for i in range(num_layers)] + transfer_kv_all_layer_direct_lf_pf( + src_ptrs=src_ptrs, + dst_ptrs=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + def load_to_device_per_layer( self, device_pool, @@ -1298,41 +1369,83 @@ def load_to_device_per_layer( layer_id, io_backend="kernel", ): - self._copy_tensor( - self._get_temporal_view(layer_id), - device_pool.mamba_cache.temporal[layer_id], - host_indices, - device_indices, - io_backend, - ) - for conv_idx in range(len(self.conv_state_shapes)): + if self.layout in ["page_first", "page_first_direct"]: + self._copy_tensor_pf_lf( + src=self.temporal_buffer, + dst=device_pool.mamba_cache.temporal[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_pf_lf( + src=self.conv_buffer[conv_idx], + dst=device_pool.mamba_cache.conv[conv_idx][layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + else: self._copy_tensor( - self._get_conv_view(layer_id, conv_idx), - device_pool.mamba_cache.conv[conv_idx][layer_id], + self.temporal_buffer[layer_id], + device_pool.mamba_cache.temporal[layer_id], host_indices, device_indices, io_backend, ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + self.conv_buffer[conv_idx][layer_id], + device_pool.mamba_cache.conv[conv_idx][layer_id], + host_indices, + device_indices, + io_backend, + ) def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend="kernel" ): - for layer_id in range(self.num_mamba_layers): - self._copy_tensor( - device_pool.mamba_cache.temporal[layer_id], - self._get_temporal_view(layer_id), - device_indices, - host_indices, - io_backend, + if self.layout in ["page_first", "page_first_direct"]: + self._copy_tensor_all_layers_lf_pf( + src_layers=device_pool.mamba_cache.temporal, + dst=self.temporal_buffer, + src_indices=device_indices, + dst_indices=host_indices, + num_layers=self.num_mamba_layers, + device=self.device_pool.device, + io_backend=io_backend, ) for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_all_layers_lf_pf( + src_layers=device_pool.mamba_cache.conv[conv_idx], + dst=self.conv_buffer[conv_idx], + src_indices=device_indices, + dst_indices=host_indices, + num_layers=self.num_mamba_layers, + device=self.device_pool.device, + io_backend=io_backend, + ) + else: + for layer_id in range(self.num_mamba_layers): self._copy_tensor( - device_pool.mamba_cache.conv[conv_idx][layer_id], - self._get_conv_view(layer_id, conv_idx), + device_pool.mamba_cache.temporal[layer_id], + self.temporal_buffer[layer_id], device_indices, host_indices, io_backend, ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + device_pool.mamba_cache.conv[conv_idx][layer_id], + self.conv_buffer[conv_idx][layer_id], + device_indices, + host_indices, + io_backend, + ) def get_data_page(self, index, flat: bool = True) -> torch.Tensor: data_page = torch.cat( From 0f2d2519c2d49d89618b346c98483a36124c0414 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 19:48:18 +0800 Subject: [PATCH 15/39] rollback format --- .../srt/mem_cache/hi_mamba_radix_cache.py | 1526 ++++++++--------- 1 file changed, 763 insertions(+), 763 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 803705d57430..85d245378117 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -1262,867 +1262,867 @@ def dec_lock_ref(self, node: TreeNode): node = node.parent return delta - def can_terminate_prefetch(self, operation: PrefetchOperation): - can_terminate = True + # ---- L3 Support ---- - if self.prefetch_stop_policy == "best_effort": - return can_terminate + def shutdown(self): + try: + if self.enable_storage: + self.detach_storage_backend() + except Exception: + logger.exception("Failed to detach storage backend on process shutdown.") - if len(operation.hash_value) == 0: - completed = False - else: - completed = ( - operation.completed_tokens == len(operation.hash_value) * self.page_size - ) + 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 + ) - 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) + 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: - return True + self.storage_metrics_collector = None - 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, + 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.", ) - 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 + 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 + ) - _, _, _, operation = self.ongoing_prefetch[req_id] - if operation.host_indices is None: - return - operation.mark_terminate() + 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}", + ) - def pop_prefetch_loaded_tokens(self, req_id: str) -> int: - return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) + 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}" - 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, + 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, ) - self.ongoing_backup[operation_id] = node - self._protect_host_node(node) + return True, "Attached HiCache storage backend successfully." - 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 + 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._protect_host_node(last_host_node) + self._drain_storage_control_queues_local() + self._force_release_pending_storage_ops() - # 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: - self._release_host_node(last_host_node) - return + 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." - # 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) - return + 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 - logger.info( - "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", - req_id, - prefetch_length // self.page_size, - 1, - ) - 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 - ) + def _force_release_pending_storage_ops(self): + cc = self.cache_controller - 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() + 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.") - 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 + 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 ) - >= 1 - ) - break + self.ongoing_backup.pop(ack_id, None) + except Exception: + logger.exception("Force release pending backup ops failed.") - 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, + 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, ) - # 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] - ) + 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 - # 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) + 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 - self._release_host_node(last_host_node) - del self.ongoing_prefetch[req_id] - self.cache_controller.prefetch_tokens_occupied -= len(token_ids) + 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 - loaded_from_storage = min_completed_tokens - matched_length - self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage + 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 + ) - 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.info( - "HiCache mamba prefetch completed for request %s: prefetched_tokens=%s mamba_states=%s", - req_id, - loaded_from_storage, - int(mamba_loaded), - ) + 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) - return True + _drain_revoke() + _drain_backup() + _drain_release() - def _insert_helper_host( - self, - node: TreeNode, - key: RadixKey, - host_value, - hash_value, - mamba_host_value: Optional[torch.Tensor] = None, - mamba_loaded: bool = False, + def _parse_storage_backend_extra_config( + self, storage_backend_extra_config: Optional[str] ): - 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 - 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) + 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 - key = key[prefix_len:] - host_value = host_value[prefix_len:] - hash_value = hash_value[prefix_len // self.page_size :] - matched_length += prefix_len + extra_config = tomllib.load(f) + elif ext in (".yaml", ".yml"): + import yaml - if prefix_len < len(node.key): - new_node = self._split_node(node.key, node, prefix_len) - node = new_node + 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 - if len(key): - child_key = self.get_child_key_fn(key) + 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 + ) - 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) - - # 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 + 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__}" + ) - def release_aborted_request(self, rid: str): - self.prefetch_loaded_tokens_by_reqid.pop(rid, None) + return ( + extra_config, + prefetch_threshold, + float(prefetch_timeout_base), + float(prefetch_timeout_per_ki_token), + hicache_storage_pass_prefix_keys, + ) - if rid not in self.ongoing_prefetch: - return + 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 - last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[rid] - if operation.host_indices is None: - return + def drain_storage_control_queues(self): + cc = self.cache_controller - completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) + 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.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 + torch.distributed.all_reduce( + qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + ) - 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) + 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, + ) - 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 _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 _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 + def can_terminate_prefetch(self, operation: PrefetchOperation): + can_terminate = True - # -- mamba PoolTransfer builders (D↔H↔S) ---------------------------------- + if self.prefetch_stop_policy == "best_effort": + return can_terminate - 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, + if len(operation.hash_value) == 0: + completed = False + else: + completed = ( + operation.completed_tokens == len(operation.hash_value) * self.page_size ) - ] - 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 - self.mamba_host_lru_list.insert_mru(node) + 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 - 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="trailing_pages", + operation_terminated = operation.is_terminated() + if self.tp_world_size > 1: + states = torch.tensor( + [1 - int(can_terminate), int(operation_terminated)], + dtype=torch.int, ) - ] - - 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="trailing_pages", + 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 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) + def terminate_prefetch(self, req_id: str): + if req_id not in self.ongoing_prefetch: + return - 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, - ) - ) + _, _, _, operation = self.ongoing_prefetch[req_id] + if operation.host_indices is None: + return + operation.mark_terminate() - 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_pool, - 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), - ) - ) + def pop_prefetch_loaded_tokens(self, req_id: str) -> int: + return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) - return transfers if transfers else None + 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, + ) + self.ongoing_backup[operation_id] = node + self._protect_host_node(node) - def mamba_restore_commit( + def prefetch_from_storage( 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: + 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 - 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 - - # ---- 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.") + self._protect_host_node(last_host_node) - 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 + # 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: + self._release_host_node(last_host_node) + return - 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) + # 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) + return - 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 + logger.info( + "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", + req_id, + prefetch_length // self.page_size, + 1, + ) + 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 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}.", - ) + def check_prefetch_progress(self, req_id: str) -> bool: + if req_id not in self.ongoing_prefetch: + return True - 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}.", - ) + last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[ + req_id + ] - 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 operation.host_indices is None: + return True - 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 - ) + if not self.can_terminate_prefetch(operation): + return False - 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}", - ) + completed_tokens, hash_value = self.cache_controller.terminate_prefetch( + operation + ) - 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, + min_completed_tokens = completed_tokens + if self.tp_world_size > 1: + completed_tokens_tensor = torch.tensor( + min_completed_tokens, dtype=torch.int ) - except Exception as e: - logger.exception( - f"Failed to attach storage backend '{storage_backend}': {e}" + torch.distributed.all_reduce( + completed_tokens_tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, ) - return False, f"Failed to attach storage backend '{storage_backend}': {e}" + min_completed_tokens = completed_tokens_tensor.item() - 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}" + 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 - self._drain_storage_control_queues_local() - self._force_release_pending_storage_ops() + 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, + ) - 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." + # 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] + ) - 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 + # 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) - def _force_release_pending_storage_ops(self): - cc = self.cache_controller + self._release_host_node(last_host_node) + del self.ongoing_prefetch[req_id] + self.cache_controller.prefetch_tokens_occupied -= len(token_ids) - 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.") + loaded_from_storage = min_completed_tokens - matched_length + self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage - 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.") + 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.info( + "HiCache mamba prefetch completed for request %s: prefetched_tokens=%s mamba_states=%s", + req_id, + loaded_from_storage, + int(mamba_loaded), + ) - 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, - ) + return True - def _drain_storage_control_queues_impl( + def _insert_helper_host( self, - n_revoke: Optional[int], - n_backup: Optional[int], - n_release: Optional[int], - log_metrics: bool, + node: TreeNode, + key: RadixKey, + host_value, + hash_value, + mamba_host_value: Optional[torch.Tensor] = None, + mamba_loaded: bool = False, ): - cc = self.cache_controller + node.last_access_time = get_last_access_time() + if len(key) == 0: + return 0 - 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 + child_key = self.get_child_key_fn(key) - 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 + 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 = self.key_match_fn(node.key, key) - 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 - ) + key = key[prefix_len:] + host_value = host_value[prefix_len:] + hash_value = hash_value[prefix_len // self.page_size :] + matched_length += prefix_len - 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) + if prefix_len < len(node.key): + new_node = self._split_node(node.key, node, prefix_len) + node = new_node - _drain_revoke() - _drain_backup() - _drain_release() + if len(key): + child_key = self.get_child_key_fn(key) - 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 + 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) - extra_config = tomllib.load(f) - elif ext in (".yaml", ".yml"): - import yaml + # 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 - 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 + def release_aborted_request(self, rid: str): + self.prefetch_loaded_tokens_by_reqid.pop(rid, None) - 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 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(), ) - 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__}" + 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, ) - 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__}" + ] + + 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 + 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="trailing_pages", ) + ] - return ( - extra_config, - prefetch_threshold, - float(prefetch_timeout_base), - float(prefetch_timeout_per_ki_token), - hicache_storage_pass_prefix_keys, + 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="trailing_pages", + ) + ] - 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 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) - def drain_storage_control_queues(self): - cc = self.cache_controller + 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, + ) + ) - 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 + 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_pool, + 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), + ) ) - 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, - ) + return transfers if transfers else None - 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 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 From 83c610458a12e79e27d97d5b6f4199b1ce81f39f Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Mon, 16 Mar 2026 21:05:25 +0800 Subject: [PATCH 16/39] Optimize code style --- .../srt/mem_cache/hi_mamba_radix_cache.py | 6 ++--- .../sglang/srt/mem_cache/hicache_storage.py | 23 ++++++++++++------- .../hybrid_cache/hybrid_cache_controller.py | 3 ++- .../sglang/srt/mem_cache/memory_pool_host.py | 9 ++++++-- 4 files changed, 27 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 85d245378117..a3648b506a34 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -18,7 +18,7 @@ MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, PrefetchOperation, @@ -2039,7 +2039,7 @@ def mamba_archive_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer] name=PoolName.MAMBA, host_indices=node.mamba_host_value, keys=[node.hash_value[-1]], - hit_policy="trailing_pages", + hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] @@ -2062,7 +2062,7 @@ def mamba_prefetch_alloc( name=PoolName.MAMBA, host_indices=host_indices, keys=["__placeholder__"], - hit_policy="trailing_pages", + hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index fcee37bd64c9..9c86181ae84c 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum -from typing import Any, List, Literal, Optional, Set +from typing import Any, List, Optional, Set import torch @@ -72,23 +72,30 @@ class PoolName(str, Enum): MAMBA = "mamba" +class PoolHitPolicy(str, Enum): + """Hit policy for batch_exists_v2 per-pool prefix matching. + + ALL_PAGES : every page in [0, kv_hit) must exist (default). + TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states). + """ + + ALL_PAGES = "all_pages" + TRAILING_PAGES = "trailing_pages" + + @dataclass class PoolTransfer: """Unified per-pool transfer descriptor for batch v2 interface. device<->host path : host_indices + device_indices host<->storage path: host_indices + keys - - hit_policy controls how batch_exists_v2 checks this pool: - "all_pages" – every page in the KV hit range must exist. - "trailing_pages"– only the last len(keys) pages must exist (e.g. mamba/SWA). """ - name: str + name: PoolName host_indices: Optional[torch.Tensor] = None device_indices: Optional[torch.Tensor] = None keys: Optional[List[str]] = None - hit_policy: Literal["all_pages", "trailing_pages"] = "all_pages" + hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES @dataclass @@ -445,7 +452,7 @@ def has_component(page_idx: int, name: str) -> bool: if final_pages == 0: break name = transfer.name - if transfer.hit_policy == "all_pages": + if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: boundary = next( (i for i in range(kv_pages) if not has_component(i, name)), kv_pages ) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 9f85416ecd65..ce40e425bea9 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -22,6 +22,7 @@ ) from sglang.srt.mem_cache.hicache_storage import ( HiCacheStorageExtraInfo, + PoolHitPolicy, PoolTransfer, PoolTransferResult, ) @@ -445,7 +446,7 @@ def _sync_trailing_keys( is a sliding window of the last N hit pages. """ for transfer in pool_transfers: - if transfer.hit_policy != "trailing_pages": + if transfer.hit_policy != PoolHitPolicy.TRAILING_PAGES: continue trailing_n = len(transfer.keys) if transfer.keys else 1 transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages] diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 065f4f274b54..0560c8fd52da 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1,10 +1,15 @@ +from __future__ import annotations + import abc import logging import threading from collections import defaultdict from dataclasses import dataclass from functools import wraps -from typing import Any, Callable, Optional +from typing import TYPE_CHECKING, Any, Callable, Optional + +if TYPE_CHECKING: + from sglang.srt.mem_cache.hicache_storage import PoolName import numpy as np import psutil @@ -1481,7 +1486,7 @@ def set_from_flat_data_page( @dataclass class PoolEntry: - name: str + name: PoolName host_pool: Any device_pool: Any layer_mapper: Callable[[int], Optional[int]] From b67edcbd0875890a1c481dbeca9f5b5a073d18df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 16 Mar 2026 21:26:43 +0800 Subject: [PATCH 17/39] Simplify Code in hicache_storage.py --- .../sglang/srt/mem_cache/hicache_storage.py | 143 ++++-------------- 1 file changed, 33 insertions(+), 110 deletions(-) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 9c86181ae84c..a60d2bdac0f5 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -470,16 +470,6 @@ def has_component(page_idx: int, name: str) -> bool: hit_count[name] = boundary final_pages = min(final_pages, boundary) - if pool_transfers: - logger.info( - "HiCacheFile batch_exists_v2: kv_pages=%s final_pages=%s hit_count=%s first_key=%s last_key=%s", - kv_pages, - final_pages, - hit_count, - keys[0] if keys else None, - keys[final_pages - 1] if final_pages > 0 else None, - ) - return PoolTransferResult(final_pages, hit_count) def _log_key(self, pool_name: str, key: str) -> str: @@ -487,127 +477,60 @@ def _log_key(self, pool_name: str, key: str) -> str: def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool: """Read one page from storage into host_pool at page_offset.""" - tensor_path = self._get_component_path(key, pool_name) - if not os.path.exists(tensor_path): - return False - try: - with open(tensor_path, "rb", buffering=0) as f: - file_data = f.read() - dummy_page = host_pool.get_dummy_flat_data_page() - expected_bytes = dummy_page.numel() * dummy_page.element_size() - if len(file_data) != expected_bytes: - logger.warning( - "File size mismatch for %s: expected %s, got %s", - self._log_key(pool_name, key), - expected_bytes, - len(file_data), - ) - return False - data_page = ( - torch.frombuffer( - file_data, dtype=dummy_page.dtype, count=dummy_page.numel() - ) - .clone() - .reshape(dummy_page.shape) - ) - host_pool.set_from_flat_data_page(page_offset, data_page) - return True - except Exception as e: - logger.error("Failed to read key %s: %s", self._log_key(pool_name, key), e) + storage_key = self._log_key(pool_name, key) + data_page = self.get(storage_key, host_pool.get_dummy_flat_data_page()) + if data_page is None: return False + host_pool.set_from_flat_data_page(page_offset, data_page) + return True def _write_page( self, pool_name: str, key: str, host_pool, page_offset: int ) -> bool: """Write one page from host_pool at page_offset to storage as raw bytes.""" - try: - data_page = host_pool.get_data_page(page_offset, flat=True) - data_bytes = ( - data_page.contiguous().view(torch.uint8).reshape(-1).numpy().tobytes() - ) - component_path = self._get_component_path(key, pool_name) - if os.path.exists(component_path): - existing_size = os.path.getsize(component_path) - if existing_size == len(data_bytes): - return True # same-size file → assume identical, skip write - logger.warning( - "Overwriting stale storage entry for %s: expected %s bytes, got %s", - self._log_key(pool_name, key), - len(data_bytes), - existing_size, - ) - with open(component_path, "wb", buffering=0) as f: - f.write(data_bytes) - if pool_name != PoolName.KV: - logger.info( - "HiCacheFile wrote auxiliary component %s for key %s", - pool_name, - key, - ) - return True - except Exception as e: - logger.error("Failed to store key %s: %s", self._log_key(pool_name, key), e) - return False + storage_key = self._log_key(pool_name, key) + data_page = host_pool.get_data_page(page_offset, flat=True) + return self.set(storage_key, data_page) - def batch_get_v2( - self, - transfers: List[PoolTransfer], - extra_info: Optional["HiCacheStorageExtraInfo"] = None, - ) -> dict[str, List[bool]]: - del extra_info + def _batch_io_v2(self, transfers: List[PoolTransfer], op_fn): results: dict[str, List[bool]] = {} for transfer in transfers: host_pool = self.registered_pools[transfer.name] + keys = transfer.keys or [] page_size = getattr(host_pool, "page_size", 1) or 1 - if transfer.host_indices.numel() != len(transfer.keys) * page_size: + expected = len(keys) * page_size + host_indices = transfer.host_indices + + if host_indices is None or host_indices.numel() != expected: logger.error( - "batch_get_v2 indices length mismatch for %s: expected %s, got %s", + "%s indices length mismatch for %s: expected %s, got %s", + op_fn.__name__, transfer.name, - len(transfer.keys) * page_size, - transfer.host_indices.numel(), + expected, + host_indices.numel() if host_indices is not None else 0, ) - results[transfer.name] = [False] * len(transfer.keys) - else: - results[transfer.name] = [ - self._read_page( - transfer.name, - key, - host_pool, - transfer.host_indices[i * page_size].item(), - ) - for i, key in enumerate(transfer.keys) - ] + results[transfer.name] = [False] * len(keys) + continue + + results[transfer.name] = [ + op_fn(transfer.name, key, host_pool, host_indices[i * page_size].item()) + for i, key in enumerate(keys) + ] return results + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + return self._batch_io_v2(transfers, self._read_page) + def batch_set_v2( self, transfers: List[PoolTransfer], extra_info: Optional["HiCacheStorageExtraInfo"] = None, ) -> dict[str, List[bool]]: - del extra_info - results: dict[str, List[bool]] = {} - for transfer in transfers: - host_pool = self.registered_pools[transfer.name] - page_size = getattr(host_pool, "page_size", 1) or 1 - if transfer.host_indices.numel() != len(transfer.keys) * page_size: - logger.error( - "batch_set_v2 indices length mismatch for %s: expected %s, got %s", - transfer.name, - len(transfer.keys) * page_size, - transfer.host_indices.numel(), - ) - results[transfer.name] = [False] * len(transfer.keys) - else: - results[transfer.name] = [ - self._write_page( - transfer.name, - key, - host_pool, - transfer.host_indices[i * page_size].item(), - ) - for i, key in enumerate(transfer.keys) - ] - return results + return self._batch_io_v2(transfers, self._write_page) def clear(self) -> bool: try: From 18d724bdb4e7c443226c5e4e23ae03defc2eefe9 Mon Sep 17 00:00:00 2001 From: ispobock Date: Mon, 16 Mar 2026 15:12:26 +0000 Subject: [PATCH 18/39] fix evict num --- .../srt/mem_cache/hi_mamba_radix_cache.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index a3648b506a34..c4eebf06d6fd 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -600,14 +600,14 @@ def _evict_regular(self, node: TreeNode) -> Tuple[int, int]: assert not node.backuped, f"backuped node, {node.id=}" assert len(node.children) == 0, f"non-leaf, {node.id=}" - num_full = len(node.value) + full_num_evicted = len(node.value) self.cache_controller.evict_device(node.value) - self.full_evictable_size_ -= num_full + self.full_evictable_size_ -= full_num_evicted if self.full_lru_list.in_list(node): self.full_lru_list.remove_node(node) - mamba_num = self._free_device_mamba(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): @@ -624,8 +624,13 @@ def _evict_regular(self, node: TreeNode) -> Tuple[int, int]: assert v == node, f"parent does not have child key, {key}" self._update_leaf_status(parent) - self._iteratively_delete_tombstone_leaf(node) - return num_full, mamba_num + _, 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 @@ -636,7 +641,7 @@ def _evict_host_leaf(self, node: TreeNode) -> int: node.host_ref_counter == 0 ), f"in use, {node.id=} {node.host_ref_counter=}" - num_evicted = self.cache_controller.evict_host(node.host_value) + full_num_evicted = self.cache_controller.evict_host(node.host_value) node.host_value = None if node.mamba_host_value is not None: @@ -652,9 +657,9 @@ def _evict_host_leaf(self, node: TreeNode) -> int: assert v == node, f"parent does not have child key, {key}" self._update_leaf_status(parent) - self._iteratively_delete_tombstone_leaf(node) + _, cascade_full_num_evicted, _ = self._iteratively_delete_tombstone_leaf(node) - return num_evicted + 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=}" From 3204f85028b78d65fd76772e009225d64439cd77 Mon Sep 17 00:00:00 2001 From: ispobock Date: Mon, 16 Mar 2026 15:13:26 +0000 Subject: [PATCH 19/39] remove unused --- python/sglang/srt/mem_cache/hi_mamba_radix_cache.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index c4eebf06d6fd..4e874a9c39f0 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -859,13 +859,6 @@ def _unevict_node(self, node: TreeNode, fresh_value: torch.Tensor): 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) From abfc3d3313d125dfb929af625d4083eeed50eed9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Tue, 17 Mar 2026 15:59:23 +0800 Subject: [PATCH 20/39] Simplify Code --- .../sglang/srt/mem_cache/hicache_storage.py | 36 ++----------------- 1 file changed, 2 insertions(+), 34 deletions(-) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index a60d2bdac0f5..815bc33a2d07 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -324,18 +324,6 @@ def __init__( def _get_suffixed_key(self, key: str) -> str: return key + self.config_suffix - def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str: - if component_name is None or component_name in ("__default__", PoolName.KV): - return self._get_suffixed_key(key) - return self._get_suffixed_key(f"{key}.{component_name}") - - def _get_component_path( - self, key: str, component_name: Optional[str] = None - ) -> str: - return os.path.join( - self.file_path, f"{self._get_component_key(key, component_name)}.bin" - ) - def get( self, key: str, @@ -405,42 +393,22 @@ def exists(self, key: str) -> bool: tensor_path = os.path.join(self.file_path, f"{key}.bin") return os.path.exists(tensor_path) - def _collect_existing_component_keys( - self, - keys: List[str], - pool_transfers: Optional[List[PoolTransfer]] = None, - ) -> Set[str]: - target_files = {f"{self._get_component_key(key)}.bin" for key in keys} - for transfer in pool_transfers or []: - for key in keys: - target_files.add(f"{self._get_component_key(key, transfer.name)}.bin") - - existing_files = set() - with os.scandir(self.file_path) as entries: - for entry in entries: - if entry.is_file() and entry.name in target_files: - existing_files.add(entry.name) - return existing_files - def batch_exists_v2( self, keys: List[str], pool_transfers: Optional[List[PoolTransfer]] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> PoolTransferResult: - existing_files = self._collect_existing_component_keys(keys, pool_transfers) def has_component(page_idx: int, name: str) -> bool: - return ( - f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files - ) + return self.exists(self._log_key(name, keys[page_idx])) # Longest contiguous KV prefix present in storage. kv_pages = next( ( i for i in range(len(keys)) - if f"{self._get_component_key(keys[i])}.bin" not in existing_files + if self.exists(keys[i]) ), len(keys), ) From 6318d9cd1b4e5dd158c3dca3b00f1c0517a16b3d Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Tue, 17 Mar 2026 17:37:48 +0800 Subject: [PATCH 21/39] Revert "Simplify Code" This reverts commit abfc3d3313d125dfb929af625d4083eeed50eed9. --- .../sglang/srt/mem_cache/hicache_storage.py | 36 +++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 898666b39bcf..8e39fd0acc14 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -325,6 +325,18 @@ def __init__( def _get_suffixed_key(self, key: str) -> str: return key + self.config_suffix + def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str: + if component_name is None or component_name in ("__default__", PoolName.KV): + return self._get_suffixed_key(key) + return self._get_suffixed_key(f"{key}.{component_name}") + + def _get_component_path( + self, key: str, component_name: Optional[str] = None + ) -> str: + return os.path.join( + self.file_path, f"{self._get_component_key(key, component_name)}.bin" + ) + def get( self, key: str, @@ -394,22 +406,42 @@ def exists(self, key: str) -> bool: tensor_path = os.path.join(self.file_path, f"{key}.bin") return os.path.exists(tensor_path) + def _collect_existing_component_keys( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + ) -> Set[str]: + target_files = {f"{self._get_component_key(key)}.bin" for key in keys} + for transfer in pool_transfers or []: + for key in keys: + target_files.add(f"{self._get_component_key(key, transfer.name)}.bin") + + existing_files = set() + with os.scandir(self.file_path) as entries: + for entry in entries: + if entry.is_file() and entry.name in target_files: + existing_files.add(entry.name) + return existing_files + def batch_exists_v2( self, keys: List[str], pool_transfers: Optional[List[PoolTransfer]] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> PoolTransferResult: + existing_files = self._collect_existing_component_keys(keys, pool_transfers) def has_component(page_idx: int, name: str) -> bool: - return self.exists(self._log_key(name, keys[page_idx])) + return ( + f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files + ) # Longest contiguous KV prefix present in storage. kv_pages = next( ( i for i in range(len(keys)) - if self.exists(keys[i]) + if f"{self._get_component_key(keys[i])}.bin" not in existing_files ), len(keys), ) From 571bce4226697590be557dd9bf55c7b890112e75 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Tue, 17 Mar 2026 17:55:49 +0800 Subject: [PATCH 22/39] Add comment for io kernel --- python/sglang/srt/mem_cache/memory_pool_host.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d7e4c92759e0..07362ef93e5b 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1274,6 +1274,9 @@ def _copy_tensor( if src_indices.numel() == 0: return if io_backend == "kernel": + # TODO: Rename the interface for clarity. + # Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state. + # This has nothing to do with MLA; it's only reused because this interface happens to transfer a single Pool. transfer_kv_per_layer_mla( src=src, dst=dst, From c825a08d1f76d917e2b3852eb5d54c379eecb1fc Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Tue, 17 Mar 2026 19:47:44 +0800 Subject: [PATCH 23/39] Fix mixed chunk ci --- test/registered/unit/managers/test_prefill_adder.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index ea1ed8bc1bff..8c5ed3625632 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -81,6 +81,7 @@ def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens) req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.finished.return_value = False + req.should_init_load_back_host = False return req def create_adder(self, running_batch, **kwargs): From 83b766339c2bae4b5e3d1c7e770e201b1a472ef1 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Wed, 18 Mar 2026 14:44:12 +0800 Subject: [PATCH 24/39] Optimize log; Add CI --- python/sglang/srt/managers/schedule_policy.py | 2 +- .../srt/mem_cache/hi_mamba_radix_cache.py | 19 +-------- .../hybrid_cache/hybrid_cache_controller.py | 1 - .../4-gpu-models/test_qwen35_models.py | 40 +++++++++++++++---- 4 files changed, 35 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index e786d5b2db73..fba8b683c3f6 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -773,7 +773,7 @@ def add_one_req( InitLoadBackParams( last_host_node=req.last_host_node, host_hit_length=req.host_hit_length, - req=req + req=req, ) ) req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index d111aab9a0a9..5ed4e2c211d7 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -308,17 +308,6 @@ def write_backup(self, node: TreeNode, write_back=False): if not write_back: # no need to lock nodes if write back self.inc_lock_ref(node) - if extra_pools is not None: - logger.info( - "HiCache mamba offload prepared for node %s: kv_tokens=%s mamba_states=%s", - node.id, - len(node.host_value), - ( - len(node.mamba_host_value) - if node.mamba_host_value is not None - else 0 - ), - ) else: return 0 @@ -366,7 +355,7 @@ def load_back( self.dec_lock_ref(ancestor_node) return None - logger.info( + 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( @@ -1778,12 +1767,6 @@ def prefetch_from_storage( self._release_host_node(last_host_node) return - logger.info( - "HiCache mamba prefetch scheduled for request %s: kv_hit_pages=%s mamba_states=%s", - req_id, - prefetch_length // self.page_size, - 1, - ) operation = self.cache_controller.prefetch( req_id, host_indices, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index ce40e425bea9..7d6abdab41ef 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -212,7 +212,6 @@ def attach_storage_backend( for entry in host_pools or []: self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name) - logger.info("Using pool-based interface for hybrid storage operations") def reset(self): super().reset() diff --git a/test/registered/4-gpu-models/test_qwen35_models.py b/test/registered/4-gpu-models/test_qwen35_models.py index d67a298b77fa..e5c640907d05 100644 --- a/test/registered/4-gpu-models/test_qwen35_models.py +++ b/test/registered/4-gpu-models/test_qwen35_models.py @@ -1,5 +1,6 @@ import shutil import tempfile +import time import unittest from types import SimpleNamespace @@ -253,6 +254,10 @@ def setUpClass(cls): other_args=[ "--tp-size", "4", + "--max-mamba-cache-size", + "500", + "--max-total-tokens", + "120000", "--chunked-prefill-size", "2048", "--mamba-scheduler-strategy", @@ -288,14 +293,14 @@ def tearDownClass(cls): kill_process_tree(cls.process.pid) shutil.rmtree(cls.storage_dir, ignore_errors=True) - def test_gsm8k(self): + def _run_gsm8k(self): args = SimpleNamespace( model=self.model, eval_name="gsm8k", num_shots=5, - num_examples=200, - max_tokens=16000, - num_threads=128, + num_examples=100, + max_tokens=4000, + num_threads=50, repeat=1, temperature=0.6, top_p=0.95, @@ -304,9 +309,30 @@ def test_gsm8k(self): host="http://127.0.0.1", port=int(self.base_url.split(":")[-1]), ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]) + return run_eval(args) + + def test_gsm8k(self): + first_metrics = self._run_gsm8k() + print(f"first_metrics={first_metrics}") + self.assertGreaterEqual( + first_metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + print(f"flush cache") + time.sleep(2) + requests.post(f"{self.base_url}/flush_cache", timeout=10) + + second_metrics = self._run_gsm8k() + print(f"second_metrics={second_metrics}") + self.assertGreaterEqual( + second_metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + self.assertLessEqual( + abs(second_metrics["score"] - first_metrics["score"]), + 0.05, + f"HiCache prefetch accuracy drift too large: " + f"first={first_metrics['score']}, second={second_metrics['score']}", + ) if __name__ == "__main__": From 61f66700fc295ec910b1cdd0716a60665c3559b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Sun, 22 Mar 2026 21:06:58 +0800 Subject: [PATCH 25/39] support mooncake backend & dsa model Co-authored-by: hzh0425 Co-authored-by: Vladislav Nosivskoy --- .../sglang/srt/mem_cache/hicache_storage.py | 3 +- python/sglang/srt/mem_cache/hiradix_cache.py | 204 ++++++++++--- .../hybrid_cache/hybrid_cache_controller.py | 18 +- .../sglang/srt/mem_cache/memory_pool_host.py | 269 ++++++++++++++---- .../storage/mooncake_store/mooncake_store.py | 233 ++++++++++++--- 5 files changed, 597 insertions(+), 130 deletions(-) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 8e39fd0acc14..ff4a0d5e904f 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -71,12 +71,13 @@ class PoolName(str, Enum): KV = "kv" MAMBA = "mamba" + INDEXER = "indexer" class PoolHitPolicy(str, Enum): """Hit policy for batch_exists_v2 per-pool prefix matching. - ALL_PAGES : every page in [0, kv_hit) must exist (default). + ALL_PAGES : every page in [0, kv_hit) must exist (e.g. DSA). TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states). """ diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e13d43c371b6..ebee9df69560 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -26,15 +26,25 @@ MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, +) +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, MLATokenToKVPool, NSATokenToKVPool, ) from sglang.srt.mem_cache.memory_pool_host import ( + HostPoolGroup, MHATokenToKVPoolHost, MLATokenToKVPoolHost, - NSATokenToKVPoolHost, + NSAIndexerPoolHost, + PoolEntry, ) from sglang.srt.mem_cache.radix_cache import ( RadixCache, @@ -54,6 +64,83 @@ logger = logging.getLogger(__name__) +def _build_nsa_hybrid_stack( + radix_cache: "HiRadixCache", + params: "CacheInitParams", + server_args: "ServerArgs", + *, + extra_config: dict, + prefetch_threshold: int, + enable_storage_metrics: bool, + load_cache_event, +) -> None: + """HostPoolGroup (KV + indexer) + HybridCacheController for NSA (DSA).""" + kv = radix_cache.kv_cache + mla_host = MLATokenToKVPoolHost( + kv, + server_args.hicache_ratio, + server_args.hicache_size, + radix_cache.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + override_kv_cache_dim=kv.kv_cache_dim, + ) + indexer_host = NSAIndexerPoolHost( + kv, + mla_host, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + layer_num = kv.layer_num + + def layer_mapper(layer_id: int): + if 0 <= layer_id < layer_num: + return layer_id + return None + + host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=mla_host, + device_pool=kv, + layer_mapper=layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.INDEXER, + host_pool=indexer_host, + device_pool=kv, + layer_mapper=layer_mapper, + share_indices_with_anchor=True, + ), + ] + ) + radix_cache.token_to_kv_pool_host = host_pool_group + radix_cache.cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + radix_cache.page_size, + radix_cache.tp_group, + load_cache_event=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=radix_cache.pp_rank, + pp_size=radix_cache.pp_size, + transfer_layer_num=layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + logger.info( + "NSA hierarchical cache: HostPoolGroup(KV + INDEXER), HybridCacheController, " + "transfer_layer_num=%s", + layer_num, + ) + + class HiRadixCache(RadixCache): def __init__(self, params: CacheInitParams, server_args: ServerArgs): @@ -82,14 +169,8 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): allocator_type=server_args.hicache_storage_backend, ) elif isinstance(self.kv_cache, NSATokenToKVPool): - self.token_to_kv_pool_host = NSATokenToKVPoolHost( - self.kv_cache, - server_args.hicache_ratio, - server_args.hicache_size, - self.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) + # Filled by _build_nsa_hybrid_stack after storage extra_config is parsed. + self.token_to_kv_pool_host = None elif isinstance(self.kv_cache, MLATokenToKVPool): self.token_to_kv_pool_host = MLATokenToKVPoolHost( self.kv_cache, @@ -100,7 +181,9 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): allocator_type=server_args.hicache_storage_backend, ) else: - raise ValueError(f"HiRadixCache only supports MHA and MLA yet") + raise ValueError( + "HiRadixCache only supports MHA, MLA, and NSA (DSA) models" + ) self.tp_group = params.tp_cache_group self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) @@ -124,22 +207,33 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): 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.token_to_kv_pool_host, - self.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=self.pp_rank, - pp_size=self.pp_size, - enable_storage_metrics=self.enable_storage_metrics, - ) + if isinstance(self.kv_cache, NSATokenToKVPool): + _build_nsa_hybrid_stack( + self, + params, + server_args, + extra_config=extra_config, + prefetch_threshold=prefetch_threshold, + enable_storage_metrics=self.enable_storage_metrics, + load_cache_event=self.load_cache_event, + ) + else: + self.cache_controller = HiCacheController( + params.token_to_kv_pool_allocator, + self.token_to_kv_pool_host, + self.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=self.pp_rank, + pp_size=self.pp_size, + enable_storage_metrics=self.enable_storage_metrics, + ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, prefetch_threshold=prefetch_threshold, @@ -339,12 +433,21 @@ def attach_storage_backend( ) 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, - ) + if isinstance(self.cache_controller, HybridCacheController): + 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.cache_controller.mem_pool_host.entries, + ) + else: + 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}" @@ -610,6 +713,22 @@ def get_height(self, node: TreeNode): height += 1 return height + def _hybrid_pools_kw(self, host_indices: Optional[torch.Tensor]=None) -> dict: + if not isinstance( + self.cache_controller, HybridCacheController + ): + return {} + if isinstance(self.kv_cache, NSATokenToKVPool): + pool = PoolTransfer( + name=PoolName.INDEXER, + host_indices=host_indices, + device_indices=None, + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + return {"extra_pools": [pool]} + else: + return {} + def clear_storage_backend(self) -> bool: if self.enable_storage: try: @@ -633,15 +752,18 @@ def clear_storage_backend(self) -> bool: return False def write_backup(self, node: TreeNode, write_back=False): + write_kw = self._hybrid_pools_kw() host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, + **write_kw, ) 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, + **write_kw, ) if host_indices is not None: node.host_value = host_indices @@ -662,8 +784,9 @@ def write_backup_storage(self, node: TreeNode): else None ) + storage_kw = self._hybrid_pools_kw(host_indices=node.host_value) operation_id = self.cache_controller.write_storage( - node.host_value, node.key, node.hash_value, prefix_keys + node.host_value, node.key, node.hash_value, prefix_keys, **storage_kw ) self.ongoing_backup[operation_id] = node node.protect_host() @@ -1030,12 +1153,16 @@ def load_back( return None device_indices = self.cache_controller.load( - host_indices=host_indices, node_id=last_hit_node.id + host_indices=host_indices, + node_id=last_hit_node.id, + **self._hybrid_pools_kw(host_indices=host_indices), ) if device_indices is None: self.evict(EvictParams(num_tokens=len(host_indices))) device_indices = self.cache_controller.load( - host_indices=host_indices, node_id=last_hit_node.id + host_indices=host_indices, + node_id=last_hit_node.id, + **self._hybrid_pools_kw(host_indices=host_indices), ) self.dec_lock_ref(ancester_node) if device_indices is None: @@ -1336,7 +1463,12 @@ def prefetch_from_storage( # no sufficient host memory for prefetch return operation = self.cache_controller.prefetch( - req_id, host_indices, new_input_tokens, last_hash, prefix_keys + req_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys, + **self._hybrid_pools_kw(host_indices=host_indices), ) self.ongoing_prefetch[req_id] = ( last_host_node, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 7d6abdab41ef..77817b1264ca 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -163,6 +163,7 @@ def __init__( pp_rank: int = 0, pp_size: int = 1, transfer_layer_num: Optional[int] = None, + enable_storage_metrics: bool = False, ): startup_storage_backend = storage_backend super().__init__( @@ -179,6 +180,7 @@ def __init__( storage_backend_extra_config=storage_backend_extra_config, pp_rank=pp_rank, pp_size=pp_size, + enable_storage_metrics=enable_storage_metrics, ) # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)), # not just the full attention layers reported by full_kv_pool. @@ -230,7 +232,10 @@ def write( if host_indices is None: return None pool_transfers = self._resolve_pool_transfers_allocation( - extra_pools, alloc_host=True + extra_pools, + alloc_host=True, + kv_device_indices=device_indices, + kv_host_indices=host_indices ) if pool_transfers is None and extra_pools: self.mem_pool_host.free(host_indices) @@ -289,7 +294,10 @@ def load( device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) pool_transfers = self._resolve_pool_transfers_allocation( - extra_pools, alloc_host=False + extra_pools, + alloc_host=False, + kv_device_indices=device_indices, + kv_host_indices=host_indices ) if pool_transfers is None and extra_pools: if need_load_kv: @@ -454,6 +462,8 @@ def _resolve_pool_transfers_allocation( self, extra_pools: Optional[list[PoolTransfer]], alloc_host: bool, + kv_device_indices: Optional[torch.Tensor] = None, + kv_host_indices: Optional[torch.Tensor] = None, ) -> Optional[list[PoolTransfer]]: """Auto-alloc host or device indices for PoolTransfers where they are None.""" if not extra_pools: @@ -463,6 +473,10 @@ def _resolve_pool_transfers_allocation( entry = self.mem_pool_host.entry_map.get(pool.name) if entry is None: continue + if entry.share_indices_with_anchor: + pool.device_indices = kv_device_indices + pool.host_indices = kv_host_indices + continue if alloc_host: if pool.host_indices is not None or pool.device_indices is None: continue diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 07362ef93e5b..d4488874af72 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1205,6 +1205,10 @@ def init_kv_buffer(self): ) ) + def get_hybrid_pool_buffer(self): + # Expose all mamba host tensors that need Mooncake buffer registration. + return [self.temporal_buffer, *self.conv_buffer] + def _iter_page_tensors(self, index: int): if self.layout in ["page_first", "page_first_direct"]: yield self.temporal_buffer[index] @@ -1487,6 +1491,65 @@ def set_from_flat_data_page( restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape) tensor.copy_(restored) + def get_page_buffer_meta(self, indices): + """Meta data for zero-copy storage I/O. + + Only page-first layouts are supported for mamba storage zero-copy because + each page slot in temporal/conv buffers is directly addressable. + """ + assert len(indices) % self.page_size == 0 + if self.layout not in ["page_first", "page_first_direct"]: + raise ValueError( + f"Mamba storage zero-copy requires page_first layout, got {self.layout}" + ) + indices = indices.tolist() + ptr_list = [] + element_size_list = [] + + # Compute base pointers once; each page pointer is offset from these bases. + temporal_base_ptr = self.temporal_buffer.data_ptr() + conv_base_ptrs = [buf.data_ptr() for buf in self.conv_buffer] + # Component sizes are constant across pages, so precompute once as well. + temporal_element_size = ( + self.page_size + * self.num_mamba_layers + * self.temporal_dtype.itemsize + * self.temporal_state_shape + ) + conv_element_sizes = [ + ( + self.page_size + * self.num_mamba_layers + * self.conv_dtype.itemsize + * conv_shape + ) + for conv_shape in self.conv_state_shapes + ] + + for i in range(0, len(indices), self.page_size): + # Emit component pointers in stable order: + # temporal first, then conv_0..conv_n for this page. + temporal_ptr = ( + temporal_base_ptr + + indices[i] + * self.num_mamba_layers + * self.temporal_state_shape + * self.temporal_dtype.itemsize + ) + ptr_list.append(temporal_ptr) + element_size_list.append(temporal_element_size) + for j in range(len(self.conv_buffer)): + conv_ptr = ( + conv_base_ptrs[j] + + indices[i] + * self.num_mamba_layers + * self.conv_state_shapes[j] + * self.conv_dtype.itemsize + ) + ptr_list.append(conv_ptr) + element_size_list.append(conv_element_sizes[j]) + return ptr_list, element_size_list + @dataclass class PoolEntry: @@ -1495,6 +1558,9 @@ class PoolEntry: device_pool: Any layer_mapper: Callable[[int], Optional[int]] is_primary_index_anchor: bool = False + # When True, host_pool uses the same logical slot indices as the anchor pool + # (e.g. DSA indexer); HostPoolGroup.free mirrors frees to this pool. + share_indices_with_anchor: bool = False # Optional eviction callbacks for auto-alloc in HybridCacheController. # host_evict_fn(n): evict n slots from the host pool (used by write()). # device_evict_fn(n): evict n slots from the device pool (used by load()). @@ -1518,6 +1584,36 @@ def __init__(self, entries: list[PoolEntry]): self.device = self.anchor_entry.host_pool.device self.size = self.anchor_entry.host_pool.size + @property + def kv_buffer(self): + return self.anchor_entry.host_pool.kv_buffer + + @property + def size_per_token(self): + return self.anchor_entry.host_pool.size_per_token + + @property + def allocator(self): + return self.anchor_entry.host_pool.allocator + + @property + def dtype(self): + return self.anchor_entry.host_pool.dtype + + @property + def start_layer(self): + return self.anchor_entry.host_pool.start_layer + + @property + def end_layer(self): + return self.anchor_entry.host_pool.end_layer + + def get_ksize_per_token(self): + return self.anchor_entry.host_pool.get_ksize_per_token() + + def get_page_buffer_meta(self, indices): + return self.anchor_entry.host_pool.get_page_buffer_meta(indices) + def clear(self) -> None: for entry in self.entries: entry.host_pool.clear() @@ -1526,7 +1622,13 @@ def alloc(self, need_size: int) -> Optional[torch.Tensor]: return self.anchor_entry.host_pool.alloc(need_size) def free(self, indices: torch.Tensor) -> int: - return self.anchor_entry.host_pool.free(indices) + n = self.anchor_entry.host_pool.free(indices) + for entry in self.entries: + if entry is self.anchor_entry: + continue + if getattr(entry, "share_indices_with_anchor", False): + entry.host_pool.free(indices) + return n def get_data_page(self, index, flat: bool = True): return self.anchor_entry.host_pool.get_data_page(index, flat) @@ -1602,21 +1704,31 @@ def backup_from_device_all_layer( ) -class NSATokenToKVPoolHost(MLATokenToKVPoolHost): +class NSAIndexerPoolHost(HostKVCache): + """Host-side NSA index buffers only. Slot layout matches the anchor MLA host pool.""" + device_pool: NSATokenToKVPool def __init__( self, device_pool: NSATokenToKVPool, - host_to_device_ratio: float, - host_size: int, - page_size: int, + anchor_host: MLATokenToKVPoolHost, layout: str, pin_memory: bool = True, device: str = "cpu", allocator_type: str = "default", ): - # Initialize indexer metadata before HostKVCache.__init__ calls get_size_per_token. + self.device_pool = device_pool + self.page_size = anchor_host.page_size + self.layout = layout + self.pin_memory = pin_memory + self.device = device + self.allocator = get_allocator_from_storage(allocator_type) + self.dtype = device_pool.store_dtype + self.start_layer = device_pool.start_layer + self.end_layer = device_pool.end_layer + self.layer_num = device_pool.layer_num + self.index_head_dim = device_pool.index_head_dim self.indexer_quant_block_size = device_pool.quant_block_size self.indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype @@ -1624,34 +1736,46 @@ def __init__( self.index_head_dim + self.index_head_dim // self.indexer_quant_block_size * 4 ) - super().__init__( - device_pool, - host_to_device_ratio, - host_size, - page_size, - layout, - pin_memory, - device, - allocator_type, - override_kv_cache_dim=device_pool.kv_cache_dim, - ) + self.size = anchor_host.size + self.page_num = anchor_host.page_num + self.indexer_page_stride_size = ( self.indexer_size_per_token * self.page_size * self.indexer_dtype.itemsize ) self.indexer_layout_dim = self.indexer_page_stride_size * self.layer_num self.indexer_page_num = (self.size + self.page_size + 1) // self.page_size - self._init_indexer_buffers() + self.size_per_token = ( + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize + ) + + buf_elem_size = self.page_num * self.layer_num * self.indexer_page_stride_size + requested_bytes = buf_elem_size * self.indexer_dtype.itemsize + host_mem = psutil.virtual_memory() + ten_gb = 10 * (1024**3) + available_bytes = host_mem.available - ten_gb + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory for NSA indexer hierarchical cache. " + f"Requesting {requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free." + ) logger.info( - f"NSATokenToKVPoolHost initialized with indexer page stride size: {self.indexer_page_stride_size}, page num: {self.indexer_page_num}" + "Allocating %.2f GB host memory for NSA indexer (layout=%s).", + requested_bytes / 1e9, + layout, ) + self.kv_buffer = self.init_kv_buffer() + self.lock = threading.RLock() + self.clear() def get_size_per_token(self): - base = super().get_size_per_token() return ( - base - + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize ) + def get_ksize_per_token(self): + return self.get_size_per_token() + def _init_indexer_buffers(self): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] self.index_k_device_ptrs = torch.tensor( @@ -1660,16 +1784,13 @@ def _init_indexer_buffers(self): device=self.device_pool.device, ) if self.layout == "layer_first": - self.index_k_with_scale_buffer = [ - alloc_func( - (self.indexer_page_num, self.indexer_page_stride_size), - dtype=self.indexer_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - for _ in range(self.layer_num) - ] + self.index_k_with_scale_buffer = alloc_func( + (self.layer_num, self.indexer_page_num, self.indexer_page_stride_size), + dtype=self.indexer_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) self.index_k_data_refs = [ self.index_k_with_scale_buffer[i] for i in range(self.layer_num) ] @@ -1694,6 +1815,9 @@ def _init_indexer_buffers(self): else: raise ValueError(f"Unsupported layout: {self.layout}") + def get_hybrid_pool_buffer(self): + return [self.index_k_with_scale_buffer] + def _get_indexer_page_indices(self, host_indices, device_indices): if host_indices.numel() == 0: return host_indices, device_indices @@ -1709,7 +1833,7 @@ def _get_indexer_page_indices(self, host_indices, device_indices): ) return host_page_indices, device_page_indices - def _load_indexer_to_device_per_layer( + def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): host_page_indices, device_page_indices = self._get_indexer_page_indices( @@ -1760,7 +1884,7 @@ def _load_indexer_to_device_per_layer( else: raise ValueError(f"Unsupported IO backend: {io_backend}") - def _backup_indexer_from_device_all_layer( + def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): host_page_indices, device_page_indices = self._get_indexer_page_indices( @@ -1811,27 +1935,60 @@ def _backup_indexer_from_device_all_layer( else: raise ValueError(f"Unsupported IO backend: {io_backend}") - def load_to_device_per_layer( - self, - device_pool, - host_indices, - device_indices, - layer_id, - io_backend, - ): - super().load_to_device_per_layer( - device_pool, host_indices, device_indices, layer_id, io_backend - ) - self._load_indexer_to_device_per_layer( - device_pool, host_indices, device_indices, layer_id, io_backend - ) + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + page_idx = int(index) // self.page_size + if self.layout == "layer_first": + data_page = self.index_k_with_scale_buffer[:, page_idx : page_idx + 1, :] + elif self.layout in ["page_first", "page_first_direct"]: + data_page = self.index_k_with_scale_buffer[page_idx : page_idx + 1, :, :, :] + else: + raise ValueError(f"Unsupported layout: {self.layout}") + if flat: + data_page = data_page.flatten() + return data_page - def backup_from_device_all_layer( - self, device_pool, host_indices, device_indices, io_backend - ): - super().backup_from_device_all_layer( - device_pool, host_indices, device_indices, io_backend - ) - self._backup_indexer_from_device_all_layer( - device_pool, host_indices, device_indices, io_backend + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + (self.layer_num, self.indexer_page_stride_size), + dtype=self.indexer_dtype, + device=self.device, + pin_memory=self.pin_memory, + ).flatten() + + def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + page_idx = int(index) // self.page_size + if self.layout == "layer_first": + self.index_k_with_scale_buffer[:, page_idx : page_idx + 1, :] = ( + data_page.reshape( + self.layer_num, + 1, + self.indexer_page_stride_size, + ) + ) + elif self.layout in ["page_first", "page_first_direct"]: + self.index_k_with_scale_buffer[page_idx : page_idx + 1, :, :, :] = ( + data_page.reshape( + 1, + self.layer_num, + 1, + self.indexer_page_stride_size, + ) + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + + def get_page_buffer_meta(self, indices): + """Meta data for zero-copy storage I/O.""" + assert len(indices) % self.page_size == 0 + if self.layout not in ["page_first", "page_first_direct"]: + raise ValueError(f"Unsupported layout: {self.layout}") + ptr_list = [] + indices = indices.tolist() + page_stride_bytes = ( + self.layer_num * self.indexer_page_stride_size * self.indexer_dtype.itemsize ) + base_ptr = self.index_k_with_scale_buffer.data_ptr() + for i in range(0, len(indices), self.page_size): + page_index = int(indices[i]) // self.page_size + ptr_list.append(base_ptr + page_index * page_stride_bytes) + return ptr_list, [page_stride_bytes] * len(ptr_list) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 2c815fd7e20d..79bcc4bf2cff 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -15,6 +15,10 @@ HiCacheStorage, HiCacheStorageConfig, HiCacheStorageExtraInfo, + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, ) from sglang.srt.mem_cache.memory_pool_host import HostKVCache, HostTensorAllocator from sglang.srt.observability.metrics_collector import StorageMetrics @@ -424,6 +428,8 @@ def __init__( else: self.mha_suffix = [f"{rank}" for rank in target_ranks] + self.registered_pools = {} + self.gb_per_page = None self.prefetch_pgs = [] self.backup_pgs = [] @@ -494,6 +500,173 @@ def register_mem_pool_host(self, mem_pool_host: HostKVCache): bytes_per_page = mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size self.gb_per_page = bytes_per_page / (1 << 30) + def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): + # Keep a name->pool mapping so batch v2 can resolve PoolTransfer.name to + # the corresponding host pool implementation at runtime. + self.registered_pools[host_pool_name] = host_pool + # KV anchor memory is already registered via register_mem_pool_host(). + # v2 here only registers additional hybrid pools. + if host_pool_name == PoolName.KV: + return + + # Hybrid pools expose the tensors that Mooncake needs for zero-copy I/O. + # The storage backend only depends on this accessor, not concrete fields. + assert getattr(host_pool, "get_hybrid_pool_buffer") is not None + buf_list = host_pool.get_hybrid_pool_buffer() + for buf in buf_list: + super().register_buffer(buf) + + def _tag_keys(self, keys: List[str]) -> List[str]: + if self.extra_backend_tag is None: + return keys + prefix = self.extra_backend_tag + return [f"{prefix}_{key}" for key in keys] + + def _get_hybrid_page_component_keys( + self, page_keys: List[str], transfer: PoolTransfer + ): + # A logical "page" may map to multiple physical objects in storage. + # - INDEXER: one key per page + # - MAMBA : one temporal key + N conv keys per page + # key_multiplier records how many component keys are generated per page, + # so post-processing can fold component results back to page-level booleans. + name = str(transfer.name) + component_keys = [] + key_multiplier = 0 + for page_key in page_keys: + if name == PoolName.INDEXER.value: + key_multiplier = 1 + component_keys.append( + f"{page_key}_{self.mla_suffix}_{PoolName.INDEXER}" + ) + elif name == PoolName.MAMBA.value: + base_key = f"{page_key}_{self.mha_suffix}" + mamba_pool = getattr(self, "registered_pools", {}).get(PoolName.MAMBA) + conv_num = len(getattr(mamba_pool, "conv_buffer", []) or []) + key_multiplier = conv_num + 1 + component_keys.extend( + [f"{base_key}_temporal"] + + [f"{base_key}_conv_{i}" for i in range(conv_num)] + ) + return component_keys, key_multiplier + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + qkeys = self._tag_keys(keys) + kv_pages = self.batch_exists(qkeys, extra_info) + + hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {} + final_pages = kv_pages + + def has_component(exist_results: List[int], page_idx: int, key_multiplier: int): + # Group component existence by page, then require all components hit. + result_groups = exist_results[ + (page_idx - 1) * key_multiplier : page_idx * key_multiplier + ] + return all(result == 1 for result in result_groups) + + for transfer in pool_transfers or []: + if final_pages == 0: + break + component_keys, key_multiplier = self._get_hybrid_page_component_keys( + qkeys, transfer + ) + ex = self._batch_exist(component_keys) + boundary = 0 + if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: + boundary = next( + ( + i + for i in range(kv_pages) + if not has_component(ex, i, key_multiplier) + ), + kv_pages, + ) + elif transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES: + trailing = max(1, len(transfer.keys) if transfer.keys else 1) + for prefix_len in range(kv_pages, 0, -1): + if all( + has_component(ex, i, key_multiplier) + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + if boundary: + hit_count[transfer.name] = boundary + final_pages = min(final_pages, boundary) + + return PoolTransferResult(final_pages, hit_count) + + def _batch_io_v2(self, transfers: List[PoolTransfer], is_set: bool): + # Unified v2 I/O path: each PoolTransfer can expand to one or more + # storage objects per logical page, but API still reports page-level result. + results: dict = {} + for transfer in transfers: + host_pool = getattr(self, "registered_pools", {}).get(transfer.name) + keys = transfer.keys or [] + page_size = getattr(host_pool, "page_size", 1) or 1 + host_indices = transfer.host_indices + assert len(keys) > 0 + assert len(keys) == host_indices // page_size + + ptr_list, element_size_list = host_pool.get_page_buffer_meta(host_indices) + key_strs, key_multiplier = self._get_hybrid_page_component_keys( + keys, transfer + ) + key_strs = self._tag_keys(key_strs) + + if is_set: + exist_result = self._batch_exist(key_strs) + set_keys = [] + set_ptrs = [] + set_sizes = [] + set_idx = [] + set_results = [-1] * len(key_strs) + for i in range(len(key_strs)): + if exist_result[i] != 1: + set_keys.append(key_strs[i]) + set_ptrs.append(ptr_list[i]) + set_sizes.append(element_size_list[i]) + set_idx.append(i) + else: + set_results[i] = 0 + if set_keys: + put_results = self._put_batch_zero_copy_impl( + set_keys, set_ptrs, set_sizes + ) + for j, orig_i in enumerate(set_idx): + set_results[orig_i] = put_results[j] + ok = self._batch_postprocess( + set_results, is_set_operate=True, key_multiplier=key_multiplier + ) + else: + get_results = self._get_batch_zero_copy_impl( + key_strs, ptr_list, element_size_list + ) + ok = self._batch_postprocess( + get_results, is_set_operate=False, key_multiplier=key_multiplier + ) + results[transfer.name] = ok + return results + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> dict: + return self._batch_io_v2(transfers, is_set=False) + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> dict: + return self._batch_io_v2(transfers, is_set=True) + def _get_mha_split_heads_buffer_meta(self, keys, indices): ptr_list, element_size_list = ( self.mem_pool_host.get_split_heads_page_buffer_meta( @@ -536,7 +709,9 @@ def _batch_preprocess(self, keys, host_indices): else: return self._get_mha_buffer_meta(keys, host_indices) - def _batch_postprocess(self, results: List[int], is_set_operate=False): + def _batch_postprocess( + self, results: List[int], is_set_operate=False, key_multiplier=None + ): """ refer to https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/include/pybind_client.h for batch_get_into, results is Vector of integers, @@ -544,32 +719,26 @@ def _batch_postprocess(self, results: List[int], is_set_operate=False): for batch_put_from, results is Vector of integers, where each element is 0 on success, or a negative value on error """ - if self.is_mla_backend: - return [k_res == 0 if is_set_operate else k_res > 0 for k_res in results] - else: - if self.storage_config.should_split_heads: - kv_groups = [ - results[i : i + self.split_factor * 2] - for i in range(0, len(results), self.split_factor * 2) - ] - return [ - ( - all(res == 0 for res in kv_group) - if is_set_operate - else all(res > 0 for res in kv_group) - ) - for kv_group in kv_groups - ] + if key_multiplier is None: + if self.is_mla_backend: + key_multiplier = 1 else: - kv_pairs = zip(results[::2], results[1::2]) - return [ - ( - (k_res == 0 and v_res == 0) - if is_set_operate - else (k_res > 0 and v_res > 0) - ) - for k_res, v_res in kv_pairs - ] + key_multiplier = 2 + if self.storage_config.should_split_heads: + key_multiplier *= self.split_factor + + result_groups = [ + results[i : i + key_multiplier] + for i in range(0, len(results), key_multiplier) + ] + return [ + ( + all(res == 0 for res in group) + if is_set_operate + else all(res > 0 for res in group) + ) + for group in result_groups + ] def batch_get_v1( self, @@ -578,9 +747,7 @@ def batch_get_v1( extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: # Apply extra_backend_tag prefix if available - if self.extra_backend_tag is not None: - prefix = self.extra_backend_tag - keys = [f"{prefix}_{key}" for key in keys] + keys = self._tag_keys(keys) key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) @@ -605,9 +772,7 @@ def batch_set_v1( extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: # Apply extra_backend_tag prefix if available - if self.extra_backend_tag is not None: - prefix = self.extra_backend_tag - keys = [f"{prefix}_{key}" for key in keys] + keys = self._tag_keys(keys) key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) exist_result = self._batch_exist(key_strs) @@ -772,9 +937,7 @@ def batch_exists( self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None ) -> int: # Apply extra_backend_tag prefix if available - if self.extra_backend_tag is not None: - prefix = self.extra_backend_tag - keys = [f"{prefix}_{key}" for key in keys] + keys = self._tag_keys(keys) if self.is_mla_backend: query_keys = [f"{key}_{self.mla_suffix}_k" for key in keys] From e9e3c8d497b9fc9c34e0522aaa0fcf44a17f102f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 23 Mar 2026 14:59:31 +0800 Subject: [PATCH 26/39] fix bugs --- .../sglang/srt/mem_cache/memory_pool_host.py | 22 ++++++++++--------- .../storage/mooncake_store/mooncake_store.py | 12 +++++----- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d4488874af72..693908d55dee 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1112,6 +1112,10 @@ def __init__( conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv ] self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:] + self.temporal_state_elem_size = int(np.prod(self.temporal_state_shape)) + self.conv_state_elem_sizes = [ + int(np.prod(conv_shape)) for conv_shape in self.conv_state_shapes + ] self.conv_dtype = device_pool.mamba_cache.conv[0].dtype self.temporal_dtype = device_pool.mamba_cache.temporal.dtype self.dtype = self.conv_dtype @@ -1250,12 +1254,10 @@ def free(self, indices: torch.Tensor) -> int: return len(indices) def get_size_per_token(self): - conv_total_size = 0 - for conv_shape in self.conv_state_shapes: - conv_total_size += int(np.prod(conv_shape)) * self.conv_dtype.itemsize - temporal_size = ( - int(np.prod(self.temporal_state_shape)) * self.temporal_dtype.itemsize + conv_total_size = sum( + conv_elem_size * self.conv_dtype.itemsize for conv_elem_size in self.conv_state_elem_sizes ) + temporal_size = self.temporal_state_elem_size * self.temporal_dtype.itemsize return (conv_total_size + temporal_size) * self.num_mamba_layers def get_ksize_per_token(self): @@ -1514,16 +1516,16 @@ def get_page_buffer_meta(self, indices): self.page_size * self.num_mamba_layers * self.temporal_dtype.itemsize - * self.temporal_state_shape + * self.temporal_state_elem_size ) conv_element_sizes = [ ( self.page_size * self.num_mamba_layers * self.conv_dtype.itemsize - * conv_shape + * self.conv_state_elem_sizes[i] ) - for conv_shape in self.conv_state_shapes + for i in range(len(self.conv_state_shapes)) ] for i in range(0, len(indices), self.page_size): @@ -1533,7 +1535,7 @@ def get_page_buffer_meta(self, indices): temporal_base_ptr + indices[i] * self.num_mamba_layers - * self.temporal_state_shape + * self.temporal_state_elem_size * self.temporal_dtype.itemsize ) ptr_list.append(temporal_ptr) @@ -1543,7 +1545,7 @@ def get_page_buffer_meta(self, indices): conv_base_ptrs[j] + indices[i] * self.num_mamba_layers - * self.conv_state_shapes[j] + * self.conv_state_elem_sizes[j] * self.conv_dtype.itemsize ) ptr_list.append(conv_ptr) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 79bcc4bf2cff..9abedea86ce8 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -530,16 +530,16 @@ def _get_hybrid_page_component_keys( # - MAMBA : one temporal key + N conv keys per page # key_multiplier records how many component keys are generated per page, # so post-processing can fold component results back to page-level booleans. - name = str(transfer.name) + name = transfer.name component_keys = [] key_multiplier = 0 for page_key in page_keys: - if name == PoolName.INDEXER.value: + if name == PoolName.INDEXER: key_multiplier = 1 component_keys.append( f"{page_key}_{self.mla_suffix}_{PoolName.INDEXER}" ) - elif name == PoolName.MAMBA.value: + elif name == PoolName.MAMBA: base_key = f"{page_key}_{self.mha_suffix}" mamba_pool = getattr(self, "registered_pools", {}).get(PoolName.MAMBA) conv_num = len(getattr(mamba_pool, "conv_buffer", []) or []) @@ -564,8 +564,10 @@ def batch_exists_v2( def has_component(exist_results: List[int], page_idx: int, key_multiplier: int): # Group component existence by page, then require all components hit. + if key_multiplier <= 0: + return False result_groups = exist_results[ - (page_idx - 1) * key_multiplier : page_idx * key_multiplier + page_idx * key_multiplier : (page_idx + 1) * key_multiplier ] return all(result == 1 for result in result_groups) @@ -611,7 +613,7 @@ def _batch_io_v2(self, transfers: List[PoolTransfer], is_set: bool): page_size = getattr(host_pool, "page_size", 1) or 1 host_indices = transfer.host_indices assert len(keys) > 0 - assert len(keys) == host_indices // page_size + assert len(keys) == len(host_indices) // page_size ptr_list, element_size_list = host_pool.get_page_buffer_meta(host_indices) key_strs, key_multiplier = self._get_hybrid_page_component_keys( From ffdf17356aa264b20adad574e99858f2265df578 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 23 Mar 2026 17:21:42 +0800 Subject: [PATCH 27/39] fix bugs while running dsa --- python/sglang/srt/mem_cache/hiradix_cache.py | 14 ++++++-------- .../hybrid_cache/hybrid_cache_controller.py | 9 +++++++++ python/sglang/srt/mem_cache/memory_pool_host.py | 4 ++-- 3 files changed, 17 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index ebee9df69560..9e2ed55a4da4 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -713,7 +713,7 @@ def get_height(self, node: TreeNode): height += 1 return height - def _hybrid_pools_kw(self, host_indices: Optional[torch.Tensor]=None) -> dict: + def _get_extra_pools_kw(self) -> dict: if not isinstance( self.cache_controller, HybridCacheController ): @@ -721,8 +721,6 @@ def _hybrid_pools_kw(self, host_indices: Optional[torch.Tensor]=None) -> dict: if isinstance(self.kv_cache, NSATokenToKVPool): pool = PoolTransfer( name=PoolName.INDEXER, - host_indices=host_indices, - device_indices=None, hit_policy=PoolHitPolicy.ALL_PAGES, ) return {"extra_pools": [pool]} @@ -752,7 +750,7 @@ def clear_storage_backend(self) -> bool: return False def write_backup(self, node: TreeNode, write_back=False): - write_kw = self._hybrid_pools_kw() + write_kw = self._get_extra_pools_kw() host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, @@ -784,7 +782,7 @@ def write_backup_storage(self, node: TreeNode): else None ) - storage_kw = self._hybrid_pools_kw(host_indices=node.host_value) + storage_kw = self._get_extra_pools_kw() operation_id = self.cache_controller.write_storage( node.host_value, node.key, node.hash_value, prefix_keys, **storage_kw ) @@ -1155,14 +1153,14 @@ def load_back( device_indices = self.cache_controller.load( host_indices=host_indices, node_id=last_hit_node.id, - **self._hybrid_pools_kw(host_indices=host_indices), + **self._get_extra_pools_kw(), ) if device_indices is None: self.evict(EvictParams(num_tokens=len(host_indices))) device_indices = self.cache_controller.load( host_indices=host_indices, node_id=last_hit_node.id, - **self._hybrid_pools_kw(host_indices=host_indices), + **self._get_extra_pools_kw(), ) self.dec_lock_ref(ancester_node) if device_indices is None: @@ -1468,7 +1466,7 @@ def prefetch_from_storage( new_input_tokens, last_hash, prefix_keys, - **self._hybrid_pools_kw(host_indices=host_indices), + **self._get_extra_pools_kw(), ) self.ongoing_prefetch[req_id] = ( last_host_node, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 77817b1264ca..168197e79466 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -423,6 +423,7 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]: def _page_transfer(self, operation): # Transfer extra pools if operation.pool_transfers and not operation.is_terminated(): + self._resolve_shared_pool_transfers(operation) results = self.storage_backend.batch_get_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) @@ -432,12 +433,20 @@ def _page_transfer(self, operation): def _page_backup(self, operation): # Backup extra pools if operation.pool_transfers: + self._resolve_shared_pool_transfers(operation) results = self.storage_backend.batch_set_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) # Backup kv pools super()._page_backup(operation) + def _resolve_shared_pool_transfers(self, operation): + for transfer in operation.pool_transfers: + entry = self.mem_pool_host.entry_map.get(transfer.name) + if entry.share_indices_with_anchor: + transfer.keys = operation.hash_value + transfer.host_indices = operation.host_indices + def _sync_trailing_keys( self, pool_transfers: list[PoolTransfer], diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 693908d55dee..2e1b7638ac3e 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1766,7 +1766,7 @@ def __init__( requested_bytes / 1e9, layout, ) - self.kv_buffer = self.init_kv_buffer() + self.init_kv_buffer() self.lock = threading.RLock() self.clear() @@ -1778,7 +1778,7 @@ def get_size_per_token(self): def get_ksize_per_token(self): return self.get_size_per_token() - def _init_indexer_buffers(self): + def init_kv_buffer(self): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] self.index_k_device_ptrs = torch.tensor( [x.data_ptr() for x in self.device_pool.index_k_with_scale_buffer], From 1d2134efda81928b9c7dd9c8db7adcdb34cc2edf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 23 Mar 2026 17:24:57 +0800 Subject: [PATCH 28/39] fix lint --- python/sglang/srt/mem_cache/hiradix_cache.py | 6 ++---- .../srt/mem_cache/hybrid_cache/hybrid_cache_controller.py | 4 ++-- python/sglang/srt/mem_cache/memory_pool_host.py | 3 ++- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 9e2ed55a4da4..36099cb4c6c2 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -714,12 +714,10 @@ def get_height(self, node: TreeNode): return height def _get_extra_pools_kw(self) -> dict: - if not isinstance( - self.cache_controller, HybridCacheController - ): + if not isinstance(self.cache_controller, HybridCacheController): return {} if isinstance(self.kv_cache, NSATokenToKVPool): - pool = PoolTransfer( + pool = PoolTransfer( name=PoolName.INDEXER, hit_policy=PoolHitPolicy.ALL_PAGES, ) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 168197e79466..7fd9580b26fd 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -235,7 +235,7 @@ def write( extra_pools, alloc_host=True, kv_device_indices=device_indices, - kv_host_indices=host_indices + kv_host_indices=host_indices, ) if pool_transfers is None and extra_pools: self.mem_pool_host.free(host_indices) @@ -297,7 +297,7 @@ def load( extra_pools, alloc_host=False, kv_device_indices=device_indices, - kv_host_indices=host_indices + kv_host_indices=host_indices, ) if pool_transfers is None and extra_pools: if need_load_kv: diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 2e1b7638ac3e..d0e3b2908c9c 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1255,7 +1255,8 @@ def free(self, indices: torch.Tensor) -> int: def get_size_per_token(self): conv_total_size = sum( - conv_elem_size * self.conv_dtype.itemsize for conv_elem_size in self.conv_state_elem_sizes + conv_elem_size * self.conv_dtype.itemsize + for conv_elem_size in self.conv_state_elem_sizes ) temporal_size = self.temporal_state_elem_size * self.temporal_dtype.itemsize return (conv_total_size + temporal_size) * self.num_mamba_layers From 934874b200014e246c2f9db03c9803cf73db4088 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Tue, 24 Mar 2026 11:27:29 +0800 Subject: [PATCH 29/39] merge fix --- python/sglang/srt/managers/schedule_batch.py | 7 ------- python/sglang/srt/managers/schedule_policy.py | 2 +- test/registered/unit/managers/test_prefill_adder.py | 1 - 3 files changed, 1 insertion(+), 9 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 65cbc08d5263..9657a72b12e5 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -642,7 +642,6 @@ def __init__( self.last_node: Any = None self.last_host_node: Any = None self.host_hit_length = 0 - self.mamba_host_hit_length = 0 # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 # The node to lock until for swa radix tree lock ref @@ -823,10 +822,6 @@ def output_ids_through_stop(self) -> List[int]: return self.output_ids[: self.finished_len] return self.output_ids - @property - def should_init_load_back_host(self) -> bool: - return self.host_hit_length > 0 or self.mamba_host_hit_length > 0 - def pop_committed_kv_cache(self) -> int: """Return the length of committed KV cache and mark them as freed.""" assert ( @@ -920,14 +915,12 @@ def init_next_round_input( self.last_node, self.last_host_node, self.host_hit_length, - self.mamba_host_hit_length, self.mamba_branching_seqlen, ) = ( match_result.device_indices, match_result.last_device_node, match_result.last_host_node, match_result.host_hit_length, - match_result.mamba_host_hit_length, match_result.mamba_branching_seqlen, ) if match_result.cache_protected_len is not None: diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d80bc684a844..2cce30600ab2 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -772,7 +772,7 @@ def add_one_req( if total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN - if req.should_init_load_back_host: + if req.host_hit_length > 0: new_indices, req.last_node = self.tree_cache.init_load_back( InitLoadBackParams( last_host_node=req.last_host_node, diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index c0e4a0f01f71..e153d9d9c635 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -77,7 +77,6 @@ def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens) req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.finished.return_value = False - req.should_init_load_back_host = False return req def create_adder(self, running_batch, **kwargs): From 98048c530f0c9fcea07f000ed0b9233e08aa3bdf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Fri, 27 Mar 2026 13:50:11 +0800 Subject: [PATCH 30/39] fix --- .../storage/mooncake_store/mooncake_store.py | 84 +++++++++---------- 1 file changed, 38 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 9abedea86ce8..2281cdd4a116 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -5,7 +5,7 @@ import time import uuid from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Optional, Tuple import requests import torch @@ -501,17 +501,16 @@ def register_mem_pool_host(self, mem_pool_host: HostKVCache): self.gb_per_page = bytes_per_page / (1 << 30) def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): - # Keep a name->pool mapping so batch v2 can resolve PoolTransfer.name to - # the corresponding host pool implementation at runtime. - self.registered_pools[host_pool_name] = host_pool # KV anchor memory is already registered via register_mem_pool_host(). # v2 here only registers additional hybrid pools. if host_pool_name == PoolName.KV: return + # Keep a name->pool mapping so batch v2 can resolve PoolTransfer.name to + # the corresponding host pool implementation at runtime. + self.registered_pools[host_pool_name] = host_pool # Hybrid pools expose the tensors that Mooncake needs for zero-copy I/O. # The storage backend only depends on this accessor, not concrete fields. - assert getattr(host_pool, "get_hybrid_pool_buffer") is not None buf_list = host_pool.get_hybrid_pool_buffer() for buf in buf_list: super().register_buffer(buf) @@ -519,35 +518,31 @@ def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): def _tag_keys(self, keys: List[str]) -> List[str]: if self.extra_backend_tag is None: return keys - prefix = self.extra_backend_tag - return [f"{prefix}_{key}" for key in keys] + return [f"{ self.extra_backend_tag}_{key}" for key in keys] def _get_hybrid_page_component_keys( self, page_keys: List[str], transfer: PoolTransfer - ): + ) -> Tuple[List[str], int]: # A logical "page" may map to multiple physical objects in storage. # - INDEXER: one key per page # - MAMBA : one temporal key + N conv keys per page - # key_multiplier records how many component keys are generated per page, - # so post-processing can fold component results back to page-level booleans. + # key_multiplier records how many component keys are generated per page. name = transfer.name - component_keys = [] - key_multiplier = 0 - for page_key in page_keys: - if name == PoolName.INDEXER: - key_multiplier = 1 - component_keys.append( - f"{page_key}_{self.mla_suffix}_{PoolName.INDEXER}" - ) - elif name == PoolName.MAMBA: - base_key = f"{page_key}_{self.mha_suffix}" - mamba_pool = getattr(self, "registered_pools", {}).get(PoolName.MAMBA) - conv_num = len(getattr(mamba_pool, "conv_buffer", []) or []) - key_multiplier = conv_num + 1 - component_keys.extend( - [f"{base_key}_temporal"] - + [f"{base_key}_conv_{i}" for i in range(conv_num)] - ) + suffixes = [] + if name == PoolName.INDEXER: + suffixes = [f"_{self.mla_suffix}_{PoolName.INDEXER}"] + elif name == PoolName.MAMBA: + pools = getattr(self, "registered_pools", {}) + mamba_pool = pools.get(PoolName.MAMBA) + conv_num = len(getattr(mamba_pool, "conv_buffer", None) or []) + base_suffix = f"_{self.mha_suffix}" + suffixes = [f"{base_suffix}_temporal"] + [ + f"{base_suffix}_conv_{i}" for i in range(conv_num) + ] + key_multiplier = len(suffixes) + component_keys = [ + f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes + ] return component_keys, key_multiplier def batch_exists_v2( @@ -562,15 +557,6 @@ def batch_exists_v2( hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {} final_pages = kv_pages - def has_component(exist_results: List[int], page_idx: int, key_multiplier: int): - # Group component existence by page, then require all components hit. - if key_multiplier <= 0: - return False - result_groups = exist_results[ - page_idx * key_multiplier : (page_idx + 1) * key_multiplier - ] - return all(result == 1 for result in result_groups) - for transfer in pool_transfers or []: if final_pages == 0: break @@ -578,21 +564,27 @@ def has_component(exist_results: List[int], page_idx: int, key_multiplier: int): qkeys, transfer ) ex = self._batch_exist(component_keys) + if key_multiplier > 0: + page_exists = [ + all( + r == 1 + for r in ex[i * key_multiplier : (i + 1) * key_multiplier] + ) + for i in range(kv_pages) + ] + else: + page_exists = [False] * kv_pages boundary = 0 if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: - boundary = next( - ( - i - for i in range(kv_pages) - if not has_component(ex, i, key_multiplier) - ), - kv_pages, - ) + try: + boundary = page_exists.index(False) + except ValueError: + boundary = kv_pages elif transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES: trailing = max(1, len(transfer.keys) if transfer.keys else 1) for prefix_len in range(kv_pages, 0, -1): if all( - has_component(ex, i, key_multiplier) + page_exists[i] for i in range(max(0, prefix_len - trailing), prefix_len) ): boundary = prefix_len @@ -609,7 +601,7 @@ def _batch_io_v2(self, transfers: List[PoolTransfer], is_set: bool): results: dict = {} for transfer in transfers: host_pool = getattr(self, "registered_pools", {}).get(transfer.name) - keys = transfer.keys or [] + keys = transfer.keys page_size = getattr(host_pool, "page_size", 1) or 1 host_indices = transfer.host_indices assert len(keys) > 0 From aacc0e93eb4c6cae16699cc41f70f1767c0e1d64 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Fri, 27 Mar 2026 14:26:28 +0300 Subject: [PATCH 31/39] address comments Signed-off-by: Vladislav Nosivskoy --- .../storage/mooncake_store/mooncake_store.py | 37 ++++++------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 2281cdd4a116..ccda0d863411 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -615,36 +615,23 @@ def _batch_io_v2(self, transfers: List[PoolTransfer], is_set: bool): if is_set: exist_result = self._batch_exist(key_strs) - set_keys = [] - set_ptrs = [] - set_sizes = [] - set_idx = [] - set_results = [-1] * len(key_strs) - for i in range(len(key_strs)): - if exist_result[i] != 1: - set_keys.append(key_strs[i]) - set_ptrs.append(ptr_list[i]) - set_sizes.append(element_size_list[i]) - set_idx.append(i) - else: - set_results[i] = 0 - if set_keys: + io_results = [0 if state == 1 else -1 for state in exist_result] + missing_idx = [i for i, state in enumerate(exist_result) if state != 1] + if missing_idx: put_results = self._put_batch_zero_copy_impl( - set_keys, set_ptrs, set_sizes + [key_strs[i] for i in missing_idx], + [ptr_list[i] for i in missing_idx], + [element_size_list[i] for i in missing_idx], ) - for j, orig_i in enumerate(set_idx): - set_results[orig_i] = put_results[j] - ok = self._batch_postprocess( - set_results, is_set_operate=True, key_multiplier=key_multiplier - ) + for i, res in zip(missing_idx, put_results): + io_results[i] = res else: - get_results = self._get_batch_zero_copy_impl( + io_results = self._get_batch_zero_copy_impl( key_strs, ptr_list, element_size_list ) - ok = self._batch_postprocess( - get_results, is_set_operate=False, key_multiplier=key_multiplier - ) - results[transfer.name] = ok + results[transfer.name] = self._batch_postprocess( + io_results, is_set_operate=is_set, key_multiplier=key_multiplier + ) return results def batch_get_v2( From 1f506f91b7a484d6446b44abd43ca7b3d27e53b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 30 Mar 2026 14:15:24 +0800 Subject: [PATCH 32/39] fix metrics --- python/sglang/srt/mem_cache/hiradix_cache.py | 1 + python/sglang/srt/mem_cache/memory_pool_host.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index f482d6ee25ab..b01dff3890b4 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -116,6 +116,7 @@ def layer_mapper(layer_id: int): ), ] ) + radix_cache.full_kv_pool_host = mla_host radix_cache.token_to_kv_pool_host = host_pool_group radix_cache.cache_controller = HybridCacheController( params.token_to_kv_pool_allocator, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index ccbdbdd6ff29..67951da87a5d 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1621,6 +1621,9 @@ def clear(self) -> None: for entry in self.entries: entry.host_pool.clear() + def available_size(self): + return self.anchor_entry.host_pool.available_size() + def alloc(self, need_size: int) -> Optional[torch.Tensor]: return self.anchor_entry.host_pool.alloc(need_size) From 9bacae6087195dc9d2722f821351db228b39024d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Tue, 31 Mar 2026 14:47:43 +0800 Subject: [PATCH 33/39] Reduce modifications in the HiRadix cache. --- python/sglang/srt/mem_cache/hiradix_cache.py | 105 +++--------------- .../srt/mem_cache/hybrid_cache/__init__.py | 0 .../hybrid_cache/hybrid_pool_assembler.py | 100 +++++++++++++++++ 3 files changed, 113 insertions(+), 92 deletions(-) create mode 100644 python/sglang/srt/mem_cache/hybrid_cache/__init__.py create mode 100644 python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 92610ef00b5f..4d0ad2a478b5 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -34,17 +34,17 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) +from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + build_nsa_hybrid_stack, +) from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, MLATokenToKVPool, NSATokenToKVPool, ) from sglang.srt.mem_cache.memory_pool_host import ( - HostPoolGroup, MHATokenToKVPoolHost, MLATokenToKVPoolHost, - NSAIndexerPoolHost, - PoolEntry, ) from sglang.srt.mem_cache.radix_cache import ( RadixCache, @@ -63,84 +63,6 @@ logger = logging.getLogger(__name__) -def _build_nsa_hybrid_stack( - radix_cache: "HiRadixCache", - params: "CacheInitParams", - server_args: "ServerArgs", - *, - extra_config: dict, - prefetch_threshold: int, - enable_storage_metrics: bool, - load_cache_event, -) -> None: - """HostPoolGroup (KV + indexer) + HybridCacheController for NSA (DSA).""" - kv = radix_cache.kv_cache - mla_host = MLATokenToKVPoolHost( - kv, - server_args.hicache_ratio, - server_args.hicache_size, - radix_cache.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - override_kv_cache_dim=kv.kv_cache_dim, - ) - indexer_host = NSAIndexerPoolHost( - kv, - mla_host, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) - layer_num = kv.layer_num - - def layer_mapper(layer_id: int): - if 0 <= layer_id < layer_num: - return layer_id - return None - - host_pool_group = HostPoolGroup( - [ - PoolEntry( - name=PoolName.KV, - host_pool=mla_host, - device_pool=kv, - layer_mapper=layer_mapper, - is_primary_index_anchor=True, - ), - PoolEntry( - name=PoolName.INDEXER, - host_pool=indexer_host, - device_pool=kv, - layer_mapper=layer_mapper, - share_indices_with_anchor=True, - ), - ] - ) - radix_cache.full_kv_pool_host = mla_host - radix_cache.token_to_kv_pool_host = host_pool_group - radix_cache.cache_controller = HybridCacheController( - params.token_to_kv_pool_allocator, - host_pool_group, - radix_cache.page_size, - radix_cache.tp_group, - load_cache_event=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=radix_cache.pp_rank, - pp_size=radix_cache.pp_size, - transfer_layer_num=layer_num, - enable_storage_metrics=enable_storage_metrics, - ) - logger.info( - "NSA hierarchical cache: HostPoolGroup(KV + INDEXER), HybridCacheController, " - "transfer_layer_num=%s", - layer_num, - ) - - class HiRadixCache(RadixCache): def __init__(self, params: CacheInitParams, server_args: ServerArgs): @@ -159,7 +81,7 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): allocator_type=server_args.hicache_storage_backend, ) elif isinstance(self.kv_cache, NSATokenToKVPool): - # Filled by _build_nsa_hybrid_stack after storage extra_config is parsed. + # Filled by build_nsa_hybrid_stack after storage extra_config is parsed. self.token_to_kv_pool_host = None elif isinstance(self.kv_cache, MLATokenToKVPool): self.token_to_kv_pool_host = MLATokenToKVPoolHost( @@ -198,7 +120,7 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): self.load_cache_event = threading.Event() if isinstance(self.kv_cache, NSATokenToKVPool): - _build_nsa_hybrid_stack( + build_nsa_hybrid_stack( self, params, server_args, @@ -703,7 +625,7 @@ def get_height(self, node: TreeNode): height += 1 return height - def _get_extra_pools_kw(self) -> dict: + def _get_extra_pools(self) -> dict: if not isinstance(self.cache_controller, HybridCacheController): return {} if isinstance(self.kv_cache, NSATokenToKVPool): @@ -738,18 +660,17 @@ def clear_storage_backend(self) -> bool: return False def write_backup(self, node: TreeNode, write_back=False): - write_kw = self._get_extra_pools_kw() host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, - **write_kw, + **self._get_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, - **write_kw, + **self._get_extra_pools(), ) if host_indices is not None: node.host_value = host_indices @@ -770,9 +691,9 @@ def write_backup_storage(self, node: TreeNode): else None ) - storage_kw = self._get_extra_pools_kw() operation_id = self.cache_controller.write_storage( - node.host_value, node.key, node.hash_value, prefix_keys, **storage_kw + node.host_value, node.key, node.hash_value, prefix_keys, + **self._get_extra_pools(), ) self.ongoing_backup[operation_id] = node node.protect_host() @@ -1141,14 +1062,14 @@ def load_back( device_indices = self.cache_controller.load( host_indices=host_indices, node_id=last_hit_node.id, - **self._get_extra_pools_kw(), + **self._get_extra_pools(), ) if device_indices is None: self.evict(EvictParams(num_tokens=len(host_indices))) device_indices = self.cache_controller.load( host_indices=host_indices, node_id=last_hit_node.id, - **self._get_extra_pools_kw(), + **self._get_extra_pools(), ) self.dec_lock_ref(ancester_node) if device_indices is None: @@ -1454,7 +1375,7 @@ def prefetch_from_storage( new_input_tokens, last_hash, prefix_keys, - **self._get_extra_pools_kw(), + **self._extra_pools_kw(), ) self.ongoing_prefetch[req_id] = ( last_host_node, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/__init__.py b/python/sglang/srt/mem_cache/hybrid_cache/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py new file mode 100644 index 000000000000..8f1bf8d55a22 --- /dev/null +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from sglang.srt.mem_cache.hicache_storage import PoolName +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) +from sglang.srt.mem_cache.memory_pool_host import ( + HostPoolGroup, + MLATokenToKVPoolHost, + NSAIndexerPoolHost, + PoolEntry, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.hiradix_cache import HiRadixCache + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +def build_nsa_hybrid_stack( + radix_cache: "HiRadixCache", + params: "CacheInitParams", + server_args: "ServerArgs", + *, + extra_config: dict, + prefetch_threshold: int, + enable_storage_metrics: bool, + load_cache_event, +) -> None: + """HostPoolGroup (KV + indexer) + HybridCacheController for NSA (DSA).""" + kv = radix_cache.kv_cache + mla_host = MLATokenToKVPoolHost( + kv, + server_args.hicache_ratio, + server_args.hicache_size, + radix_cache.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + override_kv_cache_dim=kv.kv_cache_dim, + ) + indexer_host = NSAIndexerPoolHost( + kv, + mla_host, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + layer_num = kv.layer_num + + def layer_mapper(layer_id: int): + if 0 <= layer_id < layer_num: + return layer_id + return None + + host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=mla_host, + device_pool=kv, + layer_mapper=layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.INDEXER, + host_pool=indexer_host, + device_pool=kv, + layer_mapper=layer_mapper, + share_indices_with_anchor=True, + ), + ] + ) + radix_cache.full_kv_pool_host = mla_host + radix_cache.token_to_kv_pool_host = host_pool_group + radix_cache.cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + radix_cache.page_size, + radix_cache.tp_group, + load_cache_event=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=radix_cache.pp_rank, + pp_size=radix_cache.pp_size, + transfer_layer_num=layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + logger.info( + "Hybrid hierarchical cache: HostPoolGroup(KV + INDEXER), HybridCacheController, " + "transfer_layer_num=%s", + layer_num, + ) From 75b15b16e981ef2db3cce5f27f7c660b63e7ac50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Tue, 31 Mar 2026 14:51:24 +0800 Subject: [PATCH 34/39] fix --- python/sglang/srt/mem_cache/hiradix_cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 4d0ad2a478b5..6f66d37c0c71 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1375,7 +1375,7 @@ def prefetch_from_storage( new_input_tokens, last_hash, prefix_keys, - **self._extra_pools_kw(), + **self._get_extra_pools(), ) self.ongoing_prefetch[req_id] = ( last_host_node, From 8dd2287e587462d79cedf8f5f039ea24ea20f9b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Tue, 31 Mar 2026 14:55:16 +0800 Subject: [PATCH 35/39] fix --- python/sglang/srt/mem_cache/hybrid_cache/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 python/sglang/srt/mem_cache/hybrid_cache/__init__.py diff --git a/python/sglang/srt/mem_cache/hybrid_cache/__init__.py b/python/sglang/srt/mem_cache/hybrid_cache/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 From 095baf314d024105211104c7e34b7d650926c90a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Fri, 10 Apr 2026 14:37:52 +0800 Subject: [PATCH 36/39] fix --- python/sglang/srt/mem_cache/hiradix_cache.py | 5 +- .../hybrid_cache/hybrid_pool_assembler.py | 131 +++++++++--------- .../sglang/srt/mem_cache/memory_pool_host.py | 13 +- 3 files changed, 78 insertions(+), 71 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b2b5cc364224..bc447aa229ef 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -682,7 +682,10 @@ def write_backup_storage(self, node: TreeNode): ) operation_id = self.cache_controller.write_storage( - node.host_value, node.key, node.hash_value, prefix_keys, + node.host_value, + node.key, + node.hash_value, + prefix_keys, **self._get_extra_pools(), ) self.ongoing_backup[operation_id] = node diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 8f1bf8d55a22..4bca0145c8cc 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -33,68 +33,73 @@ def build_nsa_hybrid_stack( load_cache_event, ) -> None: """HostPoolGroup (KV + indexer) + HybridCacheController for NSA (DSA).""" - kv = radix_cache.kv_cache - mla_host = MLATokenToKVPoolHost( - kv, - server_args.hicache_ratio, - server_args.hicache_size, - radix_cache.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - override_kv_cache_dim=kv.kv_cache_dim, - ) - indexer_host = NSAIndexerPoolHost( - kv, - mla_host, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) - layer_num = kv.layer_num + try: + kv = radix_cache.kv_cache + mla_host = MLATokenToKVPoolHost( + kv, + server_args.hicache_ratio, + server_args.hicache_size, + radix_cache.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + override_kv_cache_dim=kv.kv_cache_dim, + ) + indexer_host = NSAIndexerPoolHost( + kv, + mla_host, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + layer_num = kv.layer_num - def layer_mapper(layer_id: int): - if 0 <= layer_id < layer_num: - return layer_id - return None + def layer_mapper(layer_id: int): + if 0 <= layer_id < layer_num: + return layer_id + return None - host_pool_group = HostPoolGroup( - [ - PoolEntry( - name=PoolName.KV, - host_pool=mla_host, - device_pool=kv, - layer_mapper=layer_mapper, - is_primary_index_anchor=True, - ), - PoolEntry( - name=PoolName.INDEXER, - host_pool=indexer_host, - device_pool=kv, - layer_mapper=layer_mapper, - share_indices_with_anchor=True, - ), - ] - ) - radix_cache.full_kv_pool_host = mla_host - radix_cache.token_to_kv_pool_host = host_pool_group - radix_cache.cache_controller = HybridCacheController( - params.token_to_kv_pool_allocator, - host_pool_group, - radix_cache.page_size, - radix_cache.tp_group, - load_cache_event=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=radix_cache.pp_rank, - pp_size=radix_cache.pp_size, - transfer_layer_num=layer_num, - enable_storage_metrics=enable_storage_metrics, - ) - logger.info( - "Hybrid hierarchical cache: HostPoolGroup(KV + INDEXER), HybridCacheController, " - "transfer_layer_num=%s", - layer_num, - ) + host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=mla_host, + device_pool=kv, + layer_mapper=layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.INDEXER, + host_pool=indexer_host, + device_pool=kv, + layer_mapper=layer_mapper, + share_indices_with_anchor=True, + ), + ] + ) + cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + radix_cache.page_size, + radix_cache.tp_group, + load_cache_event=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=radix_cache.pp_rank, + pp_size=radix_cache.pp_size, + transfer_layer_num=layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + radix_cache.full_kv_pool_host = mla_host + radix_cache.token_to_kv_pool_host = host_pool_group + radix_cache.cache_controller = cache_controller + logger.info( + "Hybrid hierarchical cache: HostPoolGroup(KV + INDEXER), HybridCacheController, " + "transfer_layer_num=%s", + layer_num, + ) + except Exception as e: + logger.error(f"build_nsa_hybrid_stack failed: {e}", exc_info=True) + raise diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index a0ee5620ea5f..5b3e0a9424c5 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -64,6 +64,9 @@ logger = logging.getLogger(__name__) +# Host RAM to leave free when sizing HiCache pools (OS, other processes). +HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3) + def synchronized(func): @wraps(func) @@ -187,9 +190,7 @@ def __init__( # Verify there is enough available host memory. host_mem = psutil.virtual_memory() requested_bytes = self.size * self.size_per_token - # preserve at least 10GB for other usage - ten_gb = 10 * (1024**3) - available_bytes = host_mem.available - ten_gb + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES if requested_bytes > available_bytes: raise ValueError( f"Not enough host memory available. Requesting " @@ -1236,8 +1237,7 @@ def __init__( host_mem = psutil.virtual_memory() requested_bytes = self.size * self.size_per_token - ten_gb = 10 * (1024**3) - available_bytes = host_mem.available - ten_gb + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES if requested_bytes > available_bytes: raise ValueError( f"Not enough host memory available. Requesting " @@ -1858,8 +1858,7 @@ def __init__( buf_elem_size = self.page_num * self.layer_num * self.indexer_page_stride_size requested_bytes = buf_elem_size * self.indexer_dtype.itemsize host_mem = psutil.virtual_memory() - ten_gb = 10 * (1024**3) - available_bytes = host_mem.available - ten_gb + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES if requested_bytes > available_bytes: raise ValueError( f"Not enough host memory for NSA indexer hierarchical cache. " From 586b33716887e92167f0d2e99cd63d19e34cc02e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Sun, 12 Apr 2026 14:24:54 +0800 Subject: [PATCH 37/39] opt --- python/sglang/srt/mem_cache/hiradix_cache.py | 28 +++++++++----------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 5440c9231357..f9ecc0e6909b 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -338,21 +338,13 @@ def attach_storage_backend( ) try: - if isinstance(self.cache_controller, HybridCacheController): - 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.cache_controller.mem_pool_host.entries, - ) - else: - 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, - ) + 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, + **self._get_hybrid_storage_attach_kwargs(), + ) except Exception as e: logger.exception( f"Failed to attach storage backend '{storage_backend}': {e}" @@ -629,6 +621,12 @@ def _get_extra_pools(self) -> dict: else: return {} + def _get_hybrid_storage_attach_kwargs(self) -> dict: + """Extra kwargs for attach_storage_backend when controller is HybridCacheController.""" + if isinstance(self.cache_controller, HybridCacheController): + return {"host_pools": self.cache_controller.mem_pool_host.entries} + return {} + def clear_storage_backend(self) -> bool: if self.enable_storage: try: From 973110de73f3448e877af98bdcd1de49b55e17b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Sun, 12 Apr 2026 14:40:52 +0800 Subject: [PATCH 38/39] build_mamba_hybrid_stack --- .../srt/mem_cache/hi_mamba_radix_cache.py | 97 ++------------- .../hybrid_cache/hybrid_pool_assembler.py | 113 +++++++++++++++++- 2 files changed, 119 insertions(+), 91 deletions(-) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 10e63c7f8087..0b9c8f680cb6 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -24,9 +24,11 @@ ) from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( - HybridCacheController, PrefetchOperation, ) +from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + build_mamba_hybrid_stack, +) from sglang.srt.mem_cache.mamba_radix_cache import ( LRUList, MambaRadixCache, @@ -34,13 +36,6 @@ get_last_access_time, ) from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool -from sglang.srt.mem_cache.memory_pool_host import ( - HostPoolGroup, - MambaPoolHost, - MHATokenToKVPoolHost, - MLATokenToKVPoolHost, - PoolEntry, -) from sglang.srt.mem_cache.radix_cache import ( RadixKey, compute_node_hash_values, @@ -115,66 +110,6 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): ) self.kvcache = self.hybrid_kv_cache.full_kv_pool - kv_host_pool_cls = ( - MLATokenToKVPoolHost - if self.hybrid_kv_cache.use_mla - else MHATokenToKVPoolHost - ) - self.full_kv_pool_host = kv_host_pool_cls( - 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.mamba_pool_host = MambaPoolHost( - params.req_to_token_pool.mamba_pool, - server_args.hicache_ratio, - server_args.hicache_size, - allocator_type=server_args.hicache_storage_backend, - layout=server_args.hicache_mem_layout, - ) - - full_layer_ids = sorted( - self.hybrid_kv_cache.full_attention_layer_id_mapping.keys() - ) - mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) - self.transfer_layer_num = len(set(full_layer_ids) | set(mamba_layer_ids)) - - full_layer_mapping = dict(self.hybrid_kv_cache.full_attention_layer_id_mapping) - mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) - transfer_layer_num = self.transfer_layer_num - - def kv_layer_mapper(layer_id: int) -> Optional[int]: - if not 0 <= layer_id < transfer_layer_num: - return None - return full_layer_mapping.get(layer_id) - - def mamba_layer_mapper(layer_id: int) -> Optional[int]: - if not 0 <= layer_id < transfer_layer_num: - return None - return mamba_layer_mapping.get(layer_id) - - self.host_pool_group = HostPoolGroup( - [ - PoolEntry( - name=PoolName.KV, - host_pool=self.full_kv_pool_host, - device_pool=self.kvcache, - layer_mapper=kv_layer_mapper, - is_primary_index_anchor=True, - ), - PoolEntry( - name=PoolName.MAMBA, - host_pool=self.mamba_pool_host, - device_pool=params.req_to_token_pool.mamba_pool, - layer_mapper=mamba_layer_mapper, - host_evict_fn=self.evict_mamba_host, - device_evict_fn=self.evict_mamba, - ), - ] - ) self.tp_group = params.tp_cache_group self.tp_world_size = ( @@ -200,27 +135,13 @@ def mamba_layer_mapper(layer_id: int) -> Optional[int]: self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy self.load_cache_event = threading.Event() - self.cache_controller = HybridCacheController( - params.token_to_kv_pool_allocator, - self.host_pool_group, - 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, + build_mamba_hybrid_stack( + self, + params, + server_args, + extra_config=extra_config, 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, - transfer_layer_num=self.transfer_layer_num, - ) - params.req_to_token_pool.register_layer_transfer_counter( - self.cache_controller.layer_done_counter - ) - self.hybrid_kv_cache.register_layer_transfer_counter( - self.cache_controller.layer_done_counter + load_cache_event=self.load_cache_event, ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 4bca0145c8cc..99198dfece94 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from sglang.srt.mem_cache.hicache_storage import PoolName from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( @@ -9,6 +9,8 @@ ) from sglang.srt.mem_cache.memory_pool_host import ( HostPoolGroup, + MambaPoolHost, + MHATokenToKVPoolHost, MLATokenToKVPoolHost, NSAIndexerPoolHost, PoolEntry, @@ -16,6 +18,7 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache from sglang.srt.mem_cache.hiradix_cache import HiRadixCache from sglang.srt.server_args import ServerArgs @@ -100,6 +103,110 @@ def layer_mapper(layer_id: int): "transfer_layer_num=%s", layer_num, ) - except Exception as e: - logger.error(f"build_nsa_hybrid_stack failed: {e}", exc_info=True) + except Exception: + logger.exception("build_nsa_hybrid_stack failed") + raise + + +def build_mamba_hybrid_stack( + mamba_cache: "HiMambaRadixCache", + params: "CacheInitParams", + server_args: "ServerArgs", + *, + extra_config: dict, + prefetch_threshold: int, + load_cache_event, + enable_storage_metrics: bool = False, +) -> None: + """HostPoolGroup (KV + Mamba) + HybridCacheController for hybrid SSM models.""" + try: + hybrid_kv = mamba_cache.hybrid_kv_cache + kvcache = mamba_cache.kvcache + kv_host_pool_cls = ( + MLATokenToKVPoolHost if hybrid_kv.use_mla else MHATokenToKVPoolHost + ) + full_kv_pool_host = kv_host_pool_cls( + kvcache, + server_args.hicache_ratio, + server_args.hicache_size, + params.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + mamba_pool_host = MambaPoolHost( + params.req_to_token_pool.mamba_pool, + server_args.hicache_ratio, + server_args.hicache_size, + allocator_type=server_args.hicache_storage_backend, + layout=server_args.hicache_mem_layout, + ) + + full_layer_ids = sorted(hybrid_kv.full_attention_layer_id_mapping.keys()) + mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) + transfer_layer_num = len(set(full_layer_ids) | set(mamba_layer_ids)) + full_layer_mapping = dict(hybrid_kv.full_attention_layer_id_mapping) + mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) + + def kv_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: + return None + return full_layer_mapping.get(layer_id) + + def mamba_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: + return None + return mamba_layer_mapping.get(layer_id) + + host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=full_kv_pool_host, + device_pool=kvcache, + layer_mapper=kv_layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.MAMBA, + host_pool=mamba_pool_host, + device_pool=params.req_to_token_pool.mamba_pool, + layer_mapper=mamba_layer_mapper, + host_evict_fn=mamba_cache.evict_mamba_host, + device_evict_fn=mamba_cache.evict_mamba, + ), + ] + ) + cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + params.page_size, + params.tp_cache_group, + load_cache_event=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, + transfer_layer_num=transfer_layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + mamba_cache.full_kv_pool_host = full_kv_pool_host + mamba_cache.mamba_pool_host = mamba_pool_host + mamba_cache.transfer_layer_num = transfer_layer_num + mamba_cache.host_pool_group = host_pool_group + mamba_cache.cache_controller = cache_controller + params.req_to_token_pool.register_layer_transfer_counter( + cache_controller.layer_done_counter + ) + hybrid_kv.register_layer_transfer_counter(cache_controller.layer_done_counter) + logger.info( + "Hybrid hierarchical cache: HostPoolGroup(KV + MAMBA), HybridCacheController, " + "transfer_layer_num=%s", + transfer_layer_num, + ) + except Exception: + logger.exception("build_mamba_hybrid_stack failed") raise From fccef9ed356f8d8ab7b8b6c488c0fbbd4f194faa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9F=E6=B5=B7?= Date: Mon, 13 Apr 2026 11:07:56 +0800 Subject: [PATCH 39/39] fix ci --- .../sglang/srt/mem_cache/memory_pool_host.py | 2 +- .../unit/mem_cache/test_nsa_pool_host_unit.py | 24 +++++++++++++++---- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 5b3e0a9424c5..e5fd2da76f97 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -2023,7 +2023,7 @@ def backup_from_device_all_layer( if self.layout == "layer_first": transfer_kv_direct( src_layers=device_pool.index_k_with_scale_buffer, - dst_layers=self.index_k_with_scale_buffer, + dst_layers=self.index_k_data_refs, src_indices=device_page_indices, dst_indices=host_page_indices, page_size=1, diff --git a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py index f0462e267a87..2602e196a195 100644 --- a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py @@ -5,7 +5,8 @@ from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool from sglang.srt.mem_cache.memory_pool_host import ( ALLOC_MEMORY_FUNCS, - NSATokenToKVPoolHost, + MLATokenToKVPoolHost, + NSAIndexerPoolHost, alloc_with_pin_memory, ) from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu @@ -58,7 +59,7 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): if pin_memory: ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory try: - host_pool = NSATokenToKVPoolHost( + mla_host = MLATokenToKVPoolHost( device_pool=device_pool, host_to_device_ratio=2.0, host_size=0, @@ -66,6 +67,16 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): layout="layer_first", pin_memory=pin_memory, device="cpu", + allocator_type="default", + override_kv_cache_dim=device_pool.kv_cache_dim, + ) + indexer_host = NSAIndexerPoolHost( + device_pool=device_pool, + anchor_host=mla_host, + layout="layer_first", + pin_memory=pin_memory, + device="cpu", + allocator_type="default", ) finally: ALLOC_MEMORY_FUNCS["cuda"] = original_alloc @@ -97,7 +108,10 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): device="cuda" if io_backend == "kernel" else "cpu", ) - host_pool.backup_from_device_all_layer( + mla_host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend + ) + indexer_host.backup_from_device_all_layer( device_pool, host_indices, device_indices, io_backend ) @@ -105,14 +119,14 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): for host_page, device_page in zip( host_pages.tolist(), device_pages.tolist() ): - got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu() + got = indexer_host.index_k_with_scale_buffer[layer_id][host_page].cpu() expected = device_pool.index_k_with_scale_buffer[layer_id][ device_page ].cpu() self.assertTrue(torch.equal(got, expected)) host_start = host_page * page_size device_start = device_page * page_size - got_kv = host_pool.kv_buffer[layer_id][ + got_kv = mla_host.kv_buffer[layer_id][ host_start : host_start + page_size ].cpu() expected_kv = device_pool.kv_buffer[layer_id][