[UnifiedRadixTree]: Support HiCache Framework for UnifiedRadixTree (#23316)
Co-authored-by: JINZ <1023553676@qq.com> Co-authored-by: diemchai <diemchai@tencent.com>
This commit is contained in:
co-authored by
JINZ
diemchai
parent
e37f46fcf7
commit
c0f5950636
@@ -892,6 +892,26 @@ class Scheduler(
|
||||
|
||||
logger.info("Using experimental C++ radix tree implementation.")
|
||||
self.tree_cache = RadixCacheCpp(params=params, server_args=server_args)
|
||||
elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get():
|
||||
from sglang.srt.mem_cache.unified_cache_components import (
|
||||
ComponentType,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_radix_cache import (
|
||||
UnifiedRadixCache,
|
||||
)
|
||||
|
||||
tree_components = [ComponentType.FULL]
|
||||
if self.is_hybrid_swa or self.is_hybrid_ssm:
|
||||
tree_components.append(
|
||||
ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA
|
||||
)
|
||||
params.tree_components = tuple(tree_components)
|
||||
self.tree_cache = UnifiedRadixCache(params)
|
||||
if self.enable_hierarchical_cache:
|
||||
self.tree_cache.init_hicache(server_args, params)
|
||||
self.tp_worker.register_hicache_layer_transfer_counter(
|
||||
self.tree_cache.cache_controller.layer_done_counter
|
||||
)
|
||||
elif self.enable_hierarchical_cache:
|
||||
if self.is_hybrid_ssm:
|
||||
from sglang.srt.mem_cache.hi_mamba_radix_cache import (
|
||||
@@ -910,21 +930,6 @@ class Scheduler(
|
||||
self.tp_worker.register_hicache_layer_transfer_counter(
|
||||
self.tree_cache.cache_controller.layer_done_counter
|
||||
)
|
||||
elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get():
|
||||
from sglang.srt.mem_cache.unified_cache_components import (
|
||||
ComponentType,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_radix_cache import (
|
||||
UnifiedRadixCache,
|
||||
)
|
||||
|
||||
tree_components = [ComponentType.FULL]
|
||||
if self.is_hybrid_swa or self.is_hybrid_ssm:
|
||||
tree_components.append(
|
||||
ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA
|
||||
)
|
||||
params.tree_components = tuple(tree_components)
|
||||
self.tree_cache = UnifiedRadixCache(params)
|
||||
elif self.is_hybrid_swa:
|
||||
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
|
||||
|
||||
|
||||
@@ -74,7 +74,7 @@ class InsertResult:
|
||||
class EvictParams:
|
||||
"""Unified parameters for evict across different cache types"""
|
||||
|
||||
num_tokens: int
|
||||
num_tokens: int = 0
|
||||
swa_num_tokens: int = 0
|
||||
mamba_num: int = 0
|
||||
|
||||
|
||||
@@ -64,6 +64,7 @@ class PoolTransfer:
|
||||
|
||||
device<->host path : host_indices + device_indices
|
||||
host<->storage path: host_indices + keys
|
||||
nodes_to_load : evicted nodes this transfer covers
|
||||
"""
|
||||
|
||||
name: PoolName
|
||||
@@ -71,6 +72,7 @@ class PoolTransfer:
|
||||
device_indices: Optional[torch.Tensor] = None
|
||||
keys: Optional[List[str]] = None
|
||||
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
|
||||
nodes_to_load: Optional[List[Any]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Callable, Optional
|
||||
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName
|
||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||
HybridCacheController,
|
||||
)
|
||||
@@ -325,7 +325,11 @@ def attach_hybrid_pool_to_unified_cache(
|
||||
) -> None:
|
||||
"""Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache."""
|
||||
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
|
||||
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MLATokenToKVPool
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
HybridLinearKVPool,
|
||||
MLATokenToKVPool,
|
||||
NSATokenToKVPool,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_cache_components import ComponentType
|
||||
|
||||
try:
|
||||
@@ -345,6 +349,7 @@ def attach_hybrid_pool_to_unified_cache(
|
||||
}, "Non-hybrid KV pool currently only supports FULL-only UnifiedRadixCache."
|
||||
|
||||
mamba_stack = isinstance(kvcache, HybridLinearKVPool)
|
||||
nsa_stack = isinstance(kvcache, NSATokenToKVPool)
|
||||
if mamba_stack:
|
||||
full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping)
|
||||
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
|
||||
@@ -381,6 +386,45 @@ def attach_hybrid_pool_to_unified_cache(
|
||||
cache_controller.layer_done_counter
|
||||
)
|
||||
transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping)
|
||||
elif nsa_stack:
|
||||
full_layer_mapping = {
|
||||
layer_id: layer_id for layer_id in range(full_kv_pool.layer_num)
|
||||
}
|
||||
host_pool_group, cache_controller = build_shared_anchor_stack(
|
||||
params=params,
|
||||
server_args=server_args,
|
||||
kv_pool=full_kv_pool,
|
||||
shared_pool_name=PoolName.INDEXER,
|
||||
full_layer_mapping=full_layer_mapping,
|
||||
page_size=cache.page_size,
|
||||
tp_group=params.tp_cache_group,
|
||||
load_cache_event=load_cache_event,
|
||||
storage_backend=None,
|
||||
use_mla=use_mla,
|
||||
shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost(
|
||||
full_kv_pool,
|
||||
kv_host_pool,
|
||||
server_args.hicache_mem_layout,
|
||||
allocator_type=server_args.hicache_storage_backend,
|
||||
),
|
||||
pp_rank=params.pp_rank,
|
||||
pp_size=params.pp_size,
|
||||
attn_cp_rank=params.attn_cp_rank,
|
||||
attn_cp_size=params.attn_cp_size,
|
||||
)
|
||||
cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV)
|
||||
cache.host_pool_group = host_pool_group
|
||||
cache.cache_controller = cache_controller
|
||||
# Register the NSA indexer pool as sharing anchor-KV indices so
|
||||
# HiCache backup/load emits its PoolTransfer together with KV.
|
||||
cache.register_hicache_anchor_kv_shared_indices_pool(
|
||||
PoolName.INDEXER,
|
||||
hit_policy=PoolHitPolicy.ALL_PAGES,
|
||||
)
|
||||
cache.components[ComponentType.FULL]._full_kv_pool_host = (
|
||||
cache.full_kv_pool_host
|
||||
)
|
||||
transfer_layer_num = len(full_layer_mapping)
|
||||
else:
|
||||
full_layer_mapping = {
|
||||
layer_id: layer_id for layer_id in range(full_kv_pool.layer_num)
|
||||
@@ -414,7 +458,7 @@ def attach_hybrid_pool_to_unified_cache(
|
||||
|
||||
logger.info(
|
||||
"Attached hybrid pool stack to UnifiedRadixCache: pools=%s, transfer_layer_num=%s",
|
||||
"KV + MAMBA" if mamba_stack else "KV",
|
||||
"KV + MAMBA" if mamba_stack else "KV + INDEXER" if nsa_stack else "KV",
|
||||
transfer_layer_num,
|
||||
)
|
||||
except Exception:
|
||||
|
||||
@@ -4,8 +4,10 @@ from sglang.srt.mem_cache.unified_cache_components.swa_component import SWACompo
|
||||
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
|
||||
_NUM_COMPONENT_TYPES,
|
||||
BASE_COMPONENT_TYPE,
|
||||
CacheTransferPhase,
|
||||
ComponentData,
|
||||
ComponentType,
|
||||
EvictLayer,
|
||||
TreeComponent,
|
||||
get_and_increase_time_counter,
|
||||
next_component_uuid,
|
||||
@@ -15,7 +17,9 @@ __all__ = [
|
||||
"BASE_COMPONENT_TYPE",
|
||||
"ComponentData",
|
||||
"ComponentType",
|
||||
"EvictLayer",
|
||||
"FullComponent",
|
||||
"CacheTransferPhase",
|
||||
"MambaComponent",
|
||||
"SWAComponent",
|
||||
"TreeComponent",
|
||||
|
||||
@@ -1,14 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
from typing import TYPE_CHECKING, Callable, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
DecLockRefParams,
|
||||
EvictParams,
|
||||
IncLockRefResult,
|
||||
MatchPrefixParams,
|
||||
MatchResult,
|
||||
)
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
|
||||
CacheTransferPhase,
|
||||
ComponentType,
|
||||
EvictLayer,
|
||||
TreeComponent,
|
||||
)
|
||||
|
||||
@@ -30,27 +38,80 @@ class FullComponent(TreeComponent):
|
||||
self._free_full = allocator.full_attn_allocator.free
|
||||
else:
|
||||
self._free_full = allocator.free
|
||||
|
||||
def node_has_component_data(self, node: UnifiedTreeNode) -> bool:
|
||||
# Override so _for_each_component_lru includes Full in LRU operations
|
||||
return node.component_data[self.component_type].value is not None
|
||||
# HiCache state: set to host KV pool when HiCache enabled
|
||||
self._full_kv_pool_host = None
|
||||
|
||||
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
|
||||
return lambda node: True
|
||||
# HiCache: evicted + backuped nodes are valid match boundaries
|
||||
return lambda node: (
|
||||
node.component_data[self.component_type].value is not None or node.backuped
|
||||
)
|
||||
|
||||
def finalize_match_result(
|
||||
self,
|
||||
result: MatchResult,
|
||||
params: MatchPrefixParams,
|
||||
value_chunks: list[torch.Tensor],
|
||||
best_value_len: int,
|
||||
) -> MatchResult:
|
||||
# Compute Full KV host hit length: walk from last_host_node up to
|
||||
# last_device_node, summing host_value lengths of evicted nodes.
|
||||
ct = self.component_type
|
||||
kv_host_hit = 0
|
||||
node = result.last_host_node
|
||||
root_node = self.cache.root_node
|
||||
while node is not result.last_device_node and node is not root_node:
|
||||
full_host = node.component_data[ct].host_value
|
||||
if full_host is not None:
|
||||
kv_host_hit += len(full_host)
|
||||
node = node.parent
|
||||
if kv_host_hit > 0:
|
||||
return result._replace(
|
||||
host_hit_length=max(result.host_hit_length, kv_host_hit)
|
||||
)
|
||||
return result
|
||||
|
||||
def redistribute_on_node_split(
|
||||
self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
|
||||
):
|
||||
new_parent.component_data[self.component_type].lock_ref = child.component_data[
|
||||
self.component_type
|
||||
].lock_ref
|
||||
ct = self.component_type
|
||||
new_parent.component_data[ct].lock_ref = child.component_data[ct].lock_ref
|
||||
child_cd = child.component_data[ct]
|
||||
split_len = len(new_parent.key)
|
||||
if child_cd.value is not None:
|
||||
new_parent.component_data[ct].value = child_cd.value[:split_len].clone()
|
||||
child_cd.value = child_cd.value[split_len:].clone()
|
||||
if child_cd.host_value is not None:
|
||||
new_parent.component_data[ct].host_value = child_cd.host_value[
|
||||
:split_len
|
||||
].clone()
|
||||
child_cd.host_value = child_cd.host_value[split_len:].clone()
|
||||
|
||||
def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int:
|
||||
def evict_component(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
target: EvictLayer = EvictLayer.DEVICE,
|
||||
) -> tuple[int, int]:
|
||||
cd = node.component_data[self.component_type]
|
||||
self._free_full(cd.value)
|
||||
freed = len(cd.value)
|
||||
self.cache.component_evictable_size_[self.component_type] -= freed
|
||||
return freed
|
||||
freed = 0
|
||||
host_freed = 0
|
||||
|
||||
# Device layer
|
||||
if EvictLayer.DEVICE in target and cd.value is not None:
|
||||
self._free_full(cd.value)
|
||||
freed = len(cd.value)
|
||||
self.cache.component_evictable_size_[self.component_type] -= freed
|
||||
# NOTE: cd.value = None is deferred to _cascade_evict (Full as trigger)
|
||||
# because SWA's free_swa still needs to read Full.value.
|
||||
# cd.value = None
|
||||
|
||||
# Host layer
|
||||
if EvictLayer.HOST in target and cd.host_value is not None:
|
||||
host_freed = len(cd.host_value)
|
||||
if self._full_kv_pool_host is not None:
|
||||
self._full_kv_pool_host.free(cd.host_value)
|
||||
cd.host_value = None
|
||||
return freed, host_freed
|
||||
|
||||
def eviction_priority(self, is_leaf: bool) -> int:
|
||||
return 0 if is_leaf else 2
|
||||
@@ -59,30 +120,55 @@ class FullComponent(TreeComponent):
|
||||
self, params: EvictParams, tracker: dict[ComponentType, int]
|
||||
) -> None:
|
||||
request = params.num_tokens
|
||||
lru = self.cache.lru_lists[self.component_type]
|
||||
while tracker[self.component_type] < request:
|
||||
x = lru.get_leaf_lru_no_lock()
|
||||
if x is None:
|
||||
break
|
||||
self.cache._evict_component_and_detach_lru(
|
||||
x, self, is_leaf=True, tracker=tracker
|
||||
)
|
||||
self.cache._cascade_evict(x, self, tracker)
|
||||
# Heap-based eviction from evictable_device_leaves, ordered by LRU.
|
||||
heap = [(n.last_access_time, n) for n in self.cache.evictable_device_leaves]
|
||||
heapq.heapify(heap)
|
||||
ct = self.component_type
|
||||
while tracker[ct] < request and heap:
|
||||
_, x = heapq.heappop(heap)
|
||||
if x not in self.cache.evictable_device_leaves:
|
||||
continue
|
||||
self.cache._evict_device_leaf(x, tracker)
|
||||
if x.parent is not None and x.parent in self.cache.evictable_device_leaves:
|
||||
heapq.heappush(heap, (x.parent.last_access_time, x.parent))
|
||||
|
||||
def drive_host_eviction(
|
||||
self, num_tokens: int, tracker: dict[ComponentType, int]
|
||||
) -> None:
|
||||
"""Evict host leaves to free KV host pool space."""
|
||||
heap = [(n.last_access_time, n) for n in self.cache.evictable_host_leaves]
|
||||
heapq.heapify(heap)
|
||||
ct = self.component_type
|
||||
while tracker[ct] < num_tokens and heap:
|
||||
_, x = heapq.heappop(heap)
|
||||
if x not in self.cache.evictable_host_leaves:
|
||||
continue
|
||||
self.cache._evict_host_leaf(x, tracker)
|
||||
if x.parent is not None and x.parent in self.cache.evictable_host_leaves:
|
||||
heapq.heappush(heap, (x.parent.last_access_time, x.parent))
|
||||
|
||||
def acquire_component_lock(
|
||||
self, node: UnifiedTreeNode, result: IncLockRefResult
|
||||
) -> IncLockRefResult:
|
||||
ct = self.component_type
|
||||
root = self.cache.root_node
|
||||
delta = 0
|
||||
cur = node
|
||||
while cur != root:
|
||||
cd = cur.component_data[ct]
|
||||
assert cd.value is not None
|
||||
|
||||
if cd.lock_ref == 0:
|
||||
key_len = len(cd.value)
|
||||
self.cache.component_evictable_size_[ct] -= key_len
|
||||
self.cache.component_protected_size_[ct] += key_len
|
||||
delta += key_len
|
||||
cd.lock_ref += 1
|
||||
self.cache.evictable_device_leaves.discard(cur)
|
||||
cur = cur.parent
|
||||
result = IncLockRefResult(
|
||||
delta=delta, swa_uuid_for_lock=result.swa_uuid_for_lock
|
||||
)
|
||||
return result
|
||||
|
||||
def release_component_lock(
|
||||
@@ -93,10 +179,86 @@ class FullComponent(TreeComponent):
|
||||
cur = node
|
||||
while cur != root:
|
||||
cd = cur.component_data[ct]
|
||||
assert cd.value is not None
|
||||
assert cd.lock_ref > 0
|
||||
|
||||
if cd.lock_ref == 1:
|
||||
key_len = len(cd.value)
|
||||
self.cache.component_evictable_size_[ct] += key_len
|
||||
self.cache.component_protected_size_[ct] -= key_len
|
||||
cd.lock_ref -= 1
|
||||
if cd.lock_ref == 0:
|
||||
self.cache._update_evictable_leaf_sets(cur)
|
||||
cur = cur.parent
|
||||
|
||||
# ---- HiCache Hooks ----
|
||||
|
||||
def build_hicache_transfers(
|
||||
self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw
|
||||
) -> Optional[list[PoolTransfer]]:
|
||||
ct = self.component_type
|
||||
|
||||
if phase == CacheTransferPhase.BACKUP_HOST:
|
||||
# Full KV backup is handled by the main flow
|
||||
# (write_backup → cache_controller.write on host_value directly).
|
||||
# No extra PoolTransfer needed.
|
||||
return None
|
||||
|
||||
if phase == CacheTransferPhase.LOAD_BACK:
|
||||
# Walk evicted chain, collect host_values and nodes
|
||||
backed_up: list[torch.Tensor] = []
|
||||
nodes: list = []
|
||||
cur = node
|
||||
while cur.evicted:
|
||||
cd = cur.component_data[ct]
|
||||
if cd.host_value is not None:
|
||||
backed_up.append(cd.host_value)
|
||||
nodes.append(cur)
|
||||
cur = cur.parent
|
||||
backed_up.reverse()
|
||||
nodes.reverse()
|
||||
return [
|
||||
PoolTransfer(
|
||||
name=PoolName.KV,
|
||||
host_indices=(
|
||||
torch.cat(backed_up)
|
||||
if backed_up
|
||||
else torch.empty((0,), dtype=torch.int64, device="cpu")
|
||||
),
|
||||
device_indices=None,
|
||||
nodes_to_load=nodes,
|
||||
)
|
||||
]
|
||||
|
||||
return None
|
||||
|
||||
def commit_hicache_transfer(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
phase: CacheTransferPhase,
|
||||
transfers: list[PoolTransfer] = (),
|
||||
) -> None:
|
||||
ct = self.component_type
|
||||
|
||||
if phase == CacheTransferPhase.BACKUP_HOST:
|
||||
if transfers and transfers[0].host_indices is not None:
|
||||
node.component_data[ct].host_value = transfers[0].host_indices.clone()
|
||||
|
||||
elif phase == CacheTransferPhase.LOAD_BACK:
|
||||
if not transfers or transfers[0].device_indices is None:
|
||||
self.cache._update_evictable_leaf_sets(node)
|
||||
return
|
||||
|
||||
xfer = transfers[0]
|
||||
device_indices = xfer.device_indices
|
||||
offset = 0
|
||||
for n in xfer.nodes_to_load or []:
|
||||
cd = n.component_data[ct]
|
||||
n_len = len(cd.host_value)
|
||||
cd.value = device_indices[offset : offset + n_len].clone()
|
||||
offset += n_len
|
||||
# Full uses leaf sets, not LRU
|
||||
self.cache.component_evictable_size_[ct] += n_len
|
||||
self.cache._update_evictable_leaf_sets(n)
|
||||
|
||||
self.cache._update_evictable_leaf_sets(node)
|
||||
|
||||
@@ -13,8 +13,11 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
MatchPrefixParams,
|
||||
MatchResult,
|
||||
)
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
|
||||
CacheTransferPhase,
|
||||
ComponentType,
|
||||
EvictLayer,
|
||||
TreeComponent,
|
||||
get_and_increase_time_counter,
|
||||
)
|
||||
@@ -44,10 +47,16 @@ class MambaComponent(TreeComponent):
|
||||
), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {cache.page_size}"
|
||||
super().__init__(cache, params)
|
||||
self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer
|
||||
# HiCache state
|
||||
self._mamba_pool_host = None # set to host mamba pool when HiCache enabled
|
||||
|
||||
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
|
||||
ct = self.component_type
|
||||
return lambda node: node.component_data[ct].value is not None
|
||||
# HiCache: evicted + backuped (host_value present) is also a valid match
|
||||
return lambda node: (
|
||||
node.component_data[ct].value is not None
|
||||
or node.component_data[ct].host_value is not None
|
||||
)
|
||||
|
||||
def finalize_match_result(
|
||||
self,
|
||||
@@ -90,6 +99,13 @@ class MambaComponent(TreeComponent):
|
||||
mamba_value, dst_index
|
||||
)
|
||||
|
||||
# HiCache: if mamba was evicted from device but has host backup,
|
||||
# ensure host_hit_length >= 1 so load_back is triggered.
|
||||
host_node = result.last_host_node
|
||||
cd = host_node.component_data[self.component_type]
|
||||
if cd.value is None and cd.host_value is not None:
|
||||
result = result._replace(host_hit_length=max(result.host_hit_length, 1))
|
||||
|
||||
return result._replace(mamba_branching_seqlen=branching_seqlen)
|
||||
|
||||
def commit_insert_component_data(
|
||||
@@ -109,6 +125,10 @@ class MambaComponent(TreeComponent):
|
||||
return
|
||||
if node.component_data[self.component_type].value is None:
|
||||
node.component_data[self.component_type].value = params.mamba_value
|
||||
# move from host LRU to device LRU
|
||||
host_lru = self.cache.host_lru_lists[self.component_type]
|
||||
if host_lru.in_list(node):
|
||||
host_lru.remove_node(node)
|
||||
self.cache.lru_lists[self.component_type].insert_mru(node)
|
||||
self.cache.component_evictable_size_[self.component_type] += len(
|
||||
params.mamba_value
|
||||
@@ -122,41 +142,74 @@ class MambaComponent(TreeComponent):
|
||||
def redistribute_on_node_split(
|
||||
self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
|
||||
):
|
||||
new_parent.component_data[self.component_type].value = None
|
||||
new_parent.component_data[self.component_type].lock_ref = 0
|
||||
ct = self.component_type
|
||||
new_parent.component_data[ct].value = None
|
||||
new_parent.component_data[ct].lock_ref = 0
|
||||
# HiCache: mamba host_value stays on child (mamba = leaf-only data)
|
||||
new_parent.component_data[ct].host_value = None
|
||||
new_parent.component_data[ct].host_lock_ref = 0
|
||||
|
||||
def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int:
|
||||
value = node.component_data[self.component_type].value
|
||||
self.cache.req_to_token_pool.mamba_pool.free(value)
|
||||
freed = len(value)
|
||||
self.cache.component_evictable_size_[self.component_type] -= freed
|
||||
if not is_leaf:
|
||||
node.component_data[self.component_type].value = None
|
||||
return freed
|
||||
def evict_component(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
target: EvictLayer = EvictLayer.DEVICE,
|
||||
) -> tuple[int, int]:
|
||||
cd = node.component_data[self.component_type]
|
||||
freed = 0
|
||||
host_freed = 0
|
||||
|
||||
# Device layer
|
||||
if EvictLayer.DEVICE in target and cd.value is not None:
|
||||
self.cache.req_to_token_pool.mamba_pool.free(cd.value)
|
||||
freed = len(cd.value)
|
||||
self.cache.component_evictable_size_[self.component_type] -= freed
|
||||
cd.value = None
|
||||
|
||||
# Host layer
|
||||
host_lru = self.cache.host_lru_lists[self.component_type]
|
||||
if EvictLayer.HOST in target and cd.host_value is not None:
|
||||
host_freed = len(cd.host_value)
|
||||
if self._mamba_pool_host is not None:
|
||||
self._mamba_pool_host.free(cd.host_value)
|
||||
cd.host_value = None
|
||||
if host_lru.in_list(node):
|
||||
host_lru.remove_node(node)
|
||||
|
||||
# After device tombstone: if only host_value remains, insert into host LRU
|
||||
if (
|
||||
target is EvictLayer.DEVICE
|
||||
and cd.value is None
|
||||
and cd.host_value is not None
|
||||
):
|
||||
if not host_lru.in_list(node):
|
||||
host_lru.insert_mru(node)
|
||||
|
||||
return freed, host_freed
|
||||
|
||||
def drive_eviction(
|
||||
self, params: EvictParams, tracker: dict[ComponentType, int]
|
||||
) -> None:
|
||||
request = params.mamba_num
|
||||
lru = self.cache.lru_lists[self.component_type]
|
||||
ct = self.component_type
|
||||
lru = self.cache.lru_lists[ct]
|
||||
x = lru.get_lru_no_lock()
|
||||
while (
|
||||
tracker[self.component_type] < request and x is not None and lru.in_list(x)
|
||||
):
|
||||
assert x.component_data[self.component_type].value is not None
|
||||
if len(x.children) > 0:
|
||||
while tracker[ct] < request and x is not None and lru.in_list(x):
|
||||
assert x.component_data[ct].value is not None
|
||||
if x in self.cache.evictable_device_leaves:
|
||||
# D-leaf: atomic eviction of all components
|
||||
x_next = lru.get_prev_no_lock(x)
|
||||
self.cache._evict_device_leaf(x, tracker)
|
||||
if not lru.in_list(x_next):
|
||||
x_next = lru.get_lru_no_lock()
|
||||
x = x_next
|
||||
else:
|
||||
# Internal: tombstone Mamba + cascade
|
||||
x_next = lru.get_prev_no_lock(x)
|
||||
self.cache._evict_component_and_detach_lru(
|
||||
x, self, is_leaf=False, tracker=tracker
|
||||
x, self, target=EvictLayer.DEVICE, tracker=tracker
|
||||
)
|
||||
self.cache._cascade_evict(x, self, tracker)
|
||||
x = x_next
|
||||
else:
|
||||
self.cache._evict_component_and_detach_lru(
|
||||
x, self, is_leaf=True, tracker=tracker
|
||||
)
|
||||
self.cache._cascade_evict(x, self, tracker)
|
||||
x = lru.get_lru_no_lock()
|
||||
|
||||
def acquire_component_lock(
|
||||
self, node: UnifiedTreeNode, result: IncLockRefResult
|
||||
@@ -178,8 +231,7 @@ class MambaComponent(TreeComponent):
|
||||
ct = self.component_type
|
||||
cd = node.component_data[ct]
|
||||
value = cd.value
|
||||
if value is not None:
|
||||
assert cd.lock_ref > 0
|
||||
if value is not None and cd.lock_ref > 0:
|
||||
if cd.lock_ref == 1:
|
||||
vlen = len(value)
|
||||
self.cache.component_evictable_size_[ct] += vlen
|
||||
@@ -268,3 +320,114 @@ class MambaComponent(TreeComponent):
|
||||
):
|
||||
self.cache.req_to_token_pool.mamba_pool.free(insert_params.mamba_value)
|
||||
req.mamba_last_track_seqlen = None
|
||||
|
||||
# ---- HiCache Hooks ----
|
||||
|
||||
def build_hicache_transfers(
|
||||
self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw
|
||||
) -> Optional[list[PoolTransfer]]:
|
||||
ct = self.component_type
|
||||
|
||||
if phase == CacheTransferPhase.BACKUP_HOST:
|
||||
cd = node.component_data[ct]
|
||||
if cd.value is None:
|
||||
return None
|
||||
return [
|
||||
PoolTransfer(
|
||||
name=PoolName.MAMBA,
|
||||
device_indices=cd.value,
|
||||
)
|
||||
]
|
||||
|
||||
if phase == CacheTransferPhase.LOAD_BACK:
|
||||
req = kw.get("req")
|
||||
transfers: list[PoolTransfer] = []
|
||||
|
||||
cd = node.component_data[ct]
|
||||
if cd.value is not None:
|
||||
return None
|
||||
|
||||
# restore single node if host_value exists and
|
||||
if cd.host_value is not None and cd.value is None:
|
||||
transfers.append(
|
||||
PoolTransfer(
|
||||
name=PoolName.MAMBA,
|
||||
host_indices=cd.host_value,
|
||||
nodes_to_load=[node],
|
||||
)
|
||||
)
|
||||
|
||||
# Per-request mamba CoW (H→D copy into request's device slot)
|
||||
cd = node.component_data[ct]
|
||||
if req is not None and cd.host_value is not None:
|
||||
if req.mamba_pool_idx is None:
|
||||
dst = self.cache.req_to_token_pool.mamba_pool.alloc(1)
|
||||
if dst is None:
|
||||
self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
|
||||
dst = self.cache.req_to_token_pool.mamba_pool.alloc(1)
|
||||
assert dst is not None, "Cannot alloc mamba for load_back"
|
||||
req.mamba_pool_idx = dst[0]
|
||||
transfers.append(
|
||||
PoolTransfer(
|
||||
name=PoolName.MAMBA,
|
||||
host_indices=cd.host_value,
|
||||
device_indices=req.mamba_pool_idx.unsqueeze(0),
|
||||
)
|
||||
)
|
||||
|
||||
return transfers if transfers else None
|
||||
|
||||
return None
|
||||
|
||||
def commit_hicache_transfer(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
phase: CacheTransferPhase,
|
||||
transfers: list[PoolTransfer] = (),
|
||||
) -> None:
|
||||
ct = self.component_type
|
||||
|
||||
if phase == CacheTransferPhase.BACKUP_HOST:
|
||||
if transfers and transfers[0].host_indices is not None:
|
||||
cd = node.component_data[ct]
|
||||
if cd.host_value is None:
|
||||
cd.host_value = transfers[0].host_indices.clone()
|
||||
|
||||
elif phase == CacheTransferPhase.LOAD_BACK:
|
||||
if not transfers:
|
||||
return
|
||||
transfer = transfers[0]
|
||||
if transfer.device_indices is not None:
|
||||
cd = node.component_data[ct]
|
||||
cd.value = transfer.device_indices.clone()
|
||||
count = len(cd.value)
|
||||
# Move from host LRU to device LRU
|
||||
host_lru = self.cache.host_lru_lists[ct]
|
||||
if host_lru.in_list(node):
|
||||
host_lru.remove_node(node)
|
||||
self.cache.lru_lists[ct].insert_mru(node)
|
||||
self.cache.component_evictable_size_[ct] += count
|
||||
|
||||
def drive_host_eviction(
|
||||
self, num_tokens: int, tracker: dict[ComponentType, int]
|
||||
) -> None:
|
||||
"""Evict mamba host resources.
|
||||
Internal nodes: private tombstone (free host mamba only).
|
||||
Host leaves: atomic eviction via _evict_host_leaf."""
|
||||
ct = self.component_type
|
||||
host_lru = self.cache.host_lru_lists[ct]
|
||||
x = host_lru.get_lru_no_lock()
|
||||
while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x):
|
||||
x_next = host_lru.get_prev_no_lock(x)
|
||||
cd = x.component_data[ct]
|
||||
if x in self.cache.evictable_host_leaves:
|
||||
# Host leaf: atomic eviction (all components host + delete)
|
||||
self.cache._evict_host_leaf(x, tracker)
|
||||
else:
|
||||
# Internal: tombstone Mamba + cascade
|
||||
assert cd.host_value is not None
|
||||
self.cache._evict_component_and_detach_lru(
|
||||
x, self, target=EvictLayer.HOST, tracker=tracker
|
||||
)
|
||||
self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
|
||||
x = x_next
|
||||
|
||||
@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
|
||||
BASE_COMPONENT_TYPE,
|
||||
ComponentType,
|
||||
EvictLayer,
|
||||
TreeComponent,
|
||||
next_component_uuid,
|
||||
)
|
||||
@@ -186,10 +187,17 @@ class SWAComponent(TreeComponent):
|
||||
)
|
||||
child.component_data[self.component_type].metadata.pop("uuid", None)
|
||||
|
||||
def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int:
|
||||
def evict_component(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
target: EvictLayer = EvictLayer.DEVICE,
|
||||
) -> tuple[int, int]:
|
||||
if target is EvictLayer.HOST:
|
||||
return 0, 0 # TODO:SWA has no host layer currently
|
||||
|
||||
swa_value = node.component_data[self.component_type].value
|
||||
if swa_value is None:
|
||||
return 0
|
||||
return 0, 0
|
||||
# Direct swa_attn_allocator.free(swa_value) would double-free
|
||||
# free_swa(full_value) has the mapping guard to avoid double-free
|
||||
# TODO: decoupling full and swa free, need further discussion on mapping necessity
|
||||
@@ -198,9 +206,9 @@ class SWAComponent(TreeComponent):
|
||||
)
|
||||
freed = len(swa_value)
|
||||
self.cache.component_evictable_size_[self.component_type] -= freed
|
||||
if not is_leaf:
|
||||
if target is EvictLayer.DEVICE:
|
||||
node.component_data[self.component_type].value = None
|
||||
return freed
|
||||
return freed, 0
|
||||
|
||||
def eviction_priority(self, is_leaf: bool) -> int:
|
||||
return 0 if is_leaf else 1
|
||||
@@ -209,25 +217,26 @@ class SWAComponent(TreeComponent):
|
||||
self, params: EvictParams, tracker: dict[ComponentType, int]
|
||||
) -> None:
|
||||
request = params.swa_num_tokens
|
||||
lru = self.cache.lru_lists[self.component_type]
|
||||
ct = self.component_type
|
||||
lru = self.cache.lru_lists[ct]
|
||||
x = lru.get_lru_no_lock()
|
||||
while (
|
||||
tracker[self.component_type] < request and x is not None and lru.in_list(x)
|
||||
):
|
||||
assert x.component_data[self.component_type].value is not None
|
||||
if len(x.children) > 0:
|
||||
while tracker[ct] < request and x is not None and lru.in_list(x):
|
||||
assert x.component_data[ct].value is not None
|
||||
if x in self.cache.evictable_device_leaves:
|
||||
# D-leaf: atomic eviction of all components
|
||||
x_next = lru.get_prev_no_lock(x)
|
||||
self.cache._evict_device_leaf(x, tracker)
|
||||
if not lru.in_list(x_next):
|
||||
x_next = lru.get_lru_no_lock()
|
||||
x = x_next
|
||||
else:
|
||||
# Internal: tombstone SWA + cascade
|
||||
x_next = lru.get_prev_no_lock(x)
|
||||
self.cache._evict_component_and_detach_lru(
|
||||
x, self, is_leaf=False, tracker=tracker
|
||||
x, self, target=EvictLayer.DEVICE, tracker=tracker
|
||||
)
|
||||
self.cache._cascade_evict(x, self, tracker)
|
||||
x = x_next
|
||||
else:
|
||||
self.cache._evict_component_and_detach_lru(
|
||||
x, self, is_leaf=True, tracker=tracker
|
||||
)
|
||||
self.cache._cascade_evict(x, self, tracker)
|
||||
x = lru.get_lru_no_lock()
|
||||
|
||||
def acquire_component_lock(
|
||||
self, node: UnifiedTreeNode, result: IncLockRefResult
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from enum import Enum, IntFlag
|
||||
from typing import TYPE_CHECKING, Any, Callable, Optional
|
||||
|
||||
import torch
|
||||
@@ -67,6 +67,14 @@ class ComponentData:
|
||||
host_lock_ref: int = 0
|
||||
|
||||
|
||||
class EvictLayer(IntFlag):
|
||||
"""Which storage layer(s) to evict. Combinable via bitwise OR."""
|
||||
|
||||
DEVICE = 1
|
||||
HOST = 2
|
||||
ALL = DEVICE | HOST
|
||||
|
||||
|
||||
class CacheTransferPhase(str, Enum):
|
||||
|
||||
BACKUP_HOST = "backup_host" # D→H
|
||||
@@ -95,8 +103,13 @@ class TreeComponent(ABC):
|
||||
# Subclasses MUST set this as a class attribute (not @property)
|
||||
component_type: ComponentType
|
||||
|
||||
def node_has_component_data(self, node: UnifiedTreeNode) -> bool:
|
||||
return node.component_data[self.component_type].value is not None
|
||||
def node_has_component_data(
|
||||
self, node: UnifiedTreeNode, target: EvictLayer = EvictLayer.DEVICE
|
||||
) -> bool:
|
||||
cd = node.component_data[self.component_type]
|
||||
if target is EvictLayer.DEVICE:
|
||||
return cd.value is not None
|
||||
return cd.host_value is not None
|
||||
|
||||
def value_len(self, node: UnifiedTreeNode) -> int:
|
||||
value = node.component_data[self.component_type].value
|
||||
@@ -186,17 +199,22 @@ class TreeComponent(ABC):
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int:
|
||||
def evict_component(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
target: EvictLayer = EvictLayer.DEVICE,
|
||||
) -> tuple[int, int]:
|
||||
"""Free this component's KV resources on a node being evicted.
|
||||
For internal (non-leaf) nodes: free memory and tombstone the value
|
||||
(set to None); the node structure is kept.
|
||||
For leaf nodes: free memory; the node will be deleted by caller.
|
||||
Returns the number of tokens/slots freed.
|
||||
- Full: frees full_value via token_to_kv_pool_allocator.
|
||||
- SWA: frees swa value via swa_token_to_kv_pool_allocator;
|
||||
only tombstones on internal nodes.
|
||||
- Mamba: frees mamba value via mamba_token_to_kv_pool_allocator;
|
||||
only tombstones on internal nodes."""
|
||||
|
||||
*target* controls which layer(s) to evict:
|
||||
- DEVICE: free device memory and tombstone (value = None).
|
||||
Host data is untouched.
|
||||
- HOST: free host memory (host_value = None).
|
||||
Device data is untouched.
|
||||
- ALL: free both device and host memory.
|
||||
No tombstone — caller will delete the node.
|
||||
|
||||
Returns (device_freed, host_freed) token counts."""
|
||||
...
|
||||
|
||||
def eviction_priority(self, is_leaf: bool) -> int:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -262,5 +262,55 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
|
||||
# TODO(hzh): Currently, HiCache is not fully compatible with the CI CUDA13 environment; we need to wait for the fix before re-enabling the tests below.
|
||||
# class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
|
||||
# """Mamba hybrid + UnifiedRadixCache."""
|
||||
#
|
||||
# kl_threshold = 0.003
|
||||
# prefill_cache_assert = staticmethod(
|
||||
# make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE)
|
||||
# )
|
||||
# decode_cache_assert = staticmethod(
|
||||
# make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL)
|
||||
# )
|
||||
#
|
||||
# @classmethod
|
||||
# def setUpClass(cls):
|
||||
# cls.model = MAMBA_MODEL
|
||||
# cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
# cls.process = popen_launch_server(
|
||||
# cls.model,
|
||||
# cls.base_url,
|
||||
# timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
# other_args=[
|
||||
# "--tp-size",
|
||||
# "4",
|
||||
# "--chunked-prefill-size",
|
||||
# "2048",
|
||||
# "--mem-fraction-static",
|
||||
# "0.85",
|
||||
# "--mamba-scheduler-strategy",
|
||||
# "extra_buffer",
|
||||
# "--mamba-track-interval",
|
||||
# str(MAMBA_TRACK_INTERVAL),
|
||||
# "--enable-hierarchical-cache",
|
||||
# "--hicache-ratio",
|
||||
# "1.5",
|
||||
# "--hicache-write-policy",
|
||||
# "write_through",
|
||||
# "--hicache-io-backend",
|
||||
# "direct",
|
||||
# "--hicache-mem-layout",
|
||||
# "page_first_direct",
|
||||
# ],
|
||||
# env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
|
||||
# )
|
||||
# cls.input_ids = get_input_ids(cls.model, num_samples=18)
|
||||
#
|
||||
# @classmethod
|
||||
# def tearDownClass(cls):
|
||||
# kill_process_tree(cls.process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
@@ -207,17 +208,17 @@ def build_fixture(cfg: CacheConfig):
|
||||
need_sort=False,
|
||||
)
|
||||
|
||||
tree = UnifiedRadixCache(
|
||||
params=CacheInitParams(
|
||||
req_to_token_pool=req_to_token_pool,
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
page_size=cfg.page_size,
|
||||
disable=False,
|
||||
sliding_window_size=cfg.sliding_window_size,
|
||||
tree_components=cfg.components,
|
||||
enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer,
|
||||
),
|
||||
cache_init_params = CacheInitParams(
|
||||
req_to_token_pool=req_to_token_pool,
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
page_size=cfg.page_size,
|
||||
disable=False,
|
||||
sliding_window_size=cfg.sliding_window_size,
|
||||
tree_components=cfg.components,
|
||||
enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer,
|
||||
)
|
||||
tree = UnifiedRadixCache(params=cache_init_params)
|
||||
tree.cache_init_params = cache_init_params
|
||||
|
||||
return tree, allocator, req_to_token_pool
|
||||
|
||||
@@ -934,6 +935,644 @@ class UnifiedRadixCacheSuite:
|
||||
|
||||
tree.sanity_check()
|
||||
|
||||
# ================================================================
|
||||
# Evict chain tests covering demotion, cascade, and tombstone cleanup.
|
||||
# ================================================================
|
||||
|
||||
def test_evict_leaf_frees_all_components(self):
|
||||
"""Evicting a device leaf frees Full and all aux components atomically."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
seq = self._make_seq(1, 3)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq)
|
||||
|
||||
full_before = tree.full_evictable_size()
|
||||
mamba_before = tree.mamba_evictable_size() if self.cfg.has_mamba else 0
|
||||
swa_before = tree.swa_evictable_size() if self.cfg.has_swa else 0
|
||||
self.assertGreater(full_before, 0)
|
||||
|
||||
result = tree.evict(EvictParams(num_tokens=full_before * 2))
|
||||
self.assertGreaterEqual(result.num_tokens_evicted, full_before)
|
||||
self.assertEqual(tree.full_evictable_size(), 0)
|
||||
if self.cfg.has_mamba:
|
||||
self.assertEqual(tree.mamba_evictable_size(), 0)
|
||||
if self.cfg.has_swa:
|
||||
self.assertEqual(tree.swa_evictable_size(), 0)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_cascade_parent_becomes_d_leaf(self):
|
||||
"""After evicting a D-leaf child, parent may become a new D-leaf."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
base = self._make_seq(1, 2)
|
||||
leaf = base + self._make_seq(500, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, base)
|
||||
self._insert(tree, allocator, req_to_token_pool, leaf)
|
||||
|
||||
# Lock the base node to prevent it from being evicted
|
||||
m_base = tree.match_prefix(MatchPrefixParams(key=RadixKey(base)))
|
||||
lock_result = tree.inc_lock_ref(m_base.last_device_node)
|
||||
|
||||
# Evict the leaf — parent (base) should become D-leaf after unlock
|
||||
result = tree.evict(EvictParams(num_tokens=len(leaf)))
|
||||
tree.sanity_check()
|
||||
|
||||
tree.dec_lock_ref(
|
||||
m_base.last_device_node,
|
||||
DecLockRefParams(
|
||||
swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None)
|
||||
),
|
||||
)
|
||||
# After unlock, base should be in evictable_device_leaves
|
||||
self.assertIn(m_base.last_device_node, tree.evictable_device_leaves)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_iterative_tombstone_cleanup(self):
|
||||
"""Tombstone cascade: evicting a leaf triggers cleanup up the tree."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
# Create a chain: root -> A -> B -> C (3 levels)
|
||||
ps = self.cfg.page_size
|
||||
chain = self._make_seq(1, 6)
|
||||
self._insert(tree, allocator, req_to_token_pool, chain[: 2 * ps])
|
||||
self._insert(tree, allocator, req_to_token_pool, chain[: 4 * ps])
|
||||
self._insert(tree, allocator, req_to_token_pool, chain)
|
||||
|
||||
initial_evictable = tree.full_evictable_size()
|
||||
self.assertGreater(initial_evictable, 0)
|
||||
|
||||
# Evict everything — tombstone cascade should clean up all
|
||||
result = tree.evict(EvictParams(num_tokens=initial_evictable * 2))
|
||||
self.assertGreaterEqual(result.num_tokens_evicted, initial_evictable)
|
||||
self.assertEqual(tree.full_evictable_size(), 0)
|
||||
# Only root should remain
|
||||
self.assertEqual(len(tree.root_node.children), 0)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_respects_lru_order(self):
|
||||
"""Older (less recently accessed) nodes are evicted first."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
ps = self.cfg.page_size
|
||||
seq_old = self._make_seq(1, 2)
|
||||
seq_new = self._make_seq(500, 2)
|
||||
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_old)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_new)
|
||||
|
||||
# Touch seq_new to make it MRU
|
||||
tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new)))
|
||||
|
||||
# Evict just enough for one sequence
|
||||
tree.evict(EvictParams(num_tokens=len(seq_old)))
|
||||
|
||||
# seq_old should be gone (LRU), seq_new should remain
|
||||
m_old = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_old)))
|
||||
m_new = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new)))
|
||||
self.assertEqual(len(m_old.device_indices), 0)
|
||||
self.assertEqual(len(m_new.device_indices), len(seq_new))
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_multiple_independent_leaves(self):
|
||||
"""Evicting multiple independent leaves works correctly."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
seqs = [self._make_seq(i * 100, 2) for i in range(4)]
|
||||
for s in seqs:
|
||||
self._insert(tree, allocator, req_to_token_pool, s)
|
||||
|
||||
total = sum(len(s) for s in seqs)
|
||||
self.assertEqual(tree.full_evictable_size(), total)
|
||||
|
||||
# Evict half
|
||||
half = total // 2
|
||||
result = tree.evict(EvictParams(num_tokens=half))
|
||||
self.assertGreaterEqual(result.num_tokens_evicted, half)
|
||||
self.assertLessEqual(tree.full_evictable_size(), total - half)
|
||||
tree.sanity_check()
|
||||
|
||||
# Evict remainder
|
||||
remaining = tree.full_evictable_size()
|
||||
result = tree.evict(EvictParams(num_tokens=remaining * 2))
|
||||
self.assertGreaterEqual(result.num_tokens_evicted, remaining)
|
||||
self.assertEqual(tree.full_evictable_size(), 0)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_shared_prefix_keeps_common_path(self):
|
||||
"""Evicting one branch preserves the shared prefix for other branch."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
base = self._make_seq(1, 2)
|
||||
branch_a = base + self._make_seq(100, 2)
|
||||
branch_b = base + self._make_seq(200, 2)
|
||||
|
||||
self._insert(tree, allocator, req_to_token_pool, branch_a)
|
||||
self._insert(tree, allocator, req_to_token_pool, branch_b)
|
||||
|
||||
# Lock branch_b
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b)))
|
||||
lr = tree.inc_lock_ref(m.last_device_node)
|
||||
|
||||
# Evict — branch_a should go, base + branch_b stay
|
||||
tree.evict(EvictParams(num_tokens=len(branch_a)))
|
||||
|
||||
m_b = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b)))
|
||||
self.assertEqual(len(m_b.device_indices), len(branch_b))
|
||||
|
||||
tree.dec_lock_ref(
|
||||
m.last_device_node,
|
||||
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
|
||||
)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_result_accounting_matches_actual(self):
|
||||
"""EvictResult.num_tokens_evicted matches actual size change."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
seqs = [self._make_seq(i * 100, 2) for i in range(5)]
|
||||
for s in seqs:
|
||||
self._insert(tree, allocator, req_to_token_pool, s)
|
||||
|
||||
before = tree.full_evictable_size()
|
||||
result = tree.evict(EvictParams(num_tokens=before))
|
||||
after = tree.full_evictable_size()
|
||||
self.assertEqual(result.num_tokens_evicted, before - after)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_locked_subtree_skipped(self):
|
||||
"""All nodes in a locked path are skipped during eviction."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
seq_a = self._make_seq(1, 3)
|
||||
seq_b = self._make_seq(500, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_a)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_b)
|
||||
|
||||
# Lock seq_a
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
|
||||
lr = tree.inc_lock_ref(m.last_device_node)
|
||||
|
||||
# Try to evict everything
|
||||
total = tree.full_evictable_size() + tree.full_protected_size()
|
||||
result = tree.evict(EvictParams(num_tokens=total))
|
||||
|
||||
# seq_a should still be matchable (protected)
|
||||
m2 = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
|
||||
self.assertEqual(len(m2.device_indices), len(seq_a))
|
||||
|
||||
tree.dec_lock_ref(
|
||||
m.last_device_node,
|
||||
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
|
||||
)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_mamba_internal_tombstone_evict(self):
|
||||
"""Mamba eviction on internal node tombstones mamba only, keeps Full."""
|
||||
if not self.cfg.has_mamba:
|
||||
self.skipTest("requires Mamba component")
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
# Create internal node with mamba and leaf extending it
|
||||
seq_short = self._make_seq(1, 2)
|
||||
seq_long = seq_short + self._make_seq(500, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_short)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_long)
|
||||
|
||||
# Evict only mamba
|
||||
result = tree.evict(EvictParams(num_tokens=0, mamba_num=10))
|
||||
self.assertEqual(tree.mamba_evictable_size(), 0)
|
||||
|
||||
# Full should still be accessible for at least the long seq base
|
||||
# (mamba gone breaks match, but full data might still be in tree)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_reinsert_after_full_eviction(self):
|
||||
"""After evicting everything, new inserts work correctly."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
seq_a = self._make_seq(1, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_a)
|
||||
tree.evict(EvictParams(num_tokens=len(seq_a) * 2))
|
||||
self.assertEqual(tree.full_evictable_size(), 0)
|
||||
|
||||
# Re-insert
|
||||
seq_b = self._make_seq(500, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq_b)
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_b)))
|
||||
self.assertEqual(len(m.device_indices), len(seq_b))
|
||||
tree.sanity_check()
|
||||
|
||||
def test_swa_evict_internal_tombstone(self):
|
||||
"""SWA eviction on internal node cascades to lower-priority components."""
|
||||
if not self.cfg.has_swa:
|
||||
self.skipTest("requires SWA component")
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
base = self._make_seq(1, 3)
|
||||
leaf = base + self._make_seq(500, 3)
|
||||
self._insert(tree, allocator, req_to_token_pool, base)
|
||||
self._insert(tree, allocator, req_to_token_pool, leaf)
|
||||
|
||||
swa_before = tree.swa_evictable_size()
|
||||
result = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=swa_before * 2))
|
||||
self.assertEqual(tree.swa_evictable_size(), 0)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_evict_d_leaf_set_consistency(self):
|
||||
"""evictable_device_leaves is consistent after mixed operations."""
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
seqs = [self._make_seq(i * 100, 2) for i in range(6)]
|
||||
for s in seqs:
|
||||
self._insert(tree, allocator, req_to_token_pool, s)
|
||||
|
||||
# Lock some, evict some, unlock
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[0])))
|
||||
lr = tree.inc_lock_ref(m.last_device_node)
|
||||
|
||||
tree.evict(EvictParams(num_tokens=len(seqs[1])))
|
||||
tree.sanity_check()
|
||||
|
||||
tree.dec_lock_ref(
|
||||
m.last_device_node,
|
||||
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
|
||||
)
|
||||
tree.sanity_check()
|
||||
|
||||
# Insert more
|
||||
extra = self._make_seq(9000, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, extra)
|
||||
tree.sanity_check()
|
||||
|
||||
# ================================================================
|
||||
# HiCache Unit Tests (real cache_controller D<->H backup/load)
|
||||
# ================================================================
|
||||
|
||||
def _skip_unsupported_hicache_test(self):
|
||||
if self.cfg.has_swa:
|
||||
self.skipTest("HiCache tests do not run on SWA stacks")
|
||||
return False
|
||||
|
||||
def _init_hicache(self, tree):
|
||||
import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
|
||||
|
||||
orig_kv_host_pool = assembler.MHATokenToKVPoolHost
|
||||
orig_mamba_host_pool = assembler.MambaPoolHost
|
||||
|
||||
def kv_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
return orig_kv_host_pool(*args, **kwargs)
|
||||
|
||||
def mamba_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
return orig_mamba_host_pool(*args, **kwargs)
|
||||
|
||||
patchers = [
|
||||
mock.patch.object(
|
||||
assembler,
|
||||
"MHATokenToKVPoolHost",
|
||||
side_effect=kv_host_pool_wrapper,
|
||||
),
|
||||
mock.patch.object(
|
||||
assembler,
|
||||
"MambaPoolHost",
|
||||
side_effect=mamba_host_pool_wrapper,
|
||||
),
|
||||
]
|
||||
for patcher in patchers:
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
page_size=self.cfg.page_size,
|
||||
hicache_io_backend="direct",
|
||||
hicache_write_policy="write_through",
|
||||
)
|
||||
set_global_server_args_for_scheduler(server_args)
|
||||
tree.init_hicache(server_args, tree.cache_init_params)
|
||||
tree.write_through_threshold = 1 << 30
|
||||
tree.load_back_threshold = 0
|
||||
|
||||
def _build_hicache_fixture(self):
|
||||
fixture = build_fixture(self.cfg)
|
||||
tree, _, _ = fixture
|
||||
self._init_hicache(tree)
|
||||
return fixture
|
||||
|
||||
def _backup_node(self, tree, node):
|
||||
backed_up = tree.write_backup(node, write_back=True)
|
||||
self.assertGreater(backed_up, 0)
|
||||
tree.writing_check(write_back=True)
|
||||
return backed_up
|
||||
|
||||
def _backup_tree(self, tree):
|
||||
stack = [tree.root_node]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
children = list(node.children.values())
|
||||
stack.extend(reversed(children))
|
||||
if node is not tree.root_node:
|
||||
self._backup_node(tree, node)
|
||||
|
||||
def _load_back_node(self, tree, node):
|
||||
device_indices = tree.load_back(node)
|
||||
self.assertIsNotNone(device_indices)
|
||||
producer_id = tree.ready_to_load_host_cache()
|
||||
self.assertNotEqual(producer_id, -1)
|
||||
for _, finish_event, _ in list(tree.cache_controller.ack_load_queue):
|
||||
finish_event.synchronize()
|
||||
tree.loading_check()
|
||||
return device_indices
|
||||
|
||||
def _get_full_kv_pool(self, allocator):
|
||||
kv_pool = allocator.get_kvcache()
|
||||
return getattr(kv_pool, "full_kv_pool", kv_pool)
|
||||
|
||||
def _fill_full_kv(self, allocator, indices, marker):
|
||||
kv_pool = self._get_full_kv_pool(allocator)
|
||||
layer_id = kv_pool.start_layer
|
||||
k_buf = kv_pool.get_key_buffer(layer_id)
|
||||
v_buf = kv_pool.get_value_buffer(layer_id)
|
||||
k_buf[indices].fill_(marker)
|
||||
v_buf[indices].fill_(marker + 1)
|
||||
|
||||
def _snapshot_full_kv(self, allocator, indices):
|
||||
kv_pool = self._get_full_kv_pool(allocator)
|
||||
layer_id = kv_pool.start_layer
|
||||
return (
|
||||
kv_pool.get_key_buffer(layer_id)[indices].float().cpu().clone(),
|
||||
kv_pool.get_value_buffer(layer_id)[indices].float().cpu().clone(),
|
||||
)
|
||||
|
||||
def _fill_mamba_state(self, req_to_token_pool, indices, marker):
|
||||
if not self.cfg.has_mamba:
|
||||
return
|
||||
mamba_indices = indices.reshape(-1)
|
||||
mamba_cache = req_to_token_pool.mamba_pool.mamba_cache
|
||||
mamba_cache.temporal[:, mamba_indices].fill_(marker)
|
||||
for offset, conv_buf in enumerate(mamba_cache.conv, start=1):
|
||||
conv_buf[:, mamba_indices].fill_(marker + offset)
|
||||
|
||||
def _snapshot_mamba_state(self, req_to_token_pool, indices):
|
||||
mamba_indices = indices.reshape(-1)
|
||||
mamba_cache = req_to_token_pool.mamba_pool.mamba_cache
|
||||
return (
|
||||
mamba_cache.temporal[:, mamba_indices].float().cpu().clone(),
|
||||
[conv[:, mamba_indices].float().cpu().clone() for conv in mamba_cache.conv],
|
||||
)
|
||||
|
||||
def test_hicache_node_states(self):
|
||||
"""Verify device-only to device+host transition after real backup."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
seq = self._make_seq(1, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq)
|
||||
|
||||
# Find the leaf node
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
|
||||
node = m.last_device_node
|
||||
self.assertIsNot(node, tree.root_node)
|
||||
|
||||
ct = ComponentType.FULL
|
||||
# S1: device only
|
||||
self.assertIsNotNone(node.component_data[ct].value)
|
||||
self.assertIsNone(node.component_data[ct].host_value)
|
||||
self.assertFalse(node.backuped)
|
||||
self.assertFalse(node.evicted)
|
||||
|
||||
self._backup_node(tree, node)
|
||||
self.assertIsNotNone(node.component_data[ct].value)
|
||||
self.assertIsNotNone(node.component_data[ct].host_value)
|
||||
self.assertTrue(node.backuped)
|
||||
self.assertFalse(node.evicted)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_evict_to_host(self):
|
||||
"""Evicting a backed-up device leaf demotes it to host-only state."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
seq = self._make_seq(1, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq)
|
||||
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
|
||||
node = m.last_device_node
|
||||
|
||||
self._backup_node(tree, node)
|
||||
self.assertTrue(node.backuped)
|
||||
|
||||
# Evict -> should demote to host (S3)
|
||||
result = tree.evict(EvictParams(num_tokens=len(seq)))
|
||||
self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
|
||||
|
||||
# Node should now be evicted (S3)
|
||||
self.assertTrue(node.evicted)
|
||||
self.assertTrue(node.backuped)
|
||||
self.assertIsNone(node.component_data[ComponentType.FULL].value)
|
||||
self.assertIsNotNone(node.component_data[ComponentType.FULL].host_value)
|
||||
|
||||
# Should be in host_leaves, not device_leaves
|
||||
self.assertNotIn(node, tree.evictable_device_leaves)
|
||||
self.assertIn(node, tree.evictable_host_leaves)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_match_through_evicted_node(self):
|
||||
"""Match can traverse evicted (S3) nodes using host_value."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
base = self._make_seq(1, 2)
|
||||
leaf = base + self._make_seq(500, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, base)
|
||||
self._insert(tree, allocator, req_to_token_pool, leaf)
|
||||
|
||||
self._backup_tree(tree)
|
||||
|
||||
# Lock leaf so only base can be evicted
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf)))
|
||||
lr = tree.inc_lock_ref(m.last_device_node)
|
||||
|
||||
# Evict base (inner node won't be evicted while child is locked)
|
||||
tree.evict(EvictParams(num_tokens=len(base)))
|
||||
|
||||
tree.dec_lock_ref(
|
||||
m.last_device_node,
|
||||
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
|
||||
)
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf)))
|
||||
self.assertGreaterEqual(len(m.device_indices), len(base))
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_d_leaf_h_leaf_mutual_exclusion(self):
|
||||
"""D-leaf and H-leaf sets are always disjoint."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
seqs = [self._make_seq(i * 100, 2) for i in range(4)]
|
||||
for s in seqs:
|
||||
self._insert(tree, allocator, req_to_token_pool, s)
|
||||
|
||||
for i in range(2):
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i])))
|
||||
self._backup_node(tree, m.last_device_node)
|
||||
|
||||
# Evict one backed-up node
|
||||
tree.evict(EvictParams(num_tokens=len(seqs[0])))
|
||||
|
||||
# Check mutual exclusion
|
||||
overlap = tree.evictable_device_leaves & tree.evictable_host_leaves
|
||||
self.assertEqual(len(overlap), 0)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_host_leaf_eviction(self):
|
||||
"""Evicting a host leaf removes the node from the tree entirely."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
seq = self._make_seq(1, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq)
|
||||
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
|
||||
node = m.last_device_node
|
||||
|
||||
self._backup_node(tree, node)
|
||||
tree.evict(EvictParams(num_tokens=len(seq)))
|
||||
|
||||
self.assertTrue(node.evicted)
|
||||
self.assertIn(node, tree.evictable_host_leaves)
|
||||
|
||||
# Now evict host
|
||||
tree.evict_host(len(seq))
|
||||
|
||||
# Node should be removed from tree
|
||||
self.assertNotIn(node, tree.evictable_host_leaves)
|
||||
self.assertEqual(len(tree.root_node.children), 0)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_load_back_restores_data(self):
|
||||
"""Loading back an evicted node restores the backed-up cache data."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
base = self._make_seq(1, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, base)
|
||||
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(base)))
|
||||
node = m.last_device_node
|
||||
original_device_indices = m.device_indices.clone()
|
||||
self._fill_full_kv(allocator, original_device_indices, marker=3)
|
||||
expected_k, expected_v = self._snapshot_full_kv(
|
||||
allocator, original_device_indices
|
||||
)
|
||||
original_mamba_indices = None
|
||||
expected_temporal = None
|
||||
expected_conv = None
|
||||
if self.cfg.has_mamba:
|
||||
original_mamba_indices = node.component_data[
|
||||
ComponentType.MAMBA
|
||||
].value.clone()
|
||||
self._fill_mamba_state(req_to_token_pool, original_mamba_indices, marker=11)
|
||||
expected_temporal, expected_conv = self._snapshot_mamba_state(
|
||||
req_to_token_pool, original_mamba_indices
|
||||
)
|
||||
|
||||
self._backup_node(tree, node)
|
||||
tree.evict(EvictParams(num_tokens=len(base)))
|
||||
self.assertTrue(node.evicted)
|
||||
self._fill_full_kv(allocator, original_device_indices, marker=9)
|
||||
if original_mamba_indices is not None:
|
||||
self._fill_mamba_state(req_to_token_pool, original_mamba_indices, marker=21)
|
||||
|
||||
loaded_indices = self._load_back_node(tree, node)
|
||||
self.assertFalse(node.evicted)
|
||||
self.assertIsNotNone(node.component_data[ComponentType.FULL].value)
|
||||
loaded_k, loaded_v = self._snapshot_full_kv(allocator, loaded_indices)
|
||||
self.assertTrue(torch.equal(loaded_k, expected_k))
|
||||
self.assertTrue(torch.equal(loaded_v, expected_v))
|
||||
if self.cfg.has_mamba:
|
||||
loaded_mamba_indices = node.component_data[ComponentType.MAMBA].value
|
||||
loaded_temporal, loaded_conv = self._snapshot_mamba_state(
|
||||
req_to_token_pool, loaded_mamba_indices
|
||||
)
|
||||
self.assertTrue(torch.equal(loaded_temporal, expected_temporal))
|
||||
self.assertEqual(len(loaded_conv), len(expected_conv))
|
||||
for actual_conv, expected_conv_buf in zip(loaded_conv, expected_conv):
|
||||
self.assertTrue(torch.equal(actual_conv, expected_conv_buf))
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_backup_continuity(self):
|
||||
"""Backed-up nodes form a continuous prefix from the root."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
chain = self._make_seq(1, 4)
|
||||
ps = self.cfg.page_size
|
||||
self._insert(tree, allocator, req_to_token_pool, chain[: 2 * ps])
|
||||
self._insert(tree, allocator, req_to_token_pool, chain)
|
||||
|
||||
self._backup_tree(tree)
|
||||
|
||||
# Verify: every backed-up node's parent is also backed-up (or root)
|
||||
all_nodes = tree._collect_all_nodes()
|
||||
for node in all_nodes:
|
||||
if node is tree.root_node:
|
||||
continue
|
||||
if node.backuped:
|
||||
parent = node.parent
|
||||
self.assertTrue(
|
||||
parent is tree.root_node or parent.backuped,
|
||||
f"Backup continuity violated: node {node.id} backed up but parent {parent.id} not",
|
||||
)
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_evict_to_host_updates_aux_lru(self):
|
||||
"""Aux components move from device LRU to host LRU on device-to-host eviction."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
if not self.cfg.has_mamba:
|
||||
self.skipTest("requires Mamba component")
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
seq = self._make_seq(1, 2)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq)
|
||||
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
|
||||
node = m.last_device_node
|
||||
|
||||
# Check mamba is in device LRU
|
||||
mamba_lru = tree.lru_lists[ComponentType.MAMBA]
|
||||
host_mamba_lru = tree.host_lru_lists[ComponentType.MAMBA]
|
||||
self.assertTrue(mamba_lru.in_list(node))
|
||||
self.assertFalse(host_mamba_lru.in_list(node))
|
||||
|
||||
self._backup_node(tree, node)
|
||||
tree.evict(EvictParams(num_tokens=len(seq)))
|
||||
|
||||
# Mamba should move to host LRU
|
||||
self.assertFalse(mamba_lru.in_list(node))
|
||||
if node.component_data[ComponentType.MAMBA].host_value is not None:
|
||||
self.assertTrue(host_mamba_lru.in_list(node))
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_mixed_backup_evict_insert(self):
|
||||
"""Complex scenario: backup some, evict, insert new, verify invariants."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
seqs = [self._make_seq(i * 100, 2) for i in range(5)]
|
||||
|
||||
# Insert all
|
||||
for s in seqs:
|
||||
self._insert(tree, allocator, req_to_token_pool, s)
|
||||
tree.sanity_check()
|
||||
|
||||
for i in range(3):
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i])))
|
||||
self._backup_node(tree, m.last_device_node)
|
||||
|
||||
# Evict to free some tokens
|
||||
tree.evict(EvictParams(num_tokens=len(seqs[0]) * 2))
|
||||
tree.sanity_check()
|
||||
|
||||
# Insert new sequences
|
||||
new_seqs = [self._make_seq(i * 1000, 2) for i in range(3)]
|
||||
for s in new_seqs:
|
||||
self._insert(tree, allocator, req_to_token_pool, s)
|
||||
tree.sanity_check()
|
||||
|
||||
# Verify D-leaf / H-leaf mutual exclusion
|
||||
overlap = tree.evictable_device_leaves & tree.evictable_host_leaves
|
||||
self.assertEqual(len(overlap), 0)
|
||||
|
||||
|
||||
_CONFIGS: list[CacheConfig] = [
|
||||
CacheConfig(page_size=1, components=(ComponentType.FULL,)),
|
||||
|
||||
Reference in New Issue
Block a user