diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 70c1df8c8..fbdc08d64 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -71,7 +71,7 @@ class DecodeHiCachePreallocMixin: l3_storage_hit_length = 0 last_host_node = None if self.scheduler.enable_decode_hicache: - last_host_node = result.last_host_node + last_host_node = self.tree_cache.resolve_node_handle(result.last_host_node) if last_host_node.backuped or last_host_node is self.tree_cache.root_node: matched_len = l1_prefix_len + l2_host_hit_length suffix_tokens = req.origin_input_ids[matched_len:] @@ -82,7 +82,7 @@ class DecodeHiCachePreallocMixin: else None ) l3_storage_hit_length = self.tree_cache.query_storage_hit_length( - last_host_node, + result.last_host_node, suffix_tokens, last_hash, prefix_keys, @@ -93,7 +93,9 @@ class DecodeHiCachePreallocMixin: l2_host_hit_length=l2_host_hit_length, l3_storage_hit_length=l3_storage_hit_length, last_device_node=result.last_device_node, - last_host_node=last_host_node if l3_storage_hit_length > 0 else None, + last_host_node=( + result.last_host_node if l3_storage_hit_length > 0 else None + ), ) def _start_hicache_prefetch( @@ -110,7 +112,7 @@ class DecodeHiCachePreallocMixin: ): return try: - node = prefix_match.last_host_node + node = self.tree_cache.resolve_node_handle(prefix_match.last_host_node) matched_len = prefix_match.l1_prefix_len + prefix_match.l2_host_hit_length suffix = req.origin_input_ids[ matched_len : matched_len + prefix_match.l3_storage_hit_length @@ -122,7 +124,7 @@ class DecodeHiCachePreallocMixin: else None ) self.tree_cache.prefetch_from_storage( - req.rid, node, suffix, last_hash, prefix_keys + req.rid, prefix_match.last_host_node, suffix, last_hash, prefix_keys ) prefix_match.prefetch_registered = ( req.rid in self.tree_cache.ongoing_prefetch diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2753260cd..b4ff61057 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -847,6 +847,8 @@ class Envs: # Unified Radix Tree SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False) + # Registered TreeCore backend serving the unified radix cache. + SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND = EnvStr("python") # CUDA Graph SGLANG_USE_BREAKABLE_CUDA_GRAPH = EnvBool(False) diff --git a/python/sglang/srt/kv_canary/radix_cache_walker.py b/python/sglang/srt/kv_canary/radix_cache_walker.py index 1b155ff5d..718e67e58 100644 --- a/python/sglang/srt/kv_canary/radix_cache_walker.py +++ b/python/sglang/srt/kv_canary/radix_cache_walker.py @@ -1,29 +1,19 @@ from __future__ import annotations -from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch from sglang.srt.mem_cache.radix_cache import RadixCache from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache -from sglang.srt.mem_cache.unified_cache_components import ( - BASE_COMPONENT_TYPE, - ComponentType, +from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + RadixCacheWalkResult, ) from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache if TYPE_CHECKING: from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.radix_cache import TreeNode - from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode - - -@dataclass(frozen=True, slots=True, kw_only=True) -class RadixCacheWalkResult: - slot_indices: torch.Tensor - positions: torch.Tensor - prev_slot_indices: torch.Tensor def walk_radix_cache_for_canary( @@ -40,11 +30,11 @@ def walk_radix_cache_for_canary( req. ``swa_resident_only=True`` skips SWA-tombstoned nodes (slots evicted from the SWA window).""" cache_type = type(radix_cache) - if ( - cache_type is not RadixCache - and cache_type is not SWARadixCache - and cache_type is not UnifiedRadixCache - ): + if cache_type is UnifiedRadixCache: + return radix_cache.tree_core.walk_for_kv_canary( + unlocked_only=unlocked_only, swa_resident_only=swa_resident_only + ) + if cache_type is not RadixCache and cache_type is not SWARadixCache: raise NotImplementedError( f"walk_radix_cache_for_canary does not support {cache_type.__name__}" ) @@ -78,7 +68,7 @@ def walk_radix_cache_for_canary( def _walk_radix_subtree( *, - node: TreeNode | UnifiedTreeNode, + node: TreeNode, radix_cache: BasePrefixCache, depth: int, parent_last_slot: int, @@ -89,13 +79,11 @@ def _walk_radix_subtree( unlocked_only: bool, swa_resident_only: bool, ) -> None: - node_slots = _node_slots_for_canary(node=node, radix_cache=radix_cache) + node_slots = _node_slots_for_canary(node=node) if unlocked_only: emit_slots = not is_root and _node_is_unlocked_for_canary( - node=node, - radix_cache=radix_cache, - swa_resident_only=swa_resident_only, + node=node, radix_cache=radix_cache ) else: emit_slots = not is_root @@ -114,12 +102,7 @@ def _walk_radix_subtree( prev_slot_buf.append(prev) chain_last_slot = slot - child_depth = depth + _node_len_for_canary( - node=node, - radix_cache=radix_cache, - node_slots=node_slots, - is_root=is_root, - ) + child_depth = depth + len(node_slots) for child in node.children.values(): _walk_radix_subtree( node=child, @@ -135,42 +118,17 @@ def _walk_radix_subtree( ) -def _node_slots_for_canary( - *, - node: TreeNode | UnifiedTreeNode, - radix_cache: BasePrefixCache, -) -> list[int]: - value: Any - if type(radix_cache) is UnifiedRadixCache: - value = node.component_data[BASE_COMPONENT_TYPE].value - else: - value = node.value - +def _node_slots_for_canary(*, node: TreeNode) -> list[int]: + value: Any = node.value if isinstance(value, torch.Tensor): return [int(s) for s in value.tolist()] return [] -def _node_len_for_canary( - *, - node: TreeNode | UnifiedTreeNode, - radix_cache: BasePrefixCache, - node_slots: list[int], - is_root: bool, -) -> int: - if type(radix_cache) is not UnifiedRadixCache: - return len(node_slots) - - if is_root or node.key is None: - return len(node_slots) - return len(node.key) - - def _node_is_unlocked_for_canary( *, - node: TreeNode | UnifiedTreeNode, + node: TreeNode, radix_cache: BasePrefixCache, - swa_resident_only: bool, ) -> bool: if type(radix_cache) is RadixCache: return node.lock_ref == 0 @@ -178,13 +136,6 @@ def _node_is_unlocked_for_canary( if type(radix_cache) is SWARadixCache: return node.full_lock_ref == 0 - if type(radix_cache) is UnifiedRadixCache: - if swa_resident_only and radix_cache.supports_swa(): - # Unified SWA owns an independent component lock. A node can still - # hold Full KV for a running request while its SWA slots are unused. - return node.component_data[ComponentType.SWA].lock_ref == 0 - return node.component_data[BASE_COMPONENT_TYPE].lock_ref == 0 - raise NotImplementedError( f"walk_radix_cache_for_canary does not support {type(radix_cache).__name__}" ) @@ -192,15 +143,10 @@ def _node_is_unlocked_for_canary( def _node_is_swa_resident_for_canary( *, - node: TreeNode | UnifiedTreeNode, + node: TreeNode, radix_cache: BasePrefixCache, ) -> bool: if type(radix_cache) is SWARadixCache: return not node.swa_tombstone - if type(radix_cache) is UnifiedRadixCache: - if not radix_cache.supports_swa(): - return True - return node.component_data[ComponentType.SWA].value is not None - return True diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 96485fe35..f3e240c1a 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1249,7 +1249,9 @@ class Req(ReqDllmMixin): ) ) if envs.SGLANG_RADIX_FORCE_MISS.get(): - match_result = zero_match_result(tree_cache, match_result) + match_result = zero_match_result( + tree_cache, match_result, extra_key=self.extra_key + ) ( self.prefix_indices, self.last_node, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 325507926..67a9243aa 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -115,7 +115,9 @@ def match_prefix_for_req( ) ) if envs.SGLANG_RADIX_FORCE_MISS.get(): - match_result = zero_match_result(tree_cache, match_result) + match_result = zero_match_result( + tree_cache, match_result, extra_key=req.extra_key + ) ( req.prefix_indices, req.last_node, @@ -290,7 +292,7 @@ class SchedulePolicy: ) if envs.SGLANG_RADIX_FORCE_MISS.get(): match_result = zero_match_result( - self.waiting_queue_radix_tree, match_result + self.waiting_queue_radix_tree, match_result, extra_key=extra_key ) in_batch_matching_prefixes = match_result.device_indices if ( @@ -328,7 +330,8 @@ class SchedulePolicy: """Sorts the waiting queue based on a depth-first search weighting.""" last_node_to_reqs = defaultdict(list) for req in waiting_queue: - last_node_to_reqs[req.last_node].append(req) + last_node = tree_cache.resolve_node_handle(req.last_node) + last_node_to_reqs[last_node].append(req) node_to_weight = defaultdict(int) for node in last_node_to_reqs: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3fc35ee2e..aaeda72df 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2418,25 +2418,25 @@ class Scheduler( 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_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() + tree_cache = self.tree_cache + if tree_cache.is_backuped(req.last_host_node) or tree_cache.is_root( + req.last_host_node + ): matched_len = len(req.prefix_indices) + req.host_hit_length match_end = req._compute_max_prefix_len( len(req.full_untruncated_fill_ids) ) new_input_tokens = req.full_untruncated_fill_ids[matched_len:match_end] - prefix_keys = ( - last_host_node.get_prefix_hash_values(last_host_node.parent) - if self.tree_cache.hicache_storage_pass_prefix_keys + tree_cache.get_prefix_hash_values(req.last_host_node) + if tree_cache.hicache_storage_pass_prefix_keys else None ) - self.tree_cache.prefetch_from_storage( + tree_cache.prefetch_from_storage( req.rid, - last_host_node, + req.last_host_node, new_input_tokens, - last_hash, + tree_cache.get_last_hash_value(req.last_host_node), prefix_keys, ) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index d99e9adfa..244eef333 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -9,6 +9,7 @@ from typing import ( NamedTuple, Optional, Protocol, + Sequence, Tuple, runtime_checkable, ) @@ -26,6 +27,10 @@ from sglang.srt.observability.metrics_collector import ( if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.radix_cache import RadixKey + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( ComponentType, ) @@ -78,6 +83,10 @@ class InsertResult: last_device_node: Any = None mamba_exist: bool = False inserted_host_node: Any = None + # Controller-applied actions from the non-stepped channels (e.g. insert_host); the stepped insert emits via InsertStepResult.actions. + cache_actions: list[CacheAction | ComponentAction] = dataclasses.field( + default_factory=list + ) @dataclasses.dataclass @@ -191,13 +200,17 @@ class MatchResult(NamedTuple): mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None full_kv_hit_length: int = 0 + # Actions the Controller applies: CacheActions itself, ComponentActions routed to the owning component. + cache_actions: Sequence[CacheAction | ComponentAction] = () -def zero_match_result(tree_cache, match_result: MatchResult) -> MatchResult: +def zero_match_result( + tree_cache, match_result: MatchResult, extra_key: Optional[str] = None +) -> MatchResult: if tree_cache.is_chunk_cache(): # Chunk caches' match_prefix already returns a miss; no root_node to walk back to. return match_result - root = tree_cache.root_node + root = tree_cache.root_node_handle(extra_key=extra_key) return match_result._replace( # [:0] keeps dtype and device of the original tensor (e.g. CUDA int64) # without allocating a fresh empty tensor. @@ -258,6 +271,37 @@ class BasePrefixCache(ABC, PrefixCacheTrait): def supports_fast_match_prefix(self) -> bool: return False + def resolve_node_handle(self, node_handle: Any) -> Any: + """Map a node handle to its node -- e.g. UnifiedRadixCache looks up the + node object from its NodeId. Temporary API for the Unified Radix Cache + split migration. + + TODO(Jialin): Remove after the Unified Radix Cache split. + """ + return node_handle + + def root_node_handle(self, extra_key: Optional[str] = None) -> Any: + """The root handle as match results carry it -- the raw node by default, + the root's NodeId for UnifiedRadixCache. extra_key scopes the root for + implementations that shard trees per cache namespace.""" + return self.root_node + + def is_backuped(self, node: Any) -> bool: + """Whether the node's Full KV is present on host.""" + return node.backuped + + def is_root(self, node: Any) -> bool: + """Whether the node is a tree root.""" + return node is self.root_node + + def get_last_hash_value(self, node: Any) -> Optional[str]: + """The node's last page hash, or None when it was never hashed.""" + return node.get_last_hash_value() + + def get_prefix_hash_values(self, node: Any) -> list[str]: + """The hash chain of the node's ancestors, in root-to-parent order.""" + return node.get_prefix_hash_values(node.parent) + @abstractmethod def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs): pass diff --git a/python/sglang/srt/mem_cache/unified_cache/__init__.py b/python/sglang/srt/mem_cache/unified_cache/__init__.py new file mode 100644 index 000000000..15c752c54 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/__init__.py @@ -0,0 +1,18 @@ +"""Components of Unified Radix Cache. + +- **TreeCore** (``unified_tree_core.UnifiedTreeCore``): the pure tree mechanism. Owns + the tree structure, per-node values, the LRU(s), bookkeeping, and KV-cache + events -- the only object that touches that state; no scheduling policy or IO, + and never touches the cache. References the cache-owned component drivers to + drive their tree-level hooks. +- **Components** (``unified_cache_components.TreeComponent`` subclasses): the + per-component (FULL/SWA/MAMBA) drivers. Built and owned by the Controller; + hold the cache for cache-level logic and the TreeCore for tree state. +- **Controller** (``unified_radix_cache.UnifiedRadixCache``): the scheduling brain + and scheduler-facing facade. Owns the scheduling policy (when to evict / backup / + load / prefetch), the IO (HiCache controller, device/host allocators, and the + device<->host transfer), the async ack machinery, prefetch/storage, the + distributed groups, and the component drivers; holds and drives the TreeCore. + +TODO(Jialin): move all unified radix cache logic into this folder. +""" diff --git a/python/sglang/srt/mem_cache/unified_cache/cache_action.py b/python/sglang/srt/mem_cache/unified_cache/cache_action.py new file mode 100644 index 000000000..efea64b41 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/cache_action.py @@ -0,0 +1,103 @@ +"""A TreeCore emits CacheActions through the TreeCoreInterface to guide +Controller behavior. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import msgspec + +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType + +if TYPE_CHECKING: + import torch + + from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import NodeId + + +class ReplaceWriteThroughOnNodeSplit(msgspec.Struct, frozen=True): + """Replace the pending write-through node on a node split: + + parent -> node => parent -> new_node -> new_child + + old_node_id (the pre-split node) is replaced by new_node_id + new_child_node_id. + """ + + ack_id: int + old_node_id: NodeId + new_node_id: NodeId + new_child_node_id: NodeId + + +class FreeDeviceKV(msgspec.Struct, frozen=True): + """Free unreferenced device KV slots (a SWA-aware combined free).""" + + indices: list[torch.Tensor] + + +class ComponentAction(msgspec.Struct, frozen=True): + """Base for component-routed actions; the cache dispatches each one to + ``component_type``'s class-level ``apply_component_action``; every subclass + declares its ``component_type`` field.""" + + +class FreeComponentDeviceSlot(ComponentAction, frozen=True): + """Free only the given ``component_type``'s device KV slots.""" + + indices: list[torch.Tensor] + component_type: ComponentType + + +class FreeComponentHostSlot(ComponentAction, frozen=True): + """Free the given ``component_type``'s host KV pages.""" + + host_indices: list[torch.Tensor] + component_type: ComponentType + + +class BackupKV(msgspec.Struct, frozen=True): + """Back up node_ids device->host in order, stopping at the first failure; write-through + ids form a contiguous root-first parent chain (each id's parent precedes it), write-back + actions carry a single eviction victim.""" + + node_ids: list[NodeId] + + +class MambaEvictExcessPathStates(ComponentAction, frozen=True): + """Per-path Mamba state-cap eviction from the tail's root path; applied at + the insert's commit barrier, after the walk-time backups whose + write-through locks shield the backed-up chain.""" + + tail_node_id: NodeId + component_type: ComponentType = ComponentType.MAMBA + + +class RebuildFullToSWAMapping(ComponentAction, frozen=True): + """Rebuild the SWA allocator's full->swa index mapping for loaded chunks.""" + + full_indices: list[torch.Tensor] + swa_indices: list[torch.Tensor] + component_type: ComponentType = ComponentType.SWA + + +class RecoverSWAWithLockedFull(ComponentAction, frozen=True): + """Recover an SWA tombstone whose full is locked: keep the locked full, remap it + onto the incoming full's SWA translation, and free only the incoming full.""" + + node_id: NodeId + kept_full: torch.Tensor + incoming_full: torch.Tensor + component_type: ComponentType = ComponentType.SWA + + +class SWARebuild(ComponentAction, frozen=True): + """Rebuild a node's SWA value by translating its source full value, then store it.""" + + node_id: NodeId + source_value: torch.Tensor + component_type: ComponentType = ComponentType.SWA + + +# Cache-owned actions, applied by UnifiedRadixCache itself. +CacheAction = ReplaceWriteThroughOnNodeSplit | FreeDeviceKV | BackupKV diff --git a/python/sglang/srt/mem_cache/unified_cache/component_type.py b/python/sglang/srt/mem_cache/unified_cache/component_type.py new file mode 100644 index 000000000..625e4d782 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/component_type.py @@ -0,0 +1,29 @@ +"""The per-attention component identity; a leaf module importable from anywhere.""" + +from enum import Enum + + +class ComponentType(int, Enum): + """Integer enum so that per-node list/tuple storage can be indexed directly.""" + + FULL = 0 + SWA = 1 + MAMBA = 2 + + def __str__(self) -> str: # keep human-readable logging + return self.name.lower() + + @property + def is_full(self) -> bool: + return self == ComponentType.FULL + + @property + def is_swa(self) -> bool: + return self == ComponentType.SWA + + @property + def is_mamba(self) -> bool: + return self == ComponentType.MAMBA + + +BASE_COMPONENT_TYPE = ComponentType.FULL diff --git a/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py b/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py new file mode 100644 index 000000000..f7b9871fb --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py @@ -0,0 +1,73 @@ +"""Registry for pluggable TreeCore implementations. + +The unified cache constructs its TreeCore through `create_tree_core`, selected +by SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND (default "python"). To plug in a custom +implementation, register it under a string name via +`register_tree_core_backend(name, factory)`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Optional + +if TYPE_CHECKING: + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_cache.component_type import ComponentType + from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + UnifiedTreeCoreInterface, + ) + from sglang.srt.mem_cache.unified_cache_components import TreeComponent + +TreeCoreFactory = Callable[ + ["CacheInitParams", "dict[ComponentType, TreeComponent]"], + "UnifiedTreeCoreInterface", +] + +_TREE_CORE_REGISTRY: dict[str, TreeCoreFactory] = {} + + +def register_tree_core_backend(name: str, factory: TreeCoreFactory) -> None: + """Register a TreeCore factory under `name`.""" + if not name.strip(): + raise ValueError( + f"register_tree_core_backend: name must be non-empty, got {name!r}" + ) + if name in _TREE_CORE_REGISTRY: + raise ValueError(f"register_tree_core_backend: {name!r} is already registered") + _TREE_CORE_REGISTRY[name] = factory + + +def get_tree_core_factory(name: str) -> Optional[TreeCoreFactory]: + return _TREE_CORE_REGISTRY.get(name) + + +def registered_tree_core_backends() -> list[str]: + return list(_TREE_CORE_REGISTRY.keys()) + + +def _python_tree_core_factory( + params: CacheInitParams, components: dict[ComponentType, TreeComponent] +) -> UnifiedTreeCoreInterface: + """The pure-Python TreeCore.""" + from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore + + return UnifiedTreeCore(params, components) + + +register_tree_core_backend("python", _python_tree_core_factory) + + +def create_tree_core( + name: str, + params: CacheInitParams, + components: dict[ComponentType, TreeComponent], +) -> UnifiedTreeCoreInterface: + """Construct the TreeCore registered under `name`.""" + factory = get_tree_core_factory(name) + if factory is None: + raise ValueError( + f"SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND={name!r} is not registered. " + f"Registered backends: {registered_tree_core_backends()}. " + "External backends must call register_tree_core_backend(...) at import time." + ) + return factory(params, components) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py new file mode 100644 index 000000000..df086e04c --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -0,0 +1,2150 @@ +"""The radix tree mechanism. + +``UnifiedTreeCore`` owns the tree's member-var state -- the tree structure (root +node), the per-component LRU lists, the per-component size counters, the +evictable device/host leaf sets, and the empty match result -- plus ``reset()``. +It also defines the tree's building blocks (``UnifiedTreeNode``, +``UnifiedLRUList``). + +The cache builds the component drivers and passes them in; the tree holds them +(``self.components_by_type``) to drive their tree-level hooks. The components hold the +cache for cache-level logic, but the TreeCore itself never touches it. +""" + +# TODO(Jialin): Split this file into smaller cohesive modules (e.g. tree +# building blocks, insert/match walks, eviction drivers). + +from __future__ import annotations + +import logging +import sys +from array import array +from collections import defaultdict +from enum import Enum, auto +from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Sequence + +import msgspec +import torch + +from sglang.srt.disaggregation.kv_events import StorageMedium +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + DecLockRefResult, + IncLockRefResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.hicache_storage import ( + PoolName, + PoolTransfer, + PoolTransferResult, +) +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.unified_cache.cache_action import ( + BackupKV, + CacheAction, + ComponentAction, + FreeDeviceKV, + ReplaceWriteThroughOnNodeSplit, +) +from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + DecSwaLockOnlyResult, + DemoteResult, + DriveHostEvictionResult, + DropSubtreeNoHostResult, + EvictDeviceLeafResult, + EvictDeviceNextNodeResult, + InsertStepResult, + NodeId, + RadixCacheWalkResult, + UnifiedTreeCoreInterface, +) +from sglang.srt.mem_cache.unified_cache_components import ( + _NUM_COMPONENT_TYPES, + BASE_COMPONENT_TYPE, + CacheTransferPhase, + ComponentData, + ComponentType, + EvictLayer, + LRURefreshPhase, + TreeComponent, + get_and_increase_time_counter, +) +from sglang.srt.mem_cache.utils import ( + compute_node_hash_values, + get_eviction_strategy, + split_node_hash_value, +) + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + +logger = logging.getLogger(__name__) + + +class StorageBackupSpec(NamedTuple): + """A node's device->storage backup spec, gathered tree-side.""" + + host_value: torch.Tensor + token_ids: array + hash_value: list[str] + prefix_keys: Optional[list[str]] + comp_xfers: dict[ComponentType, list[PoolTransfer]] + + +class UnifiedTreeNode: + counter = 0 + + def __init__(self, tree_components: tuple[ComponentType, ...], priority: int = 0): + # Plain dict (not defaultdict): a missing-key read must raise, never + # silently mint an unregistered node outside the TreeCore arena. + self.children: dict[Any, UnifiedTreeNode] = {} + self.parent: UnifiedTreeNode | None = None + self.key: Optional[RadixKey] = None + self.component_types = tree_components + # list indexed by ComponentType (int enum 0..N-1) + self.component_data: list[ComponentData] = [ + ComponentData() for _ in range(_NUM_COMPONENT_TYPES) + ] + self.last_access_time = get_and_increase_time_counter() + self.creation_time = get_and_increase_time_counter() + self.hash_value = None + self.hit_count = 0 + self.priority = priority + self.lru_prev: list[UnifiedTreeNode | None] = [None] * ( + _NUM_COMPONENT_TYPES * 2 + ) + self.lru_next: list[UnifiedTreeNode | None] = [None] * ( + _NUM_COMPONENT_TYPES * 2 + ) + self.id = UnifiedTreeNode.counter + UnifiedTreeNode.counter += 1 + self.write_through_pending_id: Optional[int] = None + + def component(self, component_type: ComponentType) -> ComponentData: + return self.component_data[component_type] + + @property + def backuped(self) -> bool: + """Tree-level: Full KV present on host.""" + return self.component_data[ComponentType.FULL].host_value is not None + + @property + def evicted(self) -> bool: + """Tree-level: Full KV not on device (non-root with value=None).""" + return ( + self.parent is not None + and self.component_data[ComponentType.FULL].value is None + ) + + def __lt__(self, other: UnifiedTreeNode): + return self.last_access_time < other.last_access_time + + def get_last_hash_value(self) -> Optional[str]: + if self.hash_value is None or len(self.hash_value) == 0: + return None + return self.hash_value[-1] + + def get_prefix_hash_values(self, node: UnifiedTreeNode) -> list[str]: + if node is None or node.hash_value is None: + return [] + + return node.get_prefix_hash_values(node.parent) + node.hash_value + + +class UnifiedLRUList: + def __init__( + self, + component_type: ComponentType, + tree_components: tuple[ComponentType, ...], + use_host_ptr: bool = False, + ): + self.component_type = component_type + # Pointer slot: host LRU uses offset slots so device/host pointers + # never collide on the same node. + self._pt: int = component_type + (_NUM_COMPONENT_TYPES if use_host_ptr else 0) + self.head = UnifiedTreeNode(tree_components) + self.tail = UnifiedTreeNode(tree_components) + self.head.lru_next[self._pt] = self.tail + self.tail.lru_prev[self._pt] = self.head + self.cache: dict[int, UnifiedTreeNode] = {} + + def _add_node_after(self, prev_node: UnifiedTreeNode, new_node: UnifiedTreeNode): + pt = self._pt + new_node.lru_prev[pt] = prev_node + new_node.lru_next[pt] = prev_node.lru_next[pt] + prev_node.lru_next[pt].lru_prev[pt] = new_node + prev_node.lru_next[pt] = new_node + + def _add_node(self, node: UnifiedTreeNode): + self._add_node_after(self.head, node) + + def _remove_node(self, node: UnifiedTreeNode): + pt = self._pt + node.lru_prev[pt].lru_next[pt] = node.lru_next[pt] + node.lru_next[pt].lru_prev[pt] = node.lru_prev[pt] + # Clear self pointers to break reference cycles among evicted nodes. + node.lru_prev[pt] = None + node.lru_next[pt] = None + + def insert_mru(self, node: UnifiedTreeNode): + assert node.id not in self.cache + self.cache[node.id] = node + self._add_node(node) + + def remove_node(self, node: UnifiedTreeNode): + assert node.id in self.cache + del self.cache[node.id] + self._remove_node(node) + + def reset_node_mru(self, node: UnifiedTreeNode): + assert node.id in self.cache + self._remove_node(node) + self._add_node(node) + + def reset_node_and_parents_mru( + self, + node: UnifiedTreeNode, + root_node: UnifiedTreeNode, + should_include, + ): + prev_node = self.head + while node != root_node: + if should_include(node): + assert node.id in self.cache + self._remove_node(node) + self._add_node_after(prev_node, node) + prev_node = node + node = node.parent + + def reset_node_and_window_ancestors_mru( + self, + node: UnifiedTreeNode, + root_node: UnifiedTreeNode, + window_size: int, + should_include, + ): + prev_node = self.head + accumulated = 0 + while node != root_node and accumulated < window_size: + if should_include(node): + assert node.id in self.cache + self._remove_node(node) + self._add_node_after(prev_node, node) + prev_node = node + accumulated += len(node.key) + node = node.parent + + def in_list(self, node: Optional[UnifiedTreeNode]): + return node is not None and node.id in self.cache + + def get_prev_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): + if check_id: + assert node.id in self.cache + pt = self._pt + ct = self.component_type + x = node.lru_prev[pt] + while x.component_data[ct].lock_ref > 0: + x = x.lru_prev[pt] + if x == self.head: + return None + return x + + def get_prev_leaf_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): + if check_id: + assert node.id in self.cache + pt = self._pt + ct = self.component_type + x = node.lru_prev[pt] + while x.component_data[ct].lock_ref > 0 or len(x.children) > 0: + x = x.lru_prev[pt] + if x == self.head: + return None + return x + + def get_prev_no_host_lock(self, node: UnifiedTreeNode, check_id: bool = True): + """Host-LRU walker: skip nodes whose component host_lock_ref > 0.""" + if check_id: + assert node.id in self.cache + pt = self._pt + ct = self.component_type + x = node.lru_prev[pt] + while x.component_data[ct].host_lock_ref > 0: + x = x.lru_prev[pt] + if x == self.head: + return None + return x + + def get_lru_no_lock(self): + return self.get_prev_no_lock(self.tail, check_id=False) + + def get_leaf_lru_no_lock(self): + return self.get_prev_leaf_no_lock(self.tail, check_id=False) + + def get_lru_no_host_lock(self): + return self.get_prev_no_host_lock(self.tail, check_id=False) + + +# WALK (one node per step) -> COMMIT (leaf + commit hooks) -> TAIL (refresh + backup). +class _InsertPhase(Enum): + WALK = auto() + COMMIT = auto() + TAIL = auto() + + +class _InsertWalkState(msgspec.Struct): + """In-flight resumable-insert state persisted across step barriers.""" + + phase: _InsertPhase + node: UnifiedTreeNode + key: RadixKey + value: torch.Tensor + params: InsertParams + priority: int + total_prefix_length: int = 0 + is_new_leaf: bool = False + target_node: Optional[UnifiedTreeNode] = None + result: Optional[InsertResult] = None + # Emitted actions awaiting the next barrier flush (or the final step). + pending_actions: list[CacheAction | ComponentAction] = [] + + +class UnifiedTreeCore(UnifiedTreeCoreInterface): + """The radix tree mechanism: owns the tree structure, per-node values, the + per-component LRUs, the size/leaf bookkeeping, and the component drivers, + plus ``reset()``. + + TODO(Jialin): the tree operations still live on ``UnifiedRadixCache`` and + reach this state through its proxy properties; they migrate onto this class + as the TreeCore split completes. + """ + + def __init__( + self, + params: CacheInitParams, + components: dict[ComponentType, TreeComponent], + ): + self.page_size = params.page_size + self.is_eagle = params.is_eagle and ComponentType.MAMBA not in components + self.enable_hicache = False + self.enable_storage = False + self.write_through_threshold = 256 + self.is_write_back = False + self.has_swa_host_pool = False + self.eviction_strategy = get_eviction_strategy(params.eviction_policy.lower()) + + # ``device`` is derived from the construction-time allocator; the + # allocator/pool themselves are owned by the cache, not the tree. + if params.token_to_kv_pool_allocator: + self.device = params.token_to_kv_pool_allocator.device + else: + self.device = torch.device("cpu") + + # The cache builds and owns the component drivers; the tree references + # them to drive their tree-level hooks, attaching itself as tree_core. + assert components + self.component_types = tuple(components.keys()) + self.components_by_type: dict[ComponentType, TreeComponent] = components + for component in components.values(): + component.tree_core = self + self.components: tuple[TreeComponent, ...] = tuple( + self.components_by_type.values() + ) + + self.enable_kv_cache_events = params.enable_kv_cache_events + self.kv_event_queue = [] + + self.reset() + + # ==== Tree API ==== + + def reset(self) -> None: + """Rebuild the root, LRUs, sizes, evictable-leaf sets, and the empty + match result.""" + # Maintains the NodeId -> active tree node mapping. + self._node_arena: dict[NodeId, UnifiedTreeNode] = {} + + # The single in-flight resumable insert, if suspended at a barrier. + self._ongoing_insert_walk_state: Optional[_InsertWalkState] = None + + self.root_node = self._new_node() + self.root_node.priority = -sys.maxsize + self.root_node.key = RadixKey(array("q"), None) + self.root_node.component_data[BASE_COMPONENT_TYPE].value = [] + self.root_node.hash_value = [] + for ct in self.component_types: + self.root_node.component_data[ct].lock_ref = 1 + + self.component_evictable_size_ = {ct: 0 for ct in self.component_types} + self.component_protected_size_ = {ct: 0 for ct in self.component_types} + + self.lru_lists = { + ct: UnifiedLRUList(ct, self.component_types) for ct in self.component_types + } + + self.evictable_device_leaves: set[UnifiedTreeNode] = set() + self.evictable_host_leaves: set[UnifiedTreeNode] = set() + self.host_lru_lists = { + ct: UnifiedLRUList(ct, self.component_types, use_host_ptr=True) + for ct in self.component_types + } + + self._empty_match_result = MatchResult( + device_indices=torch.empty( + (0,), + dtype=torch.int64, + device=self.device, + ), + last_device_node=self.root_node.id, + last_host_node=self.root_node.id, + best_match_node=self.root_node.id, + cache_actions=[], + ) + + def node_by_id(self, node_id: NodeId) -> UnifiedTreeNode: + """Resolve a NodeId back to its tree node. + + TODO(Jialin): Make TreeCore-internal after the Unified Radix Cache split. + """ + return self._node_arena[node_id] + + def is_backuped(self, node_id: NodeId) -> bool: + """Whether the node's KV is already backed up to host.""" + return self._node_arena[node_id].backuped + + def is_root(self, node_id: NodeId) -> bool: + """Whether the node is the tree root.""" + return self.node_by_id(node_id) is self.root_node + + def get_last_hash_value(self, node_id: NodeId) -> Optional[str]: + """The node's last page hash, or None when it was never hashed.""" + return self.node_by_id(node_id).get_last_hash_value() + + def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: + """The hash chain of the node's ancestors, in root-to-parent order.""" + node = self.node_by_id(node_id) + return node.get_prefix_hash_values(node.parent) + + def _new_node(self, priority: int = 0) -> UnifiedTreeNode: + """Create and register a tree node in the arena.""" + node = UnifiedTreeNode(self.component_types, priority=priority) + self._register_node(node) + return node + + def _register_node(self, node: UnifiedTreeNode) -> None: + """Register a tree node in the arena.""" + self._node_arena[node.id] = node + + def _unregister_node(self, node: UnifiedTreeNode) -> None: + """Drop a tree node from the arena.""" + self._node_arena.pop(node.id, None) + + def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult: + node = self.node_by_id(node_id) + result = IncLockRefResult() + for component in self.components: + result = component.acquire_component_lock(node=node, result=result) + self._update_evictable_leaf_sets(node) + return result + + def dec_lock_ref( + self, + node_id: NodeId, + params: Optional[DecLockRefParams] = None, + skip_swa: bool = False, + ) -> DecLockRefResult: + node = self.node_by_id(node_id) + for component in self.components: + if skip_swa and component.component_type == ComponentType.SWA: + continue + component.release_component_lock(node=node, params=params) + self._update_evictable_leaf_sets(node) + # TODO: delta is not aggregated from components; no caller uses it yet. + return DecLockRefResult() + + def dec_swa_lock_only( + self, node_id: NodeId, swa_uuid_for_lock: Optional[int] + ) -> DecSwaLockOnlyResult: + """Early-release the SWA portion of a request's tree lock, plus any + strictly-lower-priority locks (e.g. Mamba) co-located on the node.""" + result = DecSwaLockOnlyResult() + node = self.node_by_id(node_id) + swa_component = self.components_by_type.get(ComponentType.SWA) + if swa_component is None: + return result + swa_component.release_window_lock( + node, swa_uuid_for_lock, result.device_frees, result.host_frees + ) + + # Drop strictly-lower-priority locks (e.g. Mamba) co-located on the node. + swa_priority = swa_component.eviction_priority(is_leaf=False) + dec_params = DecLockRefParams(swa_uuid_for_lock=swa_uuid_for_lock) + for comp in self.components: + if comp.eviction_priority(is_leaf=False) < swa_priority: + comp.release_component_lock(node, dec_params) + return result + + def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: + node = self.node_by_id(node_id) + result = IncLockRefResult() + for component in self.components: + result = component.acquire_component_lock( + node=node, result=result, lock_host=True + ) + self._update_evictable_leaf_sets(node) + return result + + def dec_host_lock_ref( + self, node_id: NodeId, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: + node = self.node_by_id(node_id) + for component in self.components: + component.release_component_lock(node=node, params=params, lock_host=True) + self._update_evictable_leaf_sets(node) + return DecLockRefResult() + + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: + key = params.key + key, _ = key.maybe_to_bigram_view(self.is_eagle) + if len(key) == 0: + return self._empty_match_result + key = key.page_aligned(self.page_size) + if len(key) == 0: + return self._empty_match_result + + ( + value, + best_match_node, + best_match_device_node, + best_match_device_value_len, + full_kv_hit_length, + action, + ) = self._match_prefix_helper(key) + return self._match_post_processor( + params, + value, + best_match_node, + best_match_device_node, + best_match_device_value_len, + full_kv_hit_length, + action, + ) + + def _match_prefix_helper(self, key: RadixKey) -> tuple[ + list[torch.Tensor], + UnifiedTreeNode, + UnifiedTreeNode, + int, + int, + Optional[CacheAction | ComponentAction], + ]: + # Non-HiCache mode has only device-resident matches, so the scheduler + # device anchor follows the best match. In HiCache mode, host-backed + # nodes can also match, so we separately track the best device-resident + # match for scheduler prefix indices and locking. + node = self.root_node + child_key = key.child_key(self.page_size) + value: list[torch.Tensor] = [] + best_match_node = node + best_match_device_node = node + best_match_device_value_len = 0 + full_kv_hit_length = 0 + action: Optional[CacheAction | ComponentAction] = None + separate_device_match = self.enable_hicache + if separate_device_match: + validators = tuple( + comp.create_match_validator() for comp in self.components + ) + device_validators = tuple( + comp.create_match_validator(match_device_only=True) + for comp in self.components + ) + else: + validators = tuple( + comp.create_match_validator(match_device_only=True) + for comp in self.components + ) + + def _all_valid(validators, node): + return all([v(node) for v in validators]) + + def _update_best_if_valid(node): + nonlocal best_match_node + nonlocal best_match_device_value_len, best_match_device_node + matched = _all_valid(validators, node) + if matched: + best_match_node = node + + if not separate_device_match: + if matched: + best_match_device_value_len = len(value) + best_match_device_node = node + return + if _all_valid(device_validators, node): + best_match_device_value_len = len(value) + best_match_device_node = node + + while len(key) > 0 and child_key in node.children: + child = node.children[child_key] + + # HiCache: dead node (evicted + not backuped) — stop traversal + if child.evicted and not child.backuped: + break + + prefix_len = child.key.match(key, page_size=self.page_size) + full_kv_hit_length += prefix_len + if prefix_len < len(child.key): + node, action = self._split_node(child.key, child, prefix_len) + if not node.evicted: + value.append(node.component_data[BASE_COMPONENT_TYPE].value) + _update_best_if_valid(node) + break + + if not child.evicted: + value.append(child.component_data[BASE_COMPONENT_TYPE].value) + node = child + _update_best_if_valid(node) + key = key[prefix_len:] + if len(key): + child_key = key.child_key(self.page_size) + + return ( + value, + best_match_node, + best_match_device_node, + best_match_device_value_len, + full_kv_hit_length, + action, + ) + + def _match_post_processor( + self, + params: MatchPrefixParams, + value: list[torch.Tensor], + best_match_node: UnifiedTreeNode, + best_match_device_node: UnifiedTreeNode, + best_match_device_value_len: int, + full_kv_hit_length: int, + action: Optional[CacheAction | ComponentAction], + ) -> MatchResult: + node_update = best_match_node + for comp in self.components: + if comp.component_type == BASE_COMPONENT_TYPE: + continue # Full uses last_access_time, not LRU + comp.refresh_lru(LRURefreshPhase.MATCH_END, node_update, self.root_node) + + cur_time = get_and_increase_time_counter() + while node_update: + node_update.last_access_time = cur_time + cur_time -= 0.00001 + node_update = node_update.parent + + # last_host_node will be used as the starting node for the subsequent + # `prefetch_from_storage` flow. We directly use best_match_node here, + # because best_match_node represents the node where all components + # have reached consensus on both device & host availability. + last_host_node = ( + best_match_node if self.enable_hicache else best_match_device_node + ) + + if best_match_device_value_len > 0: + device_indices = torch.cat(value[:best_match_device_value_len]) + else: + device_indices = self._empty_match_result.device_indices + result = MatchResult( + device_indices=device_indices, + last_device_node=best_match_device_node, + last_host_node=last_host_node, + best_match_node=best_match_node, + host_hit_length=0, + full_kv_hit_length=full_kv_hit_length, + ) + + for component in self.components: + result = component.finalize_match_result_in_tree_core( + result=result, + params=params, + value_chunks=value, + best_value_len=best_match_device_value_len, + ) + # Expose only NodeIds outside TreeCore. + return result._replace( + last_device_node=result.last_device_node.id, + last_host_node=result.last_host_node.id, + best_match_node=result.best_match_node.id, + cache_actions=[action] if action is not None else [], + ) + + @property + def empty_match_result(self) -> MatchResult: + """A shared empty MatchResult (empty device indices + boundary NodeIds).""" + return self._empty_match_result + + def is_full_device_evicted(self, node_id: NodeId) -> bool: + """Whether the node's FULL device value has been evicted.""" + return self.node_by_id(node_id).evicted + + def collect_full_device_indices( + self, from_node_id: NodeId, until_node_id: NodeId + ) -> torch.Tensor: + """Concatenate FULL device values from ``from_node`` up to (exclusive) + ``until_node``, in root order; empty tensor if the path is empty.""" + until_node = self.node_by_id(until_node_id) + prefix_chunks: list[torch.Tensor] = [] + node = self.node_by_id(from_node_id) + while node is not until_node: + value = node.component_data[BASE_COMPONENT_TYPE].value + assert value is not None + prefix_chunks.append(value) + node = node.parent + if not prefix_chunks: + return self._empty_match_result.device_indices + prefix_chunks.reverse() + return torch.cat(prefix_chunks) + + def _touch_node(self, node: UnifiedTreeNode): + node.last_access_time = get_and_increase_time_counter() + if node != self.root_node: + for comp in self.components: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + comp.refresh_lru(LRURefreshPhase.WALKDOWN, node, self.root_node) + + def _inc_hit_count_and_check( + self, node: UnifiedTreeNode, chunked: bool = False + ) -> bool: + """Increment hit count; check whether a write backup should be fired.""" + if node.evicted or chunked: + return False + if self.is_write_back: + return False + node.hit_count += 1 + return ( + self.enable_hicache + and not node.backuped + and node.hit_count >= self.write_through_threshold + ) + + def begin_insert(self, params: InsertParams) -> InsertStepResult: + """Start the insert, running to its first barrier or completion.""" + # Insert walks are single-flight; a live walk means re-entrancy. + assert self._ongoing_insert_walk_state is None, "concurrent insert walks" + key = params.key + value = params.value + key, value = key.maybe_to_bigram_view(self.is_eagle, value) + key = key.page_aligned(self.page_size) + if value is not None: + value = value[: len(key)] + else: + value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) + + priority = params.priority + if priority is None: + priority = 0 + self._touch_node(self.root_node) + self.root_node.priority = max(self.root_node.priority, priority) + if len(key) == 0: + return InsertStepResult( + actions=[], result=InsertResult(prefix_len=0, mamba_exist=True) + ) + + self._ongoing_insert_walk_state = _InsertWalkState( + phase=_InsertPhase.WALK, + node=self.root_node, + key=key, + value=value, + params=params, + priority=priority, + ) + return self._advance_insert() + + def resume_insert(self) -> InsertStepResult: + """Continue the suspended insert after its step actions were executed.""" + assert self._ongoing_insert_walk_state is not None, "no in-flight insert" + return self._advance_insert() + + def has_ongoing_insert(self) -> bool: + """Whether an insert walk is suspended at a barrier.""" + return self._ongoing_insert_walk_state is not None + + def end_insert(self) -> list[CacheAction | ComponentAction]: + """Finish the insert (idempotent); returns still-pending actions to drain.""" + state = self._ongoing_insert_walk_state + self._ongoing_insert_walk_state = None + return state.pending_actions if state is not None else [] + + def _advance_insert(self) -> InsertStepResult: + """Run the in-flight insert to its next barrier or to completion.""" + state = self._ongoing_insert_walk_state + while True: + flushed_len = len(state.pending_actions) + if state.phase is _InsertPhase.WALK: + self._insert_walk_step(state) + elif state.phase is _InsertPhase.COMMIT: + self._insert_commit_step(state) + elif state.phase is _InsertPhase.TAIL: + self._insert_tail_step(state) + self._ongoing_insert_walk_state = None + return InsertStepResult( + actions=state.pending_actions, result=state.result + ) + else: + raise AssertionError(f"unsupported insert phase: {state.phase}") + new_actions = state.pending_actions[flushed_len:] + # Suspend only when a step emitted a non-deferrable action. + if new_actions and not all(map(self._is_deferrable_action, new_actions)): + flushed = state.pending_actions + state.pending_actions = [] + return InsertStepResult(actions=flushed) + + @staticmethod + def _is_deferrable_action(action: CacheAction | ComponentAction) -> bool: + """Fire-and-forget actions safe to batch until the next barrier.""" + return isinstance(action, (FreeDeviceKV, ReplaceWriteThroughOnNodeSplit)) + + def _insert_walk_step(self, state: _InsertWalkState) -> None: + """Process one walked node, appending its barrier actions to the state.""" + key = state.key + child_key = key.child_key(self.page_size) if len(key) else None + if child_key not in state.node.children: + state.phase = _InsertPhase.COMMIT + return + step_actions = state.pending_actions + node = state.node.children[child_key] + self._touch_node(node) + prefix_len = node.key.match(key, page_size=self.page_size) + if prefix_len < len(node.key): + node, action = self._split_node(node.key, node, prefix_len) + if action is not None: + step_actions.append(action) + node.priority = max(node.priority, state.priority) + + if node.evicted: + self._unevict_node_on_insert(node, state.value[:prefix_len]) + # FULL was restored from the request's fresh KV. Aux + # components (e.g. SWA) may still hold tombstones and need + # to rebuild their value from the same slice. + for component in self.components: + if component.component_type == BASE_COMPONENT_TYPE: + continue + component.recover_after_unevict( + node=node, + prefix_len=prefix_len, + total_prefix_len=state.total_prefix_length, + params=state.params, + cache_actions=step_actions, + ) + else: + value_slice = state.value[:prefix_len] + consumed_from = prefix_len + # Let each component claim ownership of overlapping KV slots + for component in self.components: + comp_consumed_from = component.update_component_on_insert_overlap( + node=node, + prefix_len=prefix_len, + total_prefix_len=state.total_prefix_length, + value_slice=value_slice, + params=state.params, + cache_actions=step_actions, + ) + consumed_from = min(consumed_from, comp_consumed_from) + + dup_start = max(0, state.params.prev_prefix_len - state.total_prefix_length) + if dup_start < consumed_from: + step_actions.append( + FreeDeviceKV([value_slice[dup_start:consumed_from]]) + ) + + if self._inc_hit_count_and_check(node, state.params.chunked): + step_actions.append(self._build_backup_kv_action(node)) + state.node = node + state.total_prefix_length += prefix_len + state.key = key[prefix_len:] + state.value = state.value[prefix_len:] + + def _insert_commit_step(self, state: _InsertWalkState) -> None: + """Create the tail leaf and run the component commit hooks.""" + # Create new leaf for remaining suffix. A leaf survives on its Full + # value alone; auxiliary components (SWA, Mamba) may legitimately hold + # only a tombstone for this span (e.g. the whole leaf is outside the SWA + # window). Materialize it anyway so the Full KV stays cacheable. + if len(state.key): + state.target_node = self._add_new_node( + state.node, state.key, state.value, priority=state.priority + ) + state.is_new_leaf = True + else: + state.target_node = state.node + + # Finalize: let each component attach its data to the target node. + # e.g. Mamba attaches mamba_value to the leaf node + # All hooks run before their emitted actions execute; an action failure + # fail-stops the process, so partial-commit state is never observed. + state.result = InsertResult(prefix_len=state.total_prefix_length) + for component in self.components: + component.commit_insert_component_data( + node=state.target_node, + is_new_leaf=state.is_new_leaf, + params=state.params, + result=state.result, + cache_actions=state.pending_actions, + ) + state.phase = _InsertPhase.TAIL + + def _insert_tail_step(self, state: _InsertWalkState) -> None: + """Refresh the LRUs and append the terminal new-leaf backup.""" + if state.target_node is not self.root_node: + for component in self.components: + if component.component_type == BASE_COMPONENT_TYPE: + continue + component.refresh_lru( + LRURefreshPhase.INSERT_END, state.target_node, self.root_node + ) + + if state.is_new_leaf and self._inc_hit_count_and_check( + state.target_node, state.params.chunked + ): + state.pending_actions.append( + self._build_backup_kv_action(state.target_node) + ) + + def _split_node( + self, key: RadixKey, child: UnifiedTreeNode, split_len: int + ) -> tuple[UnifiedTreeNode, Optional[CacheAction | ComponentAction]]: + new_node = self._new_node(priority=child.priority) + new_node.children = {key[split_len:].child_key(self.page_size): child} + new_node.parent = child.parent + new_node.key = child.key[:split_len] + new_node.hit_count = child.hit_count + new_node.creation_time = child.creation_time + + self._for_each_component_lru(child, UnifiedLRUList.remove_node) + + child.parent = new_node + child.key = child.key[split_len:] + new_node.hash_value, child.hash_value = split_node_hash_value( + child.hash_value, split_len, self.page_size + ) + + for component in self.components: + component.redistribute_on_node_split(new_parent=new_node, child=child) + new_node.parent.children[key.child_key(self.page_size)] = new_node + + # A split of a backuped node tells the cache to fix its publish list. + action: Optional[CacheAction | ComponentAction] = None + if child.write_through_pending_id is not None: + ack_id = child.write_through_pending_id + new_node.write_through_pending_id = ack_id + action = ReplaceWriteThroughOnNodeSplit( + ack_id=ack_id, + old_node_id=child.id, + new_node_id=new_node.id, + new_child_node_id=child.id, + ) + + self._for_each_component_lru( + new_node, UnifiedLRUList.insert_mru, skip_existing=True + ) + self._for_each_component_lru( + child, UnifiedLRUList.insert_mru, skip_existing=True + ) + child.last_access_time = get_and_increase_time_counter() + + self._update_evictable_leaf_sets(new_node) + self._update_evictable_leaf_sets(child) + return new_node, action + + def _add_new_node( + self, + parent: UnifiedTreeNode, + key: RadixKey, + value: torch.Tensor, + priority: int = 0, + ) -> UnifiedTreeNode: + new_node = self._new_node(priority=priority) + new_node.parent = parent + new_node.key = key + new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone() + parent.children[key.child_key(self.page_size)] = new_node + self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value) + if self.enable_storage: + new_node.hash_value = compute_node_hash_values(new_node, self.page_size) + + self._update_evictable_leaf_sets(new_node) + self._update_evictable_leaf_sets(parent) + self._record_store_event(new_node) + return new_node + + def _unevict_node_on_insert( + self, node: UnifiedTreeNode, fresh_value: torch.Tensor + ) -> None: + """Restore an evicted node's Full device value from fresh KV indices + during insert.""" + ct = BASE_COMPONENT_TYPE + cd = node.component_data[ct] + assert cd.value is None + n = len(fresh_value) + cd.value = fresh_value.clone() + self.component_evictable_size_[ct] += n + self._update_evictable_leaf_sets(node) + if node.parent is not None: + self._update_evictable_leaf_sets(node.parent) + self._record_store_event(node, medium=StorageMedium.GPU) + + def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: + """Update both device and host leaf sets for a node.""" + if self._is_device_leaf(node): + self.evictable_device_leaves.add(node) + else: + self.evictable_device_leaves.discard(node) + + if self._is_host_leaf(node): + self.evictable_host_leaves.add(node) + else: + self.evictable_host_leaves.discard(node) + + def _for_each_component_lru( + self, + node: UnifiedTreeNode, + lru_op, + target: EvictLayer = EvictLayer.DEVICE, + skip_existing: bool = False, + ): + """Apply lru_op to each aux component's LRU that has data on this node. + If skip_existing=True, skip components already in the target LRU list.""" + lru_dict = self.host_lru_lists if target is EvictLayer.HOST else self.lru_lists + for ct in self.component_types: + if ct == BASE_COMPONENT_TYPE: + continue # Full uses leaf sets, not LRU + cd = node.component_data[ct] + if (cd.host_value if target is EvictLayer.HOST else cd.value) is not None: + lru = lru_dict[ct] + if skip_existing and lru.in_list(node): + continue + lru_op(lru, node) + + def evict_device_start( + self, component_type: ComponentType, request_cnt: int + ) -> None: + """Begin a component's device-eviction walk for up to request_cnt tokens.""" + self.components_by_type[component_type].evict_device_start(request_cnt) + + def evict_device_next_node( + self, component_type: ComponentType, tracker: dict[ComponentType, int] + ) -> EvictDeviceNextNodeResult: + """Return the next device leaf to evict for a component, or None when done.""" + result = EvictDeviceNextNodeResult() + # The walk reads running totals for its doneness check; the result + # carries only this step's delta. + updated_tracker = defaultdict(int, tracker) + result.node_id = self.components_by_type[component_type].evict_device_next_node( + updated_tracker, result.device_frees, result.host_frees + ) + for ct, n in updated_tracker.items(): + delta = n - tracker.get(ct, 0) + if delta: + result.tracker[ct] = delta + return result + + def evict_device_end(self, component_type: ComponentType) -> None: + """Finish a component's device-eviction walk.""" + self.components_by_type[component_type].evict_device_end() + + def evict_device_leaf( + self, node_id: NodeId, is_write_back: bool + ) -> EvictDeviceLeafResult: + """Evict one device leaf (demote if backuped, delete if write-through); + for an unbacked write-back node, the result carries the BackupKV for + the cache to execute and then demote.""" + result = EvictDeviceLeafResult() + node = self.node_by_id(node_id) + assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf" + if not node.backuped: + if is_write_back: + result.backup_kv = self._build_backup_kv_action(node, write_back=True) + return result + # Write-through: node has no backup, delete entirely. + self._delete_unbacked_device_leaf( + node, + result.tracker, + device_frees=result.device_frees, + host_frees=result.host_frees, + ) + return result + self._demote( + node, + result.tracker, + device_frees=result.device_frees, + host_frees=result.host_frees, + ) + return result + + def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult: + """Write-back fallback when a D-leaf's D->H backup fails under host + memory pressure: drop the subtree rooted at the unbacked leaf so + device eviction keeps making progress instead of leaving its KV + unevictable until host space frees up.""" + result = DropSubtreeNoHostResult(is_dropped=False) + node = self.node_by_id(node_id) + assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf" + # A failed backup never issues the D->H copy, so the subtree root has + # no host state and no in-flight DMA reading its device slots. + assert not node.backuped and node.write_through_pending_id is None + if any(cd.host_lock_ref > 0 for cd in node.component_data): + return result + descendants: list[UnifiedTreeNode] = [] + stack = list(node.children.values()) + while stack: + cur = stack.pop() + if any( + cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data + ): + return result + descendants.append(cur) + stack.extend(cur.children.values()) + for desc in reversed(descendants): + # Host-only by construction: a device descendant would contradict + # this node being a D-leaf, and D-leaves evict before ancestors. + assert desc.evicted and desc.backuped, f"node {desc.id} not host-only" + assert desc.write_through_pending_id is None + self._release_all_component_layers( + desc, + StorageMedium.CPU, + result.tracker, + result.device_frees, + result.host_frees, + ) + self._remove_leaf_from_parent(desc) + self._delete_unbacked_device_leaf( + node, + result.tracker, + device_frees=result.device_frees, + host_frees=result.host_frees, + ) + result.is_dropped = True + return result + + def _release_all_component_layers( + self, + node: UnifiedTreeNode, + medium: StorageMedium, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Free every component layer on the node and detach it from the LRU + lists and evictable leaf sets.""" + self._record_remove_event(node, medium=medium) + for comp in self.components: + self._evict_component_and_detach_lru( + node, + comp, + target=EvictLayer.ALL, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + self.evictable_device_leaves.discard(node) + self.evictable_host_leaves.discard(node) + + def _delete_unbacked_device_leaf( + self, + node: UnifiedTreeNode, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Delete a device leaf that has no host backup, freeing all layers.""" + self._release_all_component_layers( + node, StorageMedium.GPU, tracker, device_frees, host_frees + ) + parent = node.parent + self._remove_leaf_from_parent(node) + self._update_evictable_leaf_sets(parent) + self._iteratively_delete_tombstone_leaf( + node, tracker, device_frees=device_frees, host_frees=host_frees + ) + + def drive_host_eviction( + self, component_type: ComponentType, num_tokens: int + ) -> DriveHostEvictionResult: + """Evict a component's host-side resources; no-op if the component is absent.""" + result = DriveHostEvictionResult() + comp = self.components_by_type.get(component_type) + if comp is not None: + comp.drive_host_eviction( + num_tokens, + result.tracker, + result.device_frees, + result.host_frees, + ) + return result + + def _evict_host_leaf( + self, + node: UnifiedTreeNode, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Atomically evict all components on a host leaf. + + All freed tokens are accumulated into *tracker*.""" + assert self._is_host_leaf(node), f"node {node.id} is not an H-leaf" + + self._record_remove_event(node, medium=StorageMedium.CPU) + for comp in self.components: + _, hf = self._evict_component_and_detach_lru( + node, + comp, + target=EvictLayer.ALL, + tracker=None, + device_frees=device_frees, + host_frees=host_frees, + ) + tracker[comp.component_type] += hf + self.evictable_host_leaves.discard(node) + self._remove_leaf_from_parent(node) + self._iteratively_delete_tombstone_leaf(node, tracker, device_frees, host_frees) + + def demote(self, node_id: NodeId) -> DemoteResult: + """Release a node's device KV once its host copy exists; the node stays in the + tree, now host-only.""" + result = DemoteResult() + self._demote( + self.node_by_id(node_id), + result.tracker, + result.device_frees, + result.host_frees, + ) + return result + + def _demote( + self, + node: UnifiedTreeNode, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + assert not node.evicted and node.backuped + trigger = self.components_by_type[BASE_COMPONENT_TYPE] + self._evict_component_and_detach_lru( + node, + trigger, + target=EvictLayer.DEVICE, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + self._cascade_evict( + node, trigger, tracker, device_frees=device_frees, host_frees=host_frees + ) + self._record_remove_event(node, medium=StorageMedium.GPU) + + # after device eviction, insert aux components into host LRU. + self._for_each_component_lru( + node, UnifiedLRUList.insert_mru, target=EvictLayer.HOST, skip_existing=True + ) + self._update_evictable_leaf_sets(node.parent) + + def _cascade_evict( + self, + node: UnifiedTreeNode, + trigger: TreeComponent, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + target: EvictLayer = EvictLayer.DEVICE, + ): + """Cascade eviction from trigger to lower-or-equal priority components.""" + + is_leaf = False + if target == EvictLayer.DEVICE: + is_leaf = node in self.evictable_device_leaves + elif target == EvictLayer.HOST: + is_leaf = node in self.evictable_host_leaves + + trigger_priority = trigger.eviction_priority(is_leaf) + + for comp in self.components: + if comp.eviction_priority(is_leaf) <= trigger_priority: + if comp is not trigger and comp.node_has_component_data(node, target): + cd = node.component_data[comp.component_type] + # A comp whose TRUE internal priority outranks the trigger + # is only in this loop because leaf-collapse flattened + # priorities; a lock on it is a legit pin and must be + # spared. A lock on a strictly-lower-priority tier is a + # real strand — fall through to the assert below. + if comp.eviction_priority( + is_leaf=False + ) >= trigger.eviction_priority(is_leaf=False): + if EvictLayer.DEVICE in target and cd.lock_ref != 0: + continue + if EvictLayer.HOST in target and cd.host_lock_ref != 0: + continue + if EvictLayer.DEVICE in target: + assert cd.lock_ref == 0 + if EvictLayer.HOST in target: + assert cd.host_lock_ref == 0 + self._evict_component_and_detach_lru( + node, + comp, + target=target, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + + # Now that all components (including SWA which depends on Full.value) + # have been freed, we can safely tombstone Full.value. + # This is deferred from evict_component because free_swa needs it. + if ( + target is EvictLayer.DEVICE + and trigger.component_type == BASE_COMPONENT_TYPE + ): + node.component_data[trigger.component_type].value = None + + self._update_evictable_leaf_sets(node) + + def _remove_leaf_from_parent(self, node: UnifiedTreeNode): + key = node.key.child_key(self.page_size) + v = node.parent.children.pop(key, None) + assert v == node + self._unregister_node(node) + + def _evict_component_and_detach_lru( + self, + node: UnifiedTreeNode, + comp: TreeComponent, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + target: EvictLayer = EvictLayer.DEVICE, + tracker: Optional[dict[ComponentType, int]] = None, + ) -> tuple[int, int]: + device_freed, host_freed = comp.evict_component( + node, target=target, device_frees=device_frees, host_frees=host_frees + ) + if tracker is not None: + if EvictLayer.DEVICE in target: + tracker[comp.component_type] += device_freed + elif EvictLayer.HOST in target: + tracker[comp.component_type] += host_freed + + # Detach from the appropriate LRU list(s) + ct = comp.component_type + for layer, lru_lists in ( + (EvictLayer.DEVICE, self.lru_lists), + (EvictLayer.HOST, self.host_lru_lists), + ): + if layer in target: + lru = lru_lists[ct] + if lru.in_list(node): + lru.remove_node(node) + return device_freed, host_freed + + def _iteratively_delete_tombstone_leaf( + self, + deleted_node: UnifiedTreeNode, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ): + """Walk up from *deleted_node* and cascade-delete childless ancestors. + + Only the Full (base) component decides whether a node survives: + - Full device present → keep as D-leaf + - Full host present → keep as H-leaf + - neither → evict all remaining data, delete, continue up + """ + ct = BASE_COMPONENT_TYPE + cur = deleted_node.parent + while cur != self.root_node and len(cur.children) == 0: + if any( + cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data + ): + break + + has_device = cur.component_data[ct].value is not None + has_host = cur.component_data[ct].host_value is not None + + if has_device: + self._update_evictable_leaf_sets(cur) + break + + # Full device absent — clean up orphaned aux device data. + for comp in self.components_by_type.values(): + if comp.node_has_component_data(cur): + self._evict_component_and_detach_lru( + cur, + comp, + target=EvictLayer.DEVICE, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + + if has_host: + self._update_evictable_leaf_sets(cur) + break + + # Full absent on both layers — evict remaining host data, delete. + for comp in self.components_by_type.values(): + if comp.node_has_component_data(cur, target=EvictLayer.HOST): + self._evict_component_and_detach_lru( + cur, + comp, + target=EvictLayer.HOST, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + + self.evictable_host_leaves.discard(cur) + self._remove_leaf_from_parent(cur) + parent = cur.parent + self._update_evictable_leaf_sets(parent) + cur = parent + + def _is_device_leaf(self, node: UnifiedTreeNode) -> bool: + """D-leaf: Full device value present, no child with Full KV on device, + unlocked, not root. + + Only the Full (base) component is required; auxiliary components + (Mamba, SWA) are not mandatory for D-leaf membership.""" + ct = BASE_COMPONENT_TYPE + if node is self.root_node or node.evicted: + return False + if any(cd.lock_ref > 0 for cd in node.component_data): + return False + if any( + child.component_data[ct].value is not None + for child in node.children.values() + ): + return False + return True + + def _is_host_leaf(self, node: UnifiedTreeNode) -> bool: + """H-leaf: evicted, Full host value present, no children, unlocked, not root. + + Only the Full (base) component host_value is required; auxiliary + components are not mandatory for H-leaf membership.""" + if node is self.root_node or not node.evicted: + return False + if not node.backuped: + return False + if any(cd.host_lock_ref > 0 for cd in node.component_data): + return False + if len(node.children) > 0: + return False + return True + + # ==== HiCache ==== + + def set_hicache_enabled(self) -> None: + self.enable_hicache = True + + def insert_host( + self, + node_id: NodeId, + key: RadixKey, + host_value: torch.Tensor, + hash_value: list[str], + ) -> InsertResult: + """Insert a host-side (backuped) tree path descending from the given node.""" + node = self.node_by_id(node_id) + total_len = len(key) + self._touch_node(node) + if total_len == 0: + return InsertResult(prefix_len=0, mamba_exist=True) + + child_key = key.child_key(self.page_size) + matched_length = 0 + cache_actions: list[CacheAction | ComponentAction] = [] + while len(key) > 0 and child_key in node.children: + node = node.children[child_key] + self._touch_node(node) + prefix_len = node.key.match(key, page_size=self.page_size) + + key = key[prefix_len:] + host_value = host_value[prefix_len:] + hash_value = hash_value[prefix_len // self.page_size :] + matched_length += prefix_len + + if prefix_len < len(node.key): + node, action = self._split_node(node.key, node, prefix_len) + if action is not None: + cache_actions.append(action) + + if len(key): + child_key = key.child_key(self.page_size) + + result = InsertResult( + prefix_len=matched_length, + total_len=total_len, + cache_actions=cache_actions, + ) + if len(key) == 0: + if ( + node is not self.root_node + and node.component_data[BASE_COMPONENT_TYPE].host_value is not None + ): + result.inserted_host_node = node.id + return result + + new_node = self._new_node(priority=node.priority) + new_node.parent = node + new_node.key = key + new_node.hash_value = hash_value + new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone() + node.children[child_key] = new_node + self._update_evictable_leaf_sets(new_node) + self._update_evictable_leaf_sets(node) + result.inserted_host_node = new_node.id + return result + + def build_backup_spec(self, node_id: NodeId): + """Read a node's device->host backup spec (device value + component transfers) now.""" + return self._build_backup_spec(self.node_by_id(node_id)) + + def _build_backup_spec(self, node: UnifiedTreeNode): + """Gather device value backup spec.""" + device_value = node.component_data[BASE_COMPONENT_TYPE].value + comp_xfers: dict[ComponentType, list] = {} + for comp in self.components: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST) + if t: + comp_xfers[comp.component_type] = t + return device_value, comp_xfers + + def build_storage_backup_spec( + self, node_id: NodeId, pass_prefix_keys: bool + ) -> Optional[StorageBackupSpec]: + """Gather a node's device->storage backup spec; None if the node is not backuped.""" + node = self.node_by_id(node_id) + if not node.backuped: + return None + prefix_keys = None + if pass_prefix_keys: + prefix_keys = node.get_prefix_hash_values(node.parent) + comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} + for comp in self.components: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + transfers = comp.build_hicache_transfers( + node, CacheTransferPhase.BACKUP_STORAGE + ) + if transfers: + comp_xfers[comp.component_type] = transfers + return StorageBackupSpec( + host_value=node.component_data[BASE_COMPONENT_TYPE].host_value, + token_ids=node.key.token_ids, + hash_value=node.hash_value, + prefix_keys=prefix_keys, + comp_xfers=comp_xfers, + ) + + def build_hicache_transfers( + self, + component_type: ComponentType, + node_id: NodeId, + phase: CacheTransferPhase, + *, + host_indices: Optional[torch.Tensor] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, + ) -> Optional[list[PoolTransfer]]: + """Route a build_hicache_transfers call to the component for the given type.""" + return self.components_by_type[component_type].build_hicache_transfers( + self.node_by_id(node_id), + phase, + host_indices=host_indices, + token_ids=token_ids, + prefetch_tokens=prefetch_tokens, + last_hash=last_hash, + ) + + def build_load_back_spec( + self, node_id: NodeId, req: Optional[Req] = None + ) -> tuple[PoolTransfer, dict[ComponentType, list[PoolTransfer]]]: + """Build the H->D load-back KV transfer plus per-component aux transfers.""" + # Component hooks take primitives, not Req: extract its fields here. + mamba_pool_idx = req.mamba_pool_idx if req is not None else None + node = self.node_by_id(node_id) + kv_xfer = self.components_by_type[BASE_COMPONENT_TYPE].build_hicache_transfers( + node, CacheTransferPhase.LOAD_BACK + )[0] + comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} + for comp in self.components: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + t = comp.build_hicache_transfers( + node, CacheTransferPhase.LOAD_BACK, mamba_pool_idx=mamba_pool_idx + ) + if t: + comp_xfers[comp.component_type] = t + return kv_xfer, comp_xfers + + def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]: + """The anchor node's key extra_key.""" + node = self.node_by_id(node_id) + return node.key.extra_key if node.key else None + + def _build_backup_kv_action( + self, node: UnifiedTreeNode, write_back: bool = False + ) -> BackupKV: + """Build the backup action for a node and its unbacked ancestors.""" + chain = [node] + if not write_back: + ancestor = node.parent + while ( + ancestor is not None + and ancestor is not self.root_node + and not ancestor.backuped + ): + chain.append(ancestor) + ancestor = ancestor.parent + # write_through: Ancestors first to preserve backup invariant + chain.reverse() + return BackupKV([target.id for target in chain]) + + def commit_hicache_transfers( + self, + node_id: NodeId, + phase: CacheTransferPhase, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + *, + cache_actions: list[CacheAction | ComponentAction], + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, + ) -> None: + """Commit each component's HiCache transfers onto the node.""" + node = self.node_by_id(node_id) + for ct, xfers in comp_xfers.items(): + self.components_by_type[ct].commit_hicache_transfer( + node, + phase, + xfers, + cache_actions=cache_actions, + insert_result=insert_result, + pool_storage_result=pool_storage_result, + ) + + def commit_backup( + self, + node_id: NodeId, + host_indices: torch.Tensor, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> None: + """Commit a successful backup to the node.""" + node = self.node_by_id(node_id) + cache_actions: list[CacheAction | ComponentAction] = [] + kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) + self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer( + node, + CacheTransferPhase.BACKUP_HOST, + transfers=[kv_xfer], + cache_actions=cache_actions, + ) + for ct, xfers in comp_xfers.items(): + self.components_by_type[ct].commit_hicache_transfer( + node, + CacheTransferPhase.BACKUP_HOST, + transfers=xfers, + cache_actions=cache_actions, + ) + assert not cache_actions # BACKUP_HOST emits no actions + + def commit_load_back( + self, + node_id: NodeId, + device_indices: torch.Tensor, + kv_xfer: PoolTransfer, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> list[CacheAction | ComponentAction]: + """Commit a successful H->D load-back onto the node; the SWA full->swa mapping + rebuild is deferred to the orchestration layer.""" + node = self.node_by_id(node_id) + cache_actions: list[CacheAction | ComponentAction] = [] + kv_xfer.device_indices = device_indices + self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer( + node, + CacheTransferPhase.LOAD_BACK, + [kv_xfer], + cache_actions=cache_actions, + ) + for nid in kv_xfer.nodes_to_load or (): + loaded = self.node_by_id(nid) + self._record_store_event(loaded, medium=StorageMedium.GPU) + for ct, xfers in comp_xfers.items(): + self.components_by_type[ct].commit_hicache_transfer( + node, + CacheTransferPhase.LOAD_BACK, + xfers, + cache_actions=cache_actions, + ) + self._update_evictable_leaf_sets(node) + return cache_actions + + def mark_write_through_pending(self, node_id: NodeId) -> None: + """Mark a node as having an in-flight write-through backup.""" + node = self.node_by_id(node_id) + node.write_through_pending_id = node_id + + def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None: + """Clear the write-through-pending mark (when it matches ack_id) and record the + host store event for each acked node.""" + for node_id in node_ids: + node = self.node_by_id(node_id) + if node.write_through_pending_id == ack_id: + node.write_through_pending_id = None + self._record_store_event(node, medium=StorageMedium.CPU) + + def set_component_device_value( + self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor + ) -> None: + """Store an auxiliary component's device value onto a node.""" + # Full uses leaf sets, not LRU; its stores go through the insert paths. + assert component_type != BASE_COMPONENT_TYPE + node = self.node_by_id(node_id) + node.component_data[component_type].value = value + host_lru = self.host_lru_lists[component_type] + if host_lru.in_list(node): + host_lru.remove_node(node) + self.lru_lists[component_type].insert_mru(node) + self.component_evictable_size_[component_type] += len(value) + + def get_component_device_value( + self, node_id: NodeId, component_type: ComponentType + ) -> Optional[torch.Tensor]: + """The component's device value on the node, or None if evicted.""" + return self.node_by_id(node_id).component_data[component_type].value + + def component_has_host_value_only( + self, node_id: NodeId, component_type: ComponentType + ) -> bool: + """Whether the component's data is device-evicted but host-backed.""" + cd = self.node_by_id(node_id).component_data[component_type] + return cd.value is None and cd.host_value is not None + + # ==== Others ==== + + def sanity_check( + self, + ongoing_write_through: list[tuple[int, NodeId]], + ongoing_load_back: list[tuple[int, NodeId]], + ) -> None: + """Verify tree-structure, leaf-set, LRU, size, and ongoing-op invariants; raise + AssertionError on any violation. ongoing_* args are (id, node_id) pairs. + """ + errors: list[str] = [] + E = errors.append + all_nodes = self._collect_all_nodes() + all_node_set = set(all_nodes) + FCT = BASE_COMPONENT_TYPE + + # ── PART 1: Tree Structure ── + # Root state + if self.root_node.component_data[FCT].value is None: + E("[Root] root missing Full device value") + if self.root_node.component_data[FCT].lock_ref <= 0: + E( + f"[Root] root Full lock_ref={self.root_node.component_data[FCT].lock_ref}" + ) + if self.root_node.parent is not None: + E("[Root] root has a parent pointer") + # Parent ↔ child bidirectional consistency + for node in all_nodes: + for child in node.children.values(): + if child.parent is not node: + pid = child.parent.id if child.parent else None + E(f"[Tree] child {child.id} parent={pid}, expected {node.id}") + if child.key is None: + E(f"[Tree] node {child.id} has no key") + + # ── PART 2: Per-node state machine and leaf qualification ── + expected_dev_leaves: set[UnifiedTreeNode] = set() + expected_hst_leaves: set[UnifiedTreeNode] = set() + + for node in all_nodes: + if node is self.root_node: + continue + nid = node.id + full_dev = node.component_data[FCT].value is not None + full_hst = node.component_data[FCT].host_value is not None + + # Full is the tree backbone, so aux data requires Full data. + for ct in self.component_types: + if ct == FCT: + continue + cd = node.component_data[ct] + if cd.value is not None and not full_dev: + E(f"node {nid} {ct} device present but Full.value=None") + if cd.host_value is not None and not full_hst: + E(f"node {nid} {ct} host present but Full.host_value=None") + + # Every node must keep Full data on at least one layer. + if not full_dev and not full_hst: + E(f"node {nid} dead: no Full device and no Full host") + + # Parent prefixes must keep data whenever the child does. + if node.parent is not None and node.parent is not self.root_node: + p_dev = node.parent.component_data[FCT].value is not None + p_hst = node.parent.component_data[FCT].host_value is not None + if full_dev and not p_dev: + E(f"node {nid} device present but parent {node.parent.id} evicted") + if full_hst and not p_hst and not self.is_write_back: + E(f"node {nid} backed up but parent {node.parent.id} not backed up") + + # Lock hierarchy and counters must stay sane. + fl = node.component_data[FCT].lock_ref + for ct in self.component_types: + cd = node.component_data[ct] + if cd.lock_ref < 0: + E(f"node {nid} {ct} lock_ref={cd.lock_ref}") + if cd.host_lock_ref < 0: + E(f"node {nid} {ct} host_lock_ref={cd.host_lock_ref}") + if ct != FCT and fl < cd.lock_ref: + E(f"node {nid} full_lock={fl} < {ct}_lock={cd.lock_ref}") + if cd.value is None and cd.lock_ref > 0: + E(f"node {nid} {ct} evicted but lock_ref={cd.lock_ref}") + + # Collect expected leaf qualification (single pass) + if self._is_device_leaf(node): + expected_dev_leaves.add(node) + if self._is_host_leaf(node): + expected_hst_leaves.add(node) + + # ── PART 3: Tracking structures ── + + # Device leaf set must match the expected leaves. + if self.evictable_device_leaves != expected_dev_leaves: + extra = self.evictable_device_leaves - expected_dev_leaves + missing = expected_dev_leaves - self.evictable_device_leaves + if extra: + E(f"D-leaf extra: {[n.id for n in list(extra)[:5]]}") + if missing: + E(f"D-leaf missing: {[n.id for n in list(missing)[:5]]}") + + # Host leaf set must match the expected leaves. + if self.evictable_host_leaves != expected_hst_leaves: + extra = self.evictable_host_leaves - expected_hst_leaves + missing = expected_hst_leaves - self.evictable_host_leaves + if extra: + E(f"H-leaf extra: {[n.id for n in list(extra)[:5]]}") + if missing: + E(f"H-leaf missing: {[n.id for n in list(missing)[:5]]}") + + # D-leaf ∩ H-leaf = ∅ + overlap = self.evictable_device_leaves & self.evictable_host_leaves + if overlap: + E( + f"[Leaf] {len(overlap)} in both sets: {[n.id for n in list(overlap)[:5]]}" + ) + + # Stale nodes: leaf sets must only contain tree-reachable nodes + stale = self.evictable_device_leaves - all_node_set + if stale: + E( + f"{len(stale)} stale nodes in device_leaves: {[n.id for n in list(stale)[:5]]}" + ) + stale = self.evictable_host_leaves - all_node_set + if stale: + E( + f"{len(stale)} stale nodes in host_leaves: {[n.id for n in list(stale)[:5]]}" + ) + + # Per-component LRU tracking + for ct in self.component_types: + lru = self.lru_lists[ct] + if ct == FCT: + # Full uses leaf sets, not LRU + if len(lru.cache) > 0: + E(f"Full device LRU not empty: {len(lru.cache)}") + if len(self.host_lru_lists[ct].cache) > 0: + E(f"Full host LRU not empty: {len(self.host_lru_lists[ct].cache)}") + else: + # Aux device values must match the device LRU. + tree_ids = { + n.id + for n in all_nodes + if n is not self.root_node + and n.component_data[ct].value is not None + } + lru_ids = set(lru.cache.keys()) + if tree_ids != lru_ids: + E( + f"{ct} device LRU: " + f"+tree={tree_ids - lru_ids}, +lru={lru_ids - tree_ids}" + ) + # Aux host-only states must match the host LRU. + host_lru = self.host_lru_lists[ct] + s3_ids = { + n.id + for n in all_nodes + if n is not self.root_node + and n.component_data[ct].value is None + and n.component_data[ct].host_value is not None + } + host_lru_ids = set(host_lru.cache.keys()) + if s3_ids != host_lru_ids: + E( + f"{ct} host LRU: " + f"+S3={s3_ids - host_lru_ids}, +lru={host_lru_ids - s3_ids}" + ) + # The same aux node must not appear in both device and host LRU. + inv5_overlap = lru_ids & host_lru_ids + if inv5_overlap: + E(f"{ct} in both device and host LRU: {inv5_overlap}") + # Linked-list integrity + self._check_lru_linked_list(lru, ct, "device", errors) + self._check_lru_linked_list(host_lru, ct, "host", errors) + + # ── PART 4: Size Accounting ── + for ct in self.component_types: + evictable = 0 + protected = 0 + for n in all_nodes: + if n is self.root_node: + continue + cd = n.component_data[ct] + if cd.value is not None: + toks = len(cd.value) + if cd.lock_ref > 0: + protected += toks + else: + evictable += toks + if self.component_evictable_size_[ct] != evictable: + E( + f"[Size] {ct} evictable={self.component_evictable_size_[ct]} " + f"!= recomputed={evictable}" + ) + if self.component_protected_size_[ct] != protected: + E( + f"[Size] {ct} protected={self.component_protected_size_[ct]} " + f"!= recomputed={protected}" + ) + + # ── PART 5: Ongoing Operations ── + for nid, node_id in ongoing_write_through: + n = self._node_arena.get(node_id) + if n is None or n not in all_node_set: + E(f"[Ongoing] write_through node {nid} not in tree") + elif n.component_data[FCT].lock_ref <= 0: + E( + f"[Ongoing] write_through node {nid} lock_ref={n.component_data[FCT].lock_ref}" + ) + for nid, node_id in ongoing_load_back: + n = self._node_arena.get(node_id) + if n is None or n not in all_node_set: + E(f"[Ongoing] load_back node {nid} not in tree") + elif n.component_data[FCT].lock_ref <= 0: + E( + f"[Ongoing] load_back node {nid} lock_ref={n.component_data[FCT].lock_ref}" + ) + + if errors: + msg = ( + f"Sanity check FAILED ({len(errors)} violations " + f"across {len(all_nodes)} nodes):\n" + + "\n".join(f" {e}" for e in errors) + ) + logger.error(msg) + self.pretty_print() + raise AssertionError(msg) + + def _collect_all_nodes(self) -> list[UnifiedTreeNode]: + nodes = [] + stack = [self.root_node] + while stack: + node = stack.pop() + nodes.append(node) + stack.extend(node.children.values()) + return nodes + + def _check_lru_linked_list( + self, + lru: UnifiedLRUList, + ct: ComponentType, + label: str, + errors: list[str], + ) -> None: + """Walk a LRU doubly-linked list, collect integrity errors.""" + pt = lru._pt # use LRU's own pointer slot + visited: set[int] = set() + x = lru.head.lru_next[pt] + prev = lru.head + while x is not None and x != lru.tail: + if x.lru_prev[pt] != prev: + errors.append(f"[{label}][{ct}] broken prev at node {x.id}") + if x.id not in lru.cache: + errors.append(f"[{label}][{ct}] node {x.id} in list not cache") + if x.id in visited: + errors.append(f"[{label}][{ct}] cycle at node {x.id}") + break + visited.add(x.id) + prev = x + x = x.lru_next[pt] + if x is None: + errors.append( + f"[{label}][{ct}] broken chain: lru_next is None " + f"after node {prev.id if hasattr(prev, 'id') else 'head'}" + ) + if len(visited) != len(lru.cache): + errors.append( + f"[{label}][{ct}] list={len(visited)} != cache={len(lru.cache)}" + ) + + def pretty_print(self) -> None: + stack = [(self.root_node, 0)] + while stack: + node, indent = stack.pop() + component_str = " ".join( + f"{ct}={'yes' if node.component_data[ct].value is not None else 'no'}" + for ct in self.component_types + ) + print( + " " * indent, + f"[{node.id}]", + len(node.key), + f"full_lock={node.component_data[BASE_COMPONENT_TYPE].lock_ref}", + component_str, + ) + for child in node.children.values(): + stack.append((child, indent + 2)) + + def evictable_size(self) -> int: + return self.component_evictable_size_.get(BASE_COMPONENT_TYPE, 0) + + def protected_size(self) -> int: + return self.component_protected_size_.get(BASE_COMPONENT_TYPE, 0) + + def component_evictable_size(self, component_type: ComponentType) -> int: + """Evictable token count for one component (0 if the component is absent).""" + return self.component_evictable_size_.get(component_type, 0) + + def full_evictable_size(self) -> int: + return self.evictable_size() + + def full_protected_size(self) -> int: + return self.protected_size() + + def swa_evictable_size(self) -> int: + return self.component_evictable_size_.get(ComponentType.SWA, 0) + + def mamba_evictable_size(self) -> int: + return self.component_evictable_size_.get(ComponentType.MAMBA, 0) + + def swa_protected_size(self) -> int: + return self.component_protected_size_.get(ComponentType.SWA, 0) + + def mamba_protected_size(self) -> int: + return self.component_protected_size_.get(ComponentType.MAMBA, 0) + + def total_size(self) -> tuple[int, int]: + total_size = 0 + total_aux_size = 0 + stack = [self.root_node] + while stack: + node = stack.pop() + full_value = node.component_data[BASE_COMPONENT_TYPE].value + if full_value is not None: + total_size += len(full_value) + for ct in self.component_types: + if ct == BASE_COMPONENT_TYPE: + continue + value = node.component_data[ct].value + if value is not None: + total_aux_size += len(value) + for child in node.children.values(): + stack.append(child) + return total_size, total_aux_size + + def walk_for_kv_canary( + self, unlocked_only: bool, swa_resident_only: bool + ) -> RadixCacheWalkResult: + """Flatten every FULL device slot into (slot, position, prev-slot) rows for the KV-canary sweep.""" + slots: list[int] = [] + positions: list[int] = [] + prev_slots: list[int] = [] + swa_filter = swa_resident_only and ComponentType.SWA in self.components_by_type + + def _dfs(node: UnifiedTreeNode, depth: int, parent_last_slot: int) -> None: + value = node.component_data[BASE_COMPONENT_TYPE].value + node_slots = value.tolist() if isinstance(value, torch.Tensor) else [] + + emit = node is not self.root_node + if unlocked_only: + # Unified SWA owns an independent component lock. A node can still + # hold Full KV for a running request while its SWA slots are unused. + lock_ct = ComponentType.SWA if swa_filter else BASE_COMPONENT_TYPE + emit = emit and node.component_data[lock_ct].lock_ref == 0 + if swa_filter: + emit = emit and node.component_data[ComponentType.SWA].value is not None + + # Skipped nodes still advance the chain/depth so descendants stay consistent. + chain_last_slot = parent_last_slot + for j, slot in enumerate(node_slots): + if emit: + slots.append(slot) + positions.append(depth + j) + prev_slots.append(parent_last_slot if j == 0 else node_slots[j - 1]) + chain_last_slot = slot + + # Device-evicted nodes hold no slots but still span their key length. + if node is self.root_node or node.key is None: + child_depth = depth + len(node_slots) + else: + child_depth = depth + len(node.key) + for child in node.children.values(): + _dfs(child, child_depth, chain_last_slot) + + _dfs(self.root_node, 0, -1) + return RadixCacheWalkResult( + slot_indices=torch.tensor(slots, dtype=torch.int64), + positions=torch.tensor(positions, dtype=torch.int64), + prev_slot_indices=torch.tensor(prev_slots, dtype=torch.int64), + ) + + def all_values_flatten(self) -> torch.Tensor: + values = [] + + def _dfs(node: UnifiedTreeNode): + for child in node.children.values(): + v = child.component_data[BASE_COMPONENT_TYPE].value + if v is not None: + values.append(v) + _dfs(child) + + _dfs(self.root_node) + if values: + return torch.cat(values) + return torch.tensor([], dtype=torch.int64, device=self.device) + + def _all_component_values_flatten( + self, component_type: ComponentType + ) -> torch.Tensor: + if component_type not in self.components_by_type: + return torch.tensor([], dtype=torch.int64, device=self.device) + + values = [] + + def _dfs(node: UnifiedTreeNode): + value = node.component_data[component_type].value + if value is not None: + values.append(value) + for child in node.children.values(): + _dfs(child) + + _dfs(self.root_node) + if values: + return torch.cat(values) + return torch.tensor([], dtype=torch.int64, device=self.device) + + def all_mamba_values_flatten(self) -> torch.Tensor: + return self._all_component_values_flatten(ComponentType.MAMBA) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py new file mode 100644 index 000000000..f8da32f70 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -0,0 +1,485 @@ +"""Interface of the TreeCore: the radix tree mechanism that owns the tree +structure, per-node values, the LRU(s), and bookkeeping. The Controller drives a +TreeCore exclusively through this surface, so an alternative implementation (e.g. +a Rust TreeCore) can satisfy it without subclassing the Python tree. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections import defaultdict +from typing import TYPE_CHECKING, Optional, Sequence + +import msgspec + +from sglang.srt.mem_cache.events import KVCacheEventMixin + +# Tree node id -- the node handle used outside the TreeCore. The concrete tree +# node is a TreeCore-internal type. +NodeId = int + + +class BaseEvictionResult(msgspec.Struct): + """Base of the eviction step results: the component-keyed device/host + values the step freed, for the Controller to drain right after the call + returns, plus `tracker` -- the step's per-component evicted counts for the + Controller to accumulate.""" + + device_frees: dict[ComponentType, list[torch.Tensor]] = msgspec.field( + default_factory=lambda: defaultdict(list) + ) + host_frees: dict[ComponentType, list[torch.Tensor]] = msgspec.field( + default_factory=lambda: defaultdict(list) + ) + tracker: dict[ComponentType, int] = msgspec.field( + default_factory=lambda: defaultdict(int) + ) + + def __del__(self) -> None: + # Drop tripwire: every returned value must be drained before disposal. + assert ( + not self.device_frees and not self.host_frees + ), "BaseEvictionResult dropped with undrained values" + + +class EvictDeviceNextNodeResult(BaseEvictionResult): + node_id: Optional[NodeId] = None + + +class EvictDeviceLeafResult(BaseEvictionResult): + backup_kv: Optional[BackupKV] = None + + +class DemoteResult(BaseEvictionResult): + pass + + +class DropSubtreeNoHostResult(BaseEvictionResult): + is_dropped: bool = False + + +class DriveHostEvictionResult(BaseEvictionResult): + pass + + +class DecSwaLockOnlyResult(BaseEvictionResult): + pass + + +class RadixCacheWalkResult(msgspec.Struct, frozen=True, kw_only=True): + """Flat (slot, position, prev-slot) rows emitted by the KV-canary walk.""" + + slot_indices: torch.Tensor + positions: torch.Tensor + prev_slot_indices: torch.Tensor + + +class InsertStepResult(msgspec.Struct, frozen=True): + """One step of a resumable insert: the Controller executes ``actions``, then + resumes while ``result`` is None; ``result`` is set on the final step.""" + + actions: list[CacheAction | ComponentAction] + result: Optional[InsertResult] = None + + +if TYPE_CHECKING: + import torch + + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + DecLockRefResult, + IncLockRefResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, + ) + from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult + from sglang.srt.mem_cache.radix_cache import RadixKey + from sglang.srt.mem_cache.unified_cache.cache_action import ( + BackupKV, + CacheAction, + ComponentAction, + ) + from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( + StorageBackupSpec, + UnifiedTreeNode, + ) + from sglang.srt.mem_cache.unified_cache_components import ( + CacheTransferPhase, + ComponentType, + ) + + +class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC): + """Methods the Controller invokes on the Tree Core. The Controller treats the + Tree Core as opaque behind this surface, which grows as tree operations + migrate onto the TreeCore. Inherits KVCacheEventMixin for the KV-event API + (take_events, _record_* recorders).""" + + # ==== Tree-owned state the Controller reads (or, via its facade setters, writes) ==== + page_size: int + is_eagle: bool + device: torch.device + enable_hicache: bool + enable_storage: bool + write_through_threshold: int + is_write_back: bool + has_swa_host_pool: bool + + # ==== Tree API ==== + + @abstractmethod + def reset(self) -> None: + """Drop the entire tree and reinitialize empty state.""" + ... + + @abstractmethod + def node_by_id(self, node_id: NodeId) -> UnifiedTreeNode: + """Resolve a NodeId -- the tree-node identity the Controller passes + across the boundary (e.g. MatchResult fields, lock-ref args) -- back to + its tree node. + + TODO(Jialin): Remove after the Unified Radix Cache split. + """ + ... + + @abstractmethod + def is_backuped(self, node_id: NodeId) -> bool: + """Whether the node's KV is already backed up to host.""" + ... + + @abstractmethod + def is_root(self, node_id: NodeId) -> bool: + """Whether the node is the tree root.""" + ... + + @abstractmethod + def get_last_hash_value(self, node_id: NodeId) -> Optional[str]: + """The node's last page hash, or None when it was never hashed.""" + ... + + @abstractmethod + def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: + """The hash chain of the node's ancestors, in root-to-parent order.""" + ... + + @abstractmethod + def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult: + """Bump the reference count on a node's component locks.""" + ... + + @abstractmethod + def dec_lock_ref( + self, + node_id: NodeId, + params: Optional[DecLockRefParams] = None, + skip_swa: bool = False, + ) -> DecLockRefResult: + """Decrease the reference count on a node's component locks.""" + ... + + @abstractmethod + def dec_swa_lock_only( + self, node_id: NodeId, swa_uuid_for_lock: Optional[int] + ) -> DecSwaLockOnlyResult: + """Decrease only the SWA (and lower-priority co-located) reference + counts; the result carries the freed slots.""" + ... + + # ==== Device eviction (driven step-wise by the Controller's evict()) ==== + + @abstractmethod + def evict_device_start( + self, component_type: ComponentType, request_cnt: int + ) -> None: + """Begin a device-eviction walk for one component.""" + ... + + @abstractmethod + def evict_device_next_node( + self, component_type: ComponentType, tracker: dict[ComponentType, int] + ) -> EvictDeviceNextNodeResult: + """The next evictable node (None node_id when the walk is exhausted); + tracker is the caller's running totals, read for the doneness check.""" + ... + + @abstractmethod + def evict_device_leaf( + self, node_id: NodeId, is_write_back: bool + ) -> EvictDeviceLeafResult: + """Evict a leaf's device value; the result carries a BackupKV when a + D->H backup must run (write_back) before the node can be demoted.""" + ... + + @abstractmethod + def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult: + """Drop an unbacked D-leaf's subtree when its write-back backup failed + under host pressure; declines (is_dropped=False) if any node is locked.""" + ... + + @abstractmethod + def demote(self, node_id: NodeId) -> DemoteResult: + """Demote a backed-up node: drop its device value after a successful backup.""" + ... + + @abstractmethod + def evict_device_end(self, component_type: ComponentType) -> None: + """Finish the component's device-eviction walk.""" + ... + + @abstractmethod + def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: + """Bump the reference count on a node's host-side component locks.""" + ... + + @abstractmethod + def dec_host_lock_ref( + self, node_id: NodeId, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: + """Decrease the reference count on a node's host-side component locks.""" + ... + + @abstractmethod + def evictable_size(self) -> int: ... + + @abstractmethod + def protected_size(self) -> int: ... + + @abstractmethod + def component_evictable_size(self, component_type: ComponentType) -> int: + """Evictable token count for one component (0 if the component is absent).""" + ... + + @abstractmethod + def full_evictable_size(self) -> int: ... + + @abstractmethod + def full_protected_size(self) -> int: ... + + @abstractmethod + def swa_evictable_size(self) -> int: ... + + @abstractmethod + def mamba_evictable_size(self) -> int: ... + + @abstractmethod + def swa_protected_size(self) -> int: ... + + @abstractmethod + def mamba_protected_size(self) -> int: ... + + @abstractmethod + def total_size(self) -> tuple[int, int]: + """(full_tokens, aux_tokens) summed across the whole tree.""" + ... + + @abstractmethod + def all_values_flatten(self) -> torch.Tensor: ... + + @abstractmethod + def all_mamba_values_flatten(self) -> torch.Tensor: ... + + @abstractmethod + def walk_for_kv_canary( + self, unlocked_only: bool, swa_resident_only: bool + ) -> RadixCacheWalkResult: + """Flatten every FULL device slot into (slot, position, prev-slot) rows + for the KV-canary sweep.""" + ... + + @abstractmethod + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: + """Match a key against the tree; returns device indices + boundary NodeIds.""" + ... + + @property + @abstractmethod + def empty_match_result(self) -> MatchResult: + """A shared empty MatchResult (empty device indices + boundary NodeIds).""" + ... + + @abstractmethod + def is_full_device_evicted(self, node_id: NodeId) -> bool: + """Whether the node's FULL device value has been evicted.""" + ... + + @abstractmethod + def collect_full_device_indices( + self, from_node_id: NodeId, until_node_id: NodeId + ) -> torch.Tensor: + """Concatenate FULL device values from from_node up to (exclusive) until_node.""" + ... + + @abstractmethod + def begin_insert(self, params: InsertParams) -> InsertStepResult: + """Start the (single-flight) insert, running to its first barrier or completion.""" + ... + + @abstractmethod + def resume_insert(self) -> InsertStepResult: + """Continue the suspended insert after its step actions were executed.""" + ... + + @abstractmethod + def has_ongoing_insert(self) -> bool: + """Whether an insert walk is suspended at a barrier.""" + ... + + @abstractmethod + def end_insert(self) -> list[CacheAction | ComponentAction]: + """Finish the insert (idempotent); returns still-pending actions to drain.""" + ... + + @abstractmethod + def drive_host_eviction( + self, component_type: ComponentType, num_tokens: int + ) -> DriveHostEvictionResult: + """Evict a component's host-side resources; no-op if the component is absent.""" + ... + + # ==== HiCache ==== + + @abstractmethod + def set_hicache_enabled(self) -> None: + """Mark the host tier (HiCache) as wired.""" + ... + + @abstractmethod + def insert_host( + self, + node_id: NodeId, + key: RadixKey, + host_value: torch.Tensor, + hash_value: list[str], + ) -> InsertResult: + """Insert a host-side (backuped) tree path descending from the given node.""" + ... + + @abstractmethod + def build_backup_spec( + self, node_id: NodeId + ) -> tuple[torch.Tensor, dict[ComponentType, list[PoolTransfer]]]: + """Read a node's device->host backup spec (device value + transfers) now.""" + ... + + @abstractmethod + def build_storage_backup_spec( + self, node_id: NodeId, pass_prefix_keys: bool + ) -> Optional[StorageBackupSpec]: + """Gather a node's device->storage backup spec; None if not backuped.""" + ... + + @abstractmethod + def build_hicache_transfers( + self, + component_type: ComponentType, + node_id: NodeId, + phase: CacheTransferPhase, + *, + host_indices: Optional[torch.Tensor] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, + ) -> Optional[list[PoolTransfer]]: + """Build a component's HiCache transfers for the given node and phase.""" + ... + + @abstractmethod + def build_load_back_spec( + self, node_id: NodeId, req: Optional[Req] = None + ) -> tuple[PoolTransfer, dict[ComponentType, list[PoolTransfer]]]: + """Build the H->D load-back KV transfer plus per-component aux transfers.""" + ... + + @abstractmethod + def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]: + """The anchor node's key extra_key.""" + ... + + @abstractmethod + def commit_hicache_transfers( + self, + node_id: NodeId, + phase: CacheTransferPhase, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + *, + cache_actions: list[CacheAction | ComponentAction], + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, + ) -> None: + """Commit each component's HiCache transfers onto the node.""" + ... + + @abstractmethod + def commit_backup( + self, + node_id: NodeId, + host_indices: torch.Tensor, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> None: + """Commit a successful backup to the node.""" + ... + + @abstractmethod + def commit_load_back( + self, + node_id: NodeId, + device_indices: torch.Tensor, + kv_xfer: PoolTransfer, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> list[CacheAction | ComponentAction]: + """Commit a successful H->D load-back onto the node; returns any cache actions.""" + ... + + @abstractmethod + def mark_write_through_pending(self, node_id: NodeId) -> None: + """Mark a node as having an in-flight write-through backup.""" + ... + + @abstractmethod + def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None: + """Clear the write-through-pending mark (when it matches ack_id) and record the + host store event for each acked node.""" + ... + + @abstractmethod + def set_component_device_value( + self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor + ) -> None: + """Store an auxiliary (non-Full) component's device value onto a node.""" + ... + + @abstractmethod + def get_component_device_value( + self, node_id: NodeId, component_type: ComponentType + ) -> Optional[torch.Tensor]: + """The component's device value on the node, or None if evicted.""" + ... + + @abstractmethod + def component_has_host_value_only( + self, node_id: NodeId, component_type: ComponentType + ) -> bool: + """Whether the component's data is device-evicted but host-backed.""" + ... + + # ==== Others ==== + + @abstractmethod + def sanity_check( + self, + ongoing_write_through: list[tuple[int, NodeId]], + ongoing_load_back: list[tuple[int, NodeId]], + ) -> None: + """Verify tree invariants and raise AssertionError on any violation. + + ongoing_write_through/ongoing_load_back are (id, node_id) pairs for in-flight ops. + """ + ... + + @abstractmethod + def pretty_print(self) -> None: + """Print the tree structure for debugging.""" + ... diff --git a/python/sglang/srt/mem_cache/unified_cache_components/README.md b/python/sglang/srt/mem_cache/unified_cache_components/README.md index b98b21628..336e4d175 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache_components/README.md @@ -17,6 +17,15 @@ A component-based, pluggable prefix cache framework for SGLang that unifies Full │ UnifiedRadixCache │ │ (unified_radix_cache.py) │ │ │ +│ controller: executes all pool/host I/O and │ +│ drains the tree's deferred Cache/Component │ +│ Actions ("tree decides, cache executes") │ +└──────────────────────┬────────────────────────┘ + ▼ +┌───────────────────────────────────────────────┐ +│ UnifiedTreeCore │ +│ (unified_cache/unified_tree_core.py) │ +│ │ │ root_node ──► UnifiedTreeNode (radix tree) │ │ components ► {ComponentType → TreeComponent} │ │ lru_lists ─► {ComponentType → UnifiedLRUList}│ @@ -61,7 +70,10 @@ node.component_data[ComponentType.MAMBA] # MambaComponent data | File | Contents | |------|----------| -| `../unified_radix_cache.py` | `UnifiedRadixCache`, `UnifiedTreeNode`, `UnifiedLRUList` | +| `../unified_radix_cache.py` | `UnifiedRadixCache` — the controller: pool/host I/O, deferred-action draining | +| `../unified_cache/unified_tree_core.py` | `UnifiedTreeCore` — the tree, LRUs, and size counters; `UnifiedTreeNode`, `UnifiedLRUList` | +| `../unified_cache/unified_tree_core_interface.py` | `UnifiedTreeCoreInterface`, `NodeId` — the tree/cache boundary contract | +| `../unified_cache/cache_action.py` | Deferred `CacheAction`/`ComponentAction` types emitted by the tree | | `tree_component.py` | `TreeComponent` ABC, `ComponentType`, `ComponentData`, `get_and_increase_time_counter`, `next_component_uuid` | | `full_component.py` | `FullComponent` — standard full-attention KV cache component | | `swa_component.py` | `SWAComponent` — sliding-window attention component with tombstone/window tracking | @@ -99,7 +111,7 @@ Find the longest cached prefix for a token sequence. - Promotes matched path to MRU in each component's LRU via `node_has_component_data()` as filter - Updates `last_access_time` with decreasing timestamps up the path (parent < child) - Concatenates matched device indices via `torch.cat` (concat length ≤ K, subsumed by O(K)) - - Calls `finalize_match_result()` per component (Mamba performs copy-on-write: allocates new pool slot, copies SSM state) + - Calls `finalize_match_result_in_tree_core()` per component (tree-side: Full/SWA host-hit sums, Mamba `branching_seqlen`); the cache then routes the static `finalize_match_result_in_cache()` per component post-walk (Mamba performs copy-on-write: allocates new pool slot, copies SSM state) --- @@ -143,7 +155,7 @@ Free cached tokens to reclaim memory. | **Complexity** | **O(E·H + L)** — E = nodes evicted, H = tombstone chain height, L = locked nodes skipped in LRU scan. | **Algorithm detail:** -1. Calls `drive_eviction()` for each component: +1. Drives each component's walk via `evict_device_start()` / `evict_device_next_node()` / `evict_device_end()`: - Full: drives eviction from `evictable_device_leaves` using `last_access_time`; only device leaves are evicted atomically - SWA: scans SWA LRU from tail; **internal** nodes are tombstoned (evict SWA data, keep node), **leaf** nodes are fully deleted; both trigger cascade - Mamba: scans Mamba LRU from tail; **internal** nodes are tombstoned, **leaf** nodes are fully deleted; both trigger cascade @@ -254,7 +266,8 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d | Hook | Purpose | Called By | Default | |------|---------|-----------|----------| | `create_match_validator(match_device_only=False)` | Return a per-match stateful predicate that decides whether a node is a valid match boundary. Full: requires Full device data, or host backup when `match_device_only=False`. SWA: tracks accumulated window length across device/host data. Mamba: requires Mamba device data, or host backup when `match_device_only=False`. | `_match_prefix_helper` | *abstract* | -| `finalize_match_result()` | Post-process the match result after prefix matching completes. Full/SWA: pass-through. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool, records `branching_seqlen`. | `_match_post_processor` | pass-through | +| `finalize_match_result_in_tree_core()` | Tree-side post-processing inside the match walk. Full/SWA: host-hit sums. Mamba: records `branching_seqlen` + the host-hit bump. | `_match_post_processor` | pass-through | +| `finalize_match_result_in_cache()` | Static, cache-level finalize after the walk (receives the cache + NodeId-based result), dispatched class-level by `UnifiedRadixCache.match_prefix`. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool. | `UnifiedRadixCache.match_prefix` | pass-through | ### Insert Phase @@ -276,8 +289,8 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d |------|---------|-----------|----------| | `evict_component(target=EvictLayer.DEVICE)` | Free this component's device, host, or both resources on a node being evicted. Internal device eviction tombstones (`value = None`); host eviction clears `host_value`. Returns `(device_freed, host_freed)`. | `_evict_component_and_detach_lru` | *abstract* | | `eviction_priority()` | Return cascade eviction priority (higher = evicted later). Leaf: all 0. Internal: Full(2) > SWA(1) > Mamba(0). When evicting, all components with ≤ priority on the same node are cascade-evicted. | `_cascade_evict` | `0` | -| `drive_eviction()` | Drive device eviction until the target amount is freed. Full: leaf-set heap. SWA/Mamba: component LRUs with internal tombstones and atomic leaf deletion. | `evict` | *abstract* | -| `drive_host_eviction()` | Drive host eviction for this component. Full uses host leaves; SWA/Mamba use host LRUs. | `evict_host` | no-op | +| `evict_device_start()` / `evict_device_next_node()` / `evict_device_end()` | Step-wise device eviction walk the Controller drives: build the cursor/heap, return the next evictable leaf (freed values collected for the Controller to drain), clear the walk. Full: leaf-set heap. SWA/Mamba: component LRUs with internal tombstones and atomic leaf deletion. | `UnifiedRadixCache._evict_components` | *abstract* | +| `drive_host_eviction()` | Drive host eviction for this component, collecting freed values for the Controller to drain. Full uses host leaves; SWA/Mamba use host LRUs. | `evict_host` | no-op | ### Lock Phase diff --git a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py index 60d27d6d4..c8fcbc4ad 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py @@ -10,6 +10,7 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import ( EvictLayer, LRURefreshPhase, PrepareLoadBackResult, + PreparePrefetchResult, TreeComponent, get_and_increase_time_counter, next_component_uuid, @@ -25,6 +26,7 @@ __all__ = [ "LRURefreshPhase", "MambaComponent", "PrepareLoadBackResult", + "PreparePrefetchResult", "SWAComponent", "TreeComponent", "_NUM_COMPONENT_TYPES", diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index 0e62d7007..93bba1dc5 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -7,7 +7,6 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, - EvictParams, IncLockRefResult, InsertResult, MatchPrefixParams, @@ -18,6 +17,7 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransfer, PoolTransferResult, ) +from sglang.srt.mem_cache.unified_cache.cache_action import FreeComponentDeviceSlot from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, @@ -26,8 +26,12 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import ( ) if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) from sglang.srt.mem_cache.unified_radix_cache import ( + NodeId, UnifiedTreeNode, ) @@ -37,13 +41,6 @@ class FullComponent(TreeComponent): def __init__(self, cache, params): super().__init__(cache, params) - allocator = cache.token_to_kv_pool_allocator - # When SWA is present, only free full-attention KV here; - # SWA KV will be freed by cascade via SWAComponent.evict_component. - if ComponentType.SWA in cache.tree_components: - self._free_full = allocator.full_attn_allocator.free - else: - self._free_full = allocator.free # HiCache state: set to host KV pool when HiCache enabled self._full_kv_pool_host = None @@ -60,7 +57,7 @@ class FullComponent(TreeComponent): node.component_data[self.component_type].value is not None or node.backuped ) - def finalize_match_result( + def finalize_match_result_in_tree_core( self, result: MatchResult, params: MatchPrefixParams, @@ -72,7 +69,7 @@ class FullComponent(TreeComponent): ct = self.component_type kv_host_hit = 0 node = result.best_match_node - root_node = self.cache.root_node + root_node = self.tree_core.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: @@ -103,6 +100,8 @@ class FullComponent(TreeComponent): def evict_component( self, node: UnifiedTreeNode, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], target: EvictLayer = EvictLayer.DEVICE, ) -> tuple[int, int]: cd = node.component_data[self.component_type] @@ -111,9 +110,9 @@ class FullComponent(TreeComponent): # Device layer if EvictLayer.DEVICE in target and cd.value is not None: - self._free_full(cd.value) + device_frees[self.component_type].append(cd.value) freed = len(cd.value) - self.cache.component_evictable_size_[self.component_type] -= freed + self.tree_core.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 @@ -121,54 +120,78 @@ class FullComponent(TreeComponent): # 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) + host_frees[self.component_type].append(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 - def drive_eviction( - self, params: EvictParams, tracker: dict[ComponentType, int] - ) -> None: - request = params.num_tokens - heap = [ - (self.cache.eviction_strategy.get_priority(n), n) - for n in self.cache.evictable_device_leaves + def _evict_device_start(self, request_cnt: int) -> None: + self._evict_device_request_cnt = request_cnt + self._evict_device_last_node = None + self._evict_device_heap = [ + (self.tree_core.eviction_strategy.get_priority(n), n) + for n in self.tree_core.evictable_device_leaves ] - heapq.heapify(heap) + heapq.heapify(self._evict_device_heap) + + def _evict_device_next_node( + self, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> Optional[NodeId]: ct = self.component_type - while tracker[ct] < request and heap: - _, x = heapq.heappop(heap) - if x not in self.cache.evictable_device_leaves: + lv = self._evict_device_last_node + if ( + lv is not None + and lv.parent is not None + and lv.parent in self.tree_core.evictable_device_leaves + ): + heapq.heappush( + self._evict_device_heap, + (self.tree_core.eviction_strategy.get_priority(lv.parent), lv.parent), + ) + self._evict_device_last_node = None + while tracker[ct] < self._evict_device_request_cnt and self._evict_device_heap: + _, x = heapq.heappop(self._evict_device_heap) + if x not in self.tree_core.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, - (self.cache.eviction_strategy.get_priority(x.parent), x.parent), - ) + self._evict_device_last_node = x + return x.id + return None + + def _evict_device_end(self) -> None: + self._evict_device_heap = [] + self._evict_device_last_node = None def drive_host_eviction( - self, num_tokens: int, tracker: dict[ComponentType, int] + self, + num_tokens: int, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], ) -> None: """Evict host leaves to free KV host pool space.""" heap = [ - (self.cache.eviction_strategy.get_priority(n), n) - for n in self.cache.evictable_host_leaves + (self.tree_core.eviction_strategy.get_priority(n), n) + for n in self.tree_core.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: + if x not in self.tree_core.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: + self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees) + if ( + x.parent is not None + and x.parent in self.tree_core.evictable_host_leaves + ): heapq.heappush( heap, - (self.cache.eviction_strategy.get_priority(x.parent), x.parent), + (self.tree_core.eviction_strategy.get_priority(x.parent), x.parent), ) def acquire_component_lock( @@ -182,15 +205,14 @@ class FullComponent(TreeComponent): # Only the last host node needs to be protected. if lock_host: cd = node.component_data[ct] - # write_back mode: the anchor may be device-only (no host_value); - # pin it anyway so the host-pressure drop fallback cannot delete it. - if cd.host_value is None and not self.cache.is_write_back: + # write_back mode: the anchor may be device-only (no host_value); pin it anyway. + if cd.host_value is None and not self.tree_core.is_write_back: return result cd.host_lock_ref += 1 - self.cache._update_evictable_leaf_sets(node) + self.tree_core._update_evictable_leaf_sets(node) return result - root = self.cache.root_node + root = self.tree_core.root_node cur = node # Skip the bottom evicted segment @@ -207,11 +229,11 @@ class FullComponent(TreeComponent): ), f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment" 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 + self.tree_core.component_evictable_size_[ct] -= key_len + self.tree_core.component_protected_size_[ct] += key_len delta += key_len cd.lock_ref += 1 - self.cache.evictable_device_leaves.discard(cur) + self.tree_core.evictable_device_leaves.discard(cur) cur = cur.parent result.delta = delta return result @@ -228,13 +250,13 @@ class FullComponent(TreeComponent): if cd.host_lock_ref == 0: return # Mirror of `acquire`. write_back uses a pure counter. - if cd.host_value is None and not self.cache.is_write_back: + if cd.host_value is None and not self.tree_core.is_write_back: return cd.host_lock_ref -= 1 - self.cache._update_evictable_leaf_sets(node) + self.tree_core._update_evictable_leaf_sets(node) return - root = self.cache.root_node + root = self.tree_core.root_node skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () cur = node while cur != root: @@ -247,11 +269,11 @@ class FullComponent(TreeComponent): 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 + self.tree_core.component_evictable_size_[ct] += key_len + self.tree_core.component_protected_size_[ct] -= key_len cd.lock_ref -= 1 if cd.lock_ref == 0: - self.cache._update_evictable_leaf_sets(cur) + self.tree_core._update_evictable_leaf_sets(cur) cur = cur.parent # ---- HiCache Hooks ---- @@ -261,7 +283,8 @@ class FullComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, *, - req: Optional[Req] = None, + mamba_pool_idx: Optional[torch.Tensor] = None, + host_indices: Optional[torch.Tensor] = None, token_ids: Optional[Sequence[int]] = None, prefetch_tokens: int = 0, last_hash: Optional[str] = None, @@ -270,7 +293,7 @@ class FullComponent(TreeComponent): if phase == CacheTransferPhase.BACKUP_HOST: # Full KV backup is handled by the main flow - # (write_backup → cache_controller.write on host_value directly). + # (cache_controller.write on host_value directly). # No extra PoolTransfer needed. return None @@ -297,7 +320,7 @@ class FullComponent(TreeComponent): else torch.empty((0,), dtype=torch.int64, device="cpu") ), device_indices=None, - nodes_to_load=nodes, + nodes_to_load=[n.id for n in nodes], ) ] @@ -309,6 +332,7 @@ class FullComponent(TreeComponent): phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), *, + cache_actions: list[CacheAction | ComponentAction], insert_result: Optional[InsertResult] = None, pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: @@ -320,19 +344,39 @@ class FullComponent(TreeComponent): elif phase == CacheTransferPhase.LOAD_BACK: if not transfers or transfers[0].device_indices is None: - self.cache._update_evictable_leaf_sets(node) + self.tree_core._update_evictable_leaf_sets(node) return xfer = transfers[0] device_indices = xfer.device_indices offset = 0 - for n in xfer.nodes_to_load or []: + for nid in xfer.nodes_to_load or []: + n = self.tree_core.node_by_id(nid) 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.tree_core.component_evictable_size_[ct] += n_len + self.tree_core._update_evictable_leaf_sets(n) - self.cache._update_evictable_leaf_sets(node) + self.tree_core._update_evictable_leaf_sets(node) + + def free_host_values(self, host_values: list[torch.Tensor]) -> None: + if self._full_kv_pool_host is None: + return + for host_value in host_values: + self._full_kv_pool_host.free(host_value) + + def apply_component_action(self, action: ComponentAction) -> None: + if isinstance(action, FreeComponentDeviceSlot): + alloc = self.cache.token_to_kv_pool_allocator + for indices in action.indices: + if self.cache.is_swa_enabled: + alloc.full_attn_allocator.free(indices) + else: + alloc.free(indices) + return + raise AssertionError( + f"FullComponent: unhandled ComponentAction {type(action).__name__}" + ) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index ea7fa3369..9d51c6988 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections import defaultdict from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch @@ -19,12 +20,18 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransfer, PoolTransferResult, ) +from sglang.srt.mem_cache.unified_cache.cache_action import ( + FreeComponentDeviceSlot, + FreeComponentHostSlot, + MambaEvictExcessPathStates, +) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, EvictLayer, LRURefreshPhase, PrepareLoadBackResult, + PreparePrefetchResult, TreeComponent, get_and_increase_time_counter, ) @@ -33,7 +40,12 @@ from sglang.srt.runtime_context import get_server_args if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) from sglang.srt.mem_cache.unified_radix_cache import ( + NodeId, UnifiedRadixCache, UnifiedTreeNode, ) @@ -46,15 +58,13 @@ class MambaComponent(TreeComponent): from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool assert isinstance( - cache.req_to_token_pool, HybridReqToTokenPool - ), f"MambaComponent requires HybridReqToTokenPool, got {type(cache.req_to_token_pool)}" + params.req_to_token_pool, HybridReqToTokenPool + ), f"MambaComponent requires HybridReqToTokenPool, got {type(params.req_to_token_pool)}" if not params.enable_mamba_extra_buffer: assert ( - cache.page_size == 1 - ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {cache.page_size}" + params.page_size == 1 + ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {params.page_size}" super().__init__(cache, params) - self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer - self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy self.mamba_cache_chunk_size = get_server_args().mamba_cache_chunk_size self.mamba_max_states_per_path = get_server_args().mamba_max_states_per_path # HiCache state @@ -78,7 +88,7 @@ class MambaComponent(TreeComponent): return case LRURefreshPhase.MATCH_END: if node.component_data[ct].value is not None: - self.cache.lru_lists[ct].reset_node_mru(node) + self.tree_core.lru_lists[ct].reset_node_mru(node) case LRURefreshPhase.INSERT_END: return case _: @@ -97,15 +107,13 @@ class MambaComponent(TreeComponent): or node.component_data[ct].host_value is not None ) - def finalize_match_result( + def finalize_match_result_in_tree_core( self, result: MatchResult, params: MatchPrefixParams, value_chunks: list[torch.Tensor], best_value_len: int, ) -> MatchResult: - cow_mamba = params.cow_mamba - req = params.req last_node = result.best_match_node mamba_boundary_len = len(result.device_indices) + result.host_hit_length @@ -120,75 +128,102 @@ class MambaComponent(TreeComponent): aligned_seqlen if aligned_seqlen > mamba_boundary_len else None ) - mamba_value = last_node.component_data[self.component_type].value - if cow_mamba and mamba_value is not None: - assert req is not None - if req.mamba_pool_idx is None: - dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1) - if dst_index is None: - # Capture the inc result and thread swa_uuid_for_lock back - # into dec. Without it, SWA's release walks past this - # request's window boundary all the way to root and - # over-decrements SWA locks held by other resident requests - # on ancestor nodes. - lock_result = self.cache.inc_lock_ref(last_node) - self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) - dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1) - self.cache.dec_lock_ref(last_node, lock_result.to_dec_params()) - assert dst_index is not None, "Can not alloc mamba cache" - req.mamba_pool_idx = dst_index[0] - req.mamba_cow_src_index = mamba_value - req.mamba_needs_clear = False - # HiCache: if mamba was evicted from device but has host backup, # ensure mamba_host_hit_length >= 1 so load_back is triggered. - cd = last_node.component_data[self.component_type] - if cd.value is None and cd.host_value is not None: + if self.has_host_value_only(last_node): result = result._replace( mamba_host_hit_length=max(result.mamba_host_hit_length, 1) ) return result._replace(mamba_branching_seqlen=branching_seqlen) + def finalize_match_result_in_cache( + self, params: MatchPrefixParams, result: MatchResult + ) -> MatchResult: + # Copy-on-write the matched device mamba state into a per-request slot. + if not params.cow_mamba: + return result + src_index = self.tree_core.get_component_device_value( + result.best_match_node, self.component_type + ) + if src_index is None: + return result + req = params.req + assert req is not None + if req.mamba_pool_idx is None: + dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1) + if dst_index is None: + # Pin the window via inc/dec_lock_ref so evict's SWA release + # stops at this request's window boundary instead of walking to + # root and over-decrementing locks held by other requests. + lock_result = self.cache.inc_lock_ref(result.best_match_node) + self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1) + self.cache.dec_lock_ref( + result.best_match_node, lock_result.to_dec_params() + ) + assert dst_index is not None, "Can not alloc mamba cache" + req.mamba_pool_idx = dst_index[0] + req.mamba_cow_src_index = src_index + req.mamba_needs_clear = False + return result + def commit_insert_component_data( self, node: UnifiedTreeNode, is_new_leaf: bool, params: InsertParams, result: InsertResult, + cache_actions: list[CacheAction | ComponentAction], ) -> None: assert params.mamba_value is not None if is_new_leaf: node.component_data[self.component_type].value = params.mamba_value - self.cache.lru_lists[self.component_type].insert_mru(node) - self.cache.component_evictable_size_[self.component_type] += len( + self.tree_core.lru_lists[self.component_type].insert_mru(node) + self.tree_core.component_evictable_size_[self.component_type] += len( params.mamba_value ) - self._evict_excess_path_states(node) + self._emit_excess_path_states_eviction(node, cache_actions) 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] + host_lru = self.tree_core.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( + self.tree_core.lru_lists[self.component_type].insert_mru(node) + self.tree_core.component_evictable_size_[self.component_type] += len( params.mamba_value ) node.last_access_time = get_and_increase_time_counter() - self._evict_excess_path_states(node) + self._emit_excess_path_states_eviction(node, cache_actions) return - self.cache.lru_lists[self.component_type].reset_node_mru(node) + self.tree_core.lru_lists[self.component_type].reset_node_mru(node) node.last_access_time = get_and_increase_time_counter() result.mamba_exist = True - def _evict_excess_path_states(self, tail: UnifiedTreeNode) -> None: + def _emit_excess_path_states_eviction( + self, + tail: UnifiedTreeNode, + cache_actions: list[CacheAction | ComponentAction], + ) -> None: + """Defer the path-cap eviction so it runs after the insert's BackupKV.""" + if self.mamba_max_states_per_path < 0: + return + cache_actions.append(MambaEvictExcessPathStates(tail.id)) + + def _evict_excess_path_states( + self, + tail: UnifiedTreeNode, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: """Evict shallow eligible device checkpoints beyond the path cap. Full KV and any existing host backup are retained. The tail, forks, - locked nodes, and device leaves are preserved, so the cap is a - best-effort soft limit. + locked nodes (including a pending backup chain's write-through locks), + and device leaves are preserved, so the cap is a best-effort soft + limit. Freed slots are collected into the caller's dicts. """ cap = self.mamba_max_states_per_path if cap < 0: @@ -197,7 +232,7 @@ class MambaComponent(TreeComponent): ct = self.component_type holders = [] node = tail - while node is not None and node is not self.cache.root_node: + while node is not None and node is not self.tree_core.root_node: if node.component_data[ct].value is not None: holders.append(node) node = node.parent @@ -212,12 +247,17 @@ class MambaComponent(TreeComponent): break if node.component_data[ct].lock_ref > 0 or len(node.children) != 1: continue - if node in self.cache.evictable_device_leaves: + if node in self.tree_core.evictable_device_leaves: continue - self.cache._evict_component_and_detach_lru( - node, self, target=EvictLayer.DEVICE, tracker=tracker + self.tree_core._evict_component_and_detach_lru( + node, + self, + device_frees, + host_frees, + target=EvictLayer.DEVICE, + tracker=tracker, ) - self.cache._cascade_evict(node, self, tracker) + self.tree_core._cascade_evict(node, self, tracker, device_frees, host_frees) excess -= 1 def redistribute_on_node_split( @@ -233,6 +273,8 @@ class MambaComponent(TreeComponent): def evict_component( self, node: UnifiedTreeNode, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], target: EvictLayer = EvictLayer.DEVICE, ) -> tuple[int, int]: cd = node.component_data[self.component_type] @@ -241,17 +283,16 @@ class MambaComponent(TreeComponent): # Device layer if EvictLayer.DEVICE in target and cd.value is not None: - self._free_mamba_value(cd.value) + device_frees[self.component_type].append(cd.value) freed = len(cd.value) - self.cache.component_evictable_size_[self.component_type] -= freed + self.tree_core.component_evictable_size_[self.component_type] -= freed cd.value = None # Host layer - host_lru = self.cache.host_lru_lists[self.component_type] + host_lru = self.tree_core.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) + host_frees[self.component_type].append(cd.host_value) cd.host_value = None if host_lru.in_list(node): host_lru.remove_node(node) @@ -267,30 +308,56 @@ class MambaComponent(TreeComponent): return freed, host_freed - def drive_eviction( - self, params: EvictParams, tracker: dict[ComponentType, int] - ) -> None: - request = params.mamba_num + def _evict_device_start(self, request_cnt: int) -> None: + """Begin the device-eviction walk from this component's LRU cursor.""" + self._evict_device_request_cnt = request_cnt + self._evict_device_cursor = self.tree_core.lru_lists[ + self.component_type + ].get_lru_no_lock() + + def _evict_device_next_node( + self, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> Optional[NodeId]: + """Return the next device-leaf node for the driver to evict, or None. + Internal nodes are tombstoned inline (no IO); the cursor is re-validated + (reset to LRU head) if the previous node's eviction removed it.""" ct = self.component_type - lru = self.cache.lru_lists[ct] - x = lru.get_lru_no_lock() - while tracker[ct] < request and x is not None and lru.in_list(x): + lru = self.tree_core.lru_lists[ct] + if self._evict_device_cursor is not None and not lru.in_list( + self._evict_device_cursor + ): + self._evict_device_cursor = lru.get_lru_no_lock() + while ( + tracker[ct] < self._evict_device_request_cnt + and self._evict_device_cursor is not None + and lru.in_list(self._evict_device_cursor) + ): + x = self._evict_device_cursor 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, target=EvictLayer.DEVICE, tracker=tracker - ) - self.cache._cascade_evict(x, self, tracker) - x = x_next + if x in self.tree_core.evictable_device_leaves: + self._evict_device_cursor = lru.get_prev_no_lock(x) + return x.id + x_next = lru.get_prev_no_lock(x) + self.tree_core._evict_component_and_detach_lru( + x, + self, + target=EvictLayer.DEVICE, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + self.tree_core._cascade_evict( + x, self, tracker, device_frees=device_frees, host_frees=host_frees + ) + self._evict_device_cursor = x_next + return None + + def _evict_device_end(self) -> None: + """Clear the device-eviction walk cursor state.""" + self._evict_device_cursor = None def acquire_component_lock( self, @@ -299,7 +366,7 @@ class MambaComponent(TreeComponent): lock_host: bool = False, ) -> IncLockRefResult: ct = self.component_type - if node is self.cache.root_node: + if node is self.tree_core.root_node: return result cd = node.component_data[ct] value = cd.host_value if lock_host else cd.value @@ -310,15 +377,15 @@ class MambaComponent(TreeComponent): if lock_host: if cd.host_lock_ref == 0: - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] if host_lru.in_list(node): host_lru.remove_node(node) cd.host_lock_ref += 1 else: if cd.lock_ref == 0: vlen = len(value) - self.cache.component_evictable_size_[ct] -= vlen - self.cache.component_protected_size_[ct] += vlen + self.tree_core.component_evictable_size_[ct] -= vlen + self.tree_core.component_protected_size_[ct] += vlen cd.lock_ref += 1 return result @@ -329,7 +396,7 @@ class MambaComponent(TreeComponent): lock_host: bool = False, ) -> None: ct = self.component_type - if node is self.cache.root_node: + if node is self.tree_core.root_node: return cd = node.component_data[ct] skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () @@ -340,7 +407,7 @@ class MambaComponent(TreeComponent): if lock_host: cd.host_lock_ref -= 1 if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None: - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] if not host_lru.in_list(node): host_lru.insert_mru(node) return @@ -348,8 +415,8 @@ class MambaComponent(TreeComponent): if cd.lock_ref > 0: if cd.lock_ref == 1: vlen = len(value) - self.cache.component_evictable_size_[ct] += vlen - self.cache.component_protected_size_[ct] -= vlen + self.tree_core.component_evictable_size_[ct] += vlen + self.tree_core.component_protected_size_[ct] -= vlen cd.lock_ref -= 1 def _alloc_mamba_slot(self) -> torch.Tensor: @@ -395,7 +462,7 @@ class MambaComponent(TreeComponent): token_ids_len: int, is_finished: bool, ) -> Optional[int]: - if self.enable_mamba_extra_buffer: + if self.cache.enable_mamba_extra_buffer: cache_len = req.mamba_last_track_seqlen else: cache_len = token_ids_len @@ -416,7 +483,7 @@ class MambaComponent(TreeComponent): if is_finished: if cache_len is None: cache_len = 0 - if self.enable_mamba_extra_buffer: + if self.cache.enable_mamba_extra_buffer: keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_keep_idx( req ) @@ -435,7 +502,7 @@ class MambaComponent(TreeComponent): return 0 # Donate the mamba index to the radix cache instead of copying. if self.int8_ckpt_pool is not None: - if self.enable_mamba_extra_buffer: + if self.cache.enable_mamba_extra_buffer: new_slot = self._alloc_mamba_slot() src_active = ( self.cache.req_to_token_pool.donate_mamba_ping_pong_slot( @@ -448,7 +515,7 @@ class MambaComponent(TreeComponent): mamba_value_donated = self._commit_int8_checkpoint( req.mamba_pool_idx.view(-1) ) - elif self.enable_mamba_extra_buffer: + elif self.cache.enable_mamba_extra_buffer: new_slot = self._alloc_mamba_slot() mamba_value_donated = ( self.cache.req_to_token_pool.donate_mamba_ping_pong_slot( @@ -491,7 +558,7 @@ class MambaComponent(TreeComponent): pool.free_mamba_cache(req) return - if self.enable_mamba_extra_buffer: + if self.cache.enable_mamba_extra_buffer: keep_idx = ( pool.get_mamba_ping_pong_keep_idx(req) if mamba_value_inserted @@ -515,17 +582,16 @@ class MambaComponent(TreeComponent): def prepare_load_back( self, - node: UnifiedTreeNode, + node_id: NodeId, *, req: Optional[Req] = None, ) -> PrepareLoadBackResult: - cd = node.component_data[self.component_type] - # skip unless the node needs a load-back (device value absent), like build_hicache_transfers if ( req is None or req.mamba_pool_idx is not None - or cd.host_value is None - or cd.value is not None + or not self.tree_core.component_has_host_value_only( + node_id, self.component_type + ) ): return PrepareLoadBackResult() dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1) @@ -544,12 +610,27 @@ class MambaComponent(TreeComponent): self.cache.req_to_token_pool.mamba_allocator.free(prep.allocated_mamba_slot) req.mamba_pool_idx = None + def prepare_prefetch( + self, + node_id: NodeId, + *, + prefetch_tokens: int = 0, + ) -> PreparePrefetchResult: + host_indices = self._mamba_pool_host.alloc(1) + if host_indices is None: + self.cache.evict_host(1, ComponentType.MAMBA) + host_indices = self._mamba_pool_host.alloc(1) + if host_indices is None: + return PreparePrefetchResult(alloc_failed=True) + return PreparePrefetchResult(host_indices=host_indices) + def build_hicache_transfers( self, node: UnifiedTreeNode, phase: CacheTransferPhase, *, - req: Optional[Req] = None, + mamba_pool_idx: Optional[torch.Tensor] = None, + host_indices: Optional[torch.Tensor] = None, token_ids: Optional[Sequence[int]] = None, prefetch_tokens: int = 0, last_hash: Optional[str] = None, @@ -580,19 +661,18 @@ class MambaComponent(TreeComponent): PoolTransfer( name=PoolName.MAMBA, host_indices=cd.host_value, - nodes_to_load=[node], + nodes_to_load=[node.id], ) ) - # Per-request mamba CoW: H→D copy into the request's device slot allocated by prepare_load_back. + # Per-request mamba CoW: H→D copy into the request's device slot pre-allocated on the caller side. cd = node.component_data[ct] - if req is not None and cd.host_value is not None: - assert req.mamba_pool_idx is not None + if mamba_pool_idx is not None and cd.host_value is not None: transfers.append( PoolTransfer( name=PoolName.MAMBA, host_indices=cd.host_value, - device_indices=req.mamba_pool_idx.unsqueeze(0), + device_indices=mamba_pool_idx.unsqueeze(0), ) ) @@ -612,12 +692,7 @@ class MambaComponent(TreeComponent): ] if phase == CacheTransferPhase.PREFETCH: - host_indices = self._mamba_pool_host.alloc(1) - if host_indices is None: - self.cache.evict_host(1, ComponentType.MAMBA) - host_indices = self._mamba_pool_host.alloc(1) - if host_indices is None: - return [] + assert host_indices is not None return [ PoolTransfer( name=PoolName.MAMBA, @@ -635,6 +710,7 @@ class MambaComponent(TreeComponent): phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), *, + cache_actions: list[CacheAction | ComponentAction], insert_result: Optional[InsertResult] = None, pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: @@ -655,11 +731,11 @@ class MambaComponent(TreeComponent): 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] + host_lru = self.tree_core.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 + self.tree_core.lru_lists[ct].insert_mru(node) + self.tree_core.component_evictable_size_[ct] += count elif phase == CacheTransferPhase.PREFETCH: if not transfers: @@ -671,7 +747,10 @@ class MambaComponent(TreeComponent): and pool_storage_result.extra_pool_hit_pages.get(PoolName.MAMBA, 0) >= 1 ) target_node = ( - insert_result.inserted_host_node if insert_result is not None else None + self.tree_core.node_by_id(insert_result.inserted_host_node) + if insert_result is not None + and insert_result.inserted_host_node is not None + else None ) if ( host_indices is None @@ -679,8 +758,10 @@ class MambaComponent(TreeComponent): or not loaded or target_node.component_data[ct].host_value is not None ): - self.cache.cache_controller.append_host_mem_release( - extra_pools=[transfer] + cache_actions.append( + FreeComponentHostSlot( + [host_indices], component_type=ComponentType.MAMBA + ) ) if insert_result is not None: insert_result.mamba_exist = True @@ -688,33 +769,86 @@ class MambaComponent(TreeComponent): target_node.component_data[ct].host_value = host_indices.clone() if target_node.component_data[ct].value is None: - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] if not host_lru.in_list(target_node): host_lru.insert_mru(target_node) if insert_result is not None: insert_result.mamba_exist = False def drive_host_eviction( - self, num_tokens: int, tracker: dict[ComponentType, int] + self, + num_tokens: int, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], ) -> 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] + host_lru = self.tree_core.host_lru_lists[ct] x = host_lru.get_lru_no_host_lock() while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x): x_next = host_lru.get_prev_no_host_lock(x) cd = x.component_data[ct] - if x in self.cache.evictable_host_leaves: + if x in self.tree_core.evictable_host_leaves: # Host leaf: atomic eviction (all components host + delete) - self.cache._evict_host_leaf(x, tracker) + self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees) 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.tree_core._evict_component_and_detach_lru( + x, + self, + target=EvictLayer.HOST, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, ) - self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST) - self.cache._update_evictable_leaf_sets(x) + self.tree_core._cascade_evict( + x, + self, + tracker, + device_frees=device_frees, + host_frees=host_frees, + target=EvictLayer.HOST, + ) + self.tree_core._update_evictable_leaf_sets(x) x = x_next + + def free_host_values(self, host_values: list[torch.Tensor]) -> None: + if self._mamba_pool_host is None: + return + for host_value in host_values: + self._mamba_pool_host.free(host_value) + + def apply_component_action(self, action: ComponentAction) -> None: + if isinstance(action, MambaEvictExcessPathStates): + device_frees: dict[ComponentType, list[torch.Tensor]] = defaultdict(list) + host_frees: dict[ComponentType, list[torch.Tensor]] = defaultdict(list) + # Drain even if the walk raises so tombstoned slots are not leaked. + try: + self._evict_excess_path_states( + self.tree_core.node_by_id(action.tail_node_id), + device_frees, + host_frees, + ) + finally: + self.cache._free_values(device_frees, host_frees) + return + if isinstance(action, FreeComponentDeviceSlot): + for indices in action.indices: + self._free_mamba_value(indices) + return + if isinstance(action, FreeComponentHostSlot): + for host_indices in action.host_indices: + if host_indices is not None and host_indices.numel() > 0: + self.cache.cache_controller.append_host_mem_release( + extra_pools=[ + PoolTransfer(name=PoolName.MAMBA, host_indices=host_indices) + ] + ) + return + raise AssertionError( + f"MambaComponent: unhandled ComponentAction {type(action).__name__}" + ) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 827b0a8e0..12621dce7 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -6,7 +6,6 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, - EvictParams, IncLockRefResult, InsertParams, InsertResult, @@ -20,12 +19,21 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransfer, PoolTransferResult, ) +from sglang.srt.mem_cache.unified_cache.cache_action import ( + FreeComponentDeviceSlot, + FreeComponentHostSlot, + FreeDeviceKV, + RebuildFullToSWAMapping, + RecoverSWAWithLockedFull, + SWARebuild, +) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( BASE_COMPONENT_TYPE, CacheTransferPhase, ComponentType, EvictLayer, LRURefreshPhase, + PreparePrefetchResult, TreeComponent, next_component_uuid, ) @@ -33,7 +41,12 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import ( if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) from sglang.srt.mem_cache.unified_radix_cache import ( + NodeId, UnifiedRadixCache, UnifiedTreeNode, ) @@ -53,8 +66,8 @@ class SWAComponent(TreeComponent): from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator assert isinstance( - cache.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator - ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(cache.token_to_kv_pool_allocator)}" + params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator + ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(params.token_to_kv_pool_allocator)}" super().__init__(cache, params) self.sliding_window_size = params.sliding_window_size # HiCache state: set to host SWA pool when HiCache enabled @@ -81,39 +94,17 @@ class SWAComponent(TreeComponent): # MATCH_END / INSERT_END instead. return case LRURefreshPhase.MATCH_END | LRURefreshPhase.INSERT_END: - self.cache.lru_lists[ + self.tree_core.lru_lists[ self.component_type ].reset_node_and_window_ancestors_mru( node, root_node, - self.sliding_window_size + self.cache.page_size, + self.sliding_window_size + self.tree_core.page_size, self.node_has_component_data, ) case _: raise ValueError(f"Unknown LRURefreshPhase: {phase}") - def _restore_device_value(self, node: UnifiedTreeNode, value: torch.Tensor) -> None: - ct = self.component_type - node.component_data[ct].value = value - 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] += len(value) - - def _restore_device_value_with_locked_full( - self, - node: UnifiedTreeNode, - full_value: torch.Tensor, - incoming_full_value: torch.Tensor, - ) -> None: - allocator = self.cache.token_to_kv_pool_allocator - swa_value = self._translate_full_to_swa(incoming_full_value) - allocator.set_full_to_swa_mapping(full_value, swa_value) - allocator.full_to_swa_index_mapping[incoming_full_value.to(torch.int64)] = 0 - allocator.full_attn_allocator.free(incoming_full_value) - self._restore_device_value(node, swa_value) - def create_match_validator( self, match_device_only: bool = False ) -> Callable[[UnifiedTreeNode], bool]: @@ -124,7 +115,7 @@ class SWAComponent(TreeComponent): # unified_kv never caches the SWA ring (per-request, not content-stable), # so SWA bookkeeping must not gate the match here. swa_device_only_hicache = ( - self._swa_kv_pool_host is None and self.cache.cache_controller is not None + not self.tree_core.has_swa_host_pool and self.tree_core.enable_hicache ) def validator(node: UnifiedTreeNode) -> bool: @@ -141,7 +132,7 @@ class SWAComponent(TreeComponent): return validator - def finalize_match_result( + def finalize_match_result_in_tree_core( self, result: MatchResult, params: MatchPrefixParams, @@ -152,7 +143,7 @@ class SWAComponent(TreeComponent): n_swa = 0 swa_host_hit = 0 node = result.best_match_node - root = self.cache.root_node + root = self.tree_core.root_node while node is not root and n_swa < self.sliding_window_size: cd = node.component_data[ct] if cd.value is not None: @@ -182,6 +173,7 @@ class SWAComponent(TreeComponent): total_prefix_len: int, value_slice: torch.Tensor, params: InsertParams, + cache_actions: list[CacheAction | ComponentAction], ) -> int: if params.prev_prefix_len >= total_prefix_len + prefix_len: return prefix_len @@ -196,40 +188,38 @@ class SWAComponent(TreeComponent): node.component_data[self.component_type].lock_ref == 0 ), f"tombstone {self.component_type} lock_ref should be 0, node {node.id}" assert ( - swa_evicted_seqlen % self.cache.page_size == 0 + swa_evicted_seqlen % self.tree_core.page_size == 0 ), f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" if swa_evicted_seqlen <= total_prefix_len: # Branch 1: entire value_slice is within SWA window — recover + old_full = full_cd.value if full_cd.lock_ref > 0: - self._restore_device_value_with_locked_full( - node, full_cd.value, value_slice + cache_actions.append( + RecoverSWAWithLockedFull(node.id, old_full, value_slice) ) return 0 - self.cache.token_to_kv_pool_allocator.free(full_cd.value) full_cd.value = value_slice.clone() - swa_value = self._translate_full_to_swa(full_cd.value) - self._restore_device_value(node, swa_value) + cache_actions.append(FreeDeviceKV([old_full])) + cache_actions.append(SWARebuild(node.id, value_slice)) return 0 elif swa_evicted_seqlen < total_prefix_len + prefix_len: # Branch 2: value_slice[start_idx:] is within SWA window — partial recover start_idx = swa_evicted_seqlen - total_prefix_len - if full_cd.lock_ref > 0: - self.cache._split_node(node.key, node, start_idx) - full_cd = node.component_data[BASE_COMPONENT_TYPE] - self._restore_device_value_with_locked_full( - node, full_cd.value, value_slice[start_idx:] + is_locked = full_cd.lock_ref > 0 + old_full = full_cd.value[start_idx:] + _, action = self.tree_core._split_node(node.key, node, start_idx) + if action is not None: + cache_actions.append(action) + new_full = value_slice[start_idx:] + if is_locked: + cache_actions.append( + RecoverSWAWithLockedFull(node.id, old_full, new_full) ) return start_idx - self.cache.token_to_kv_pool_allocator.free(full_cd.value[start_idx:]) - self.cache._split_node(node.key, node, start_idx) - node.component_data[BASE_COMPONENT_TYPE].value = value_slice[ - start_idx: - ].clone() - swa_value = self._translate_full_to_swa( - node.component_data[BASE_COMPONENT_TYPE].value - ) - self._restore_device_value(node, swa_value) + node.component_data[BASE_COMPONENT_TYPE].value = new_full.clone() + cache_actions.append(FreeDeviceKV([old_full])) + cache_actions.append(SWARebuild(node.id, new_full)) return start_idx else: # Branch 3: entire value_slice is outside SWA window — not consumed @@ -241,6 +231,7 @@ class SWAComponent(TreeComponent): prefix_len: int, total_prefix_len: int, params: InsertParams, + cache_actions: list[CacheAction | ComponentAction], ) -> None: # _unevict_node_on_insert already wrote the request's fresh KV slice # into the base value. We just need to rebuild SWA from that slice for @@ -253,20 +244,24 @@ class SWAComponent(TreeComponent): ), f"tombstone {ct} lock_ref should be 0 on unevict, node {node.id}" swa_evicted_seqlen = params.swa_evicted_seqlen assert ( - swa_evicted_seqlen % self.cache.page_size == 0 + swa_evicted_seqlen % self.tree_core.page_size == 0 ), f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" - full_value = node.component_data[BASE_COMPONENT_TYPE].value if swa_evicted_seqlen <= total_prefix_len: - swa_value = self._translate_full_to_swa(full_value) + pass # entire node is within the SWA window elif swa_evicted_seqlen < total_prefix_len + prefix_len: start_idx = swa_evicted_seqlen - total_prefix_len - self.cache._split_node(node.key, node, start_idx) - full_value = node.component_data[BASE_COMPONENT_TYPE].value - swa_value = self._translate_full_to_swa(full_value) + _, action = self.tree_core._split_node(node.key, node, start_idx) + if action is not None: + cache_actions.append(action) else: return - self._restore_device_value(node, swa_value) + cache_actions.append( + SWARebuild( + node.id, + node.component_data[BASE_COMPONENT_TYPE].value, + ) + ) def commit_insert_component_data( self, @@ -274,57 +269,65 @@ class SWAComponent(TreeComponent): is_new_leaf: bool, params: InsertParams, result: InsertResult, + cache_actions: list[CacheAction | ComponentAction], ) -> None: if not is_new_leaf: return node_start = result.prefix_len split_pos = params.swa_evicted_seqlen - node_start - - if split_pos <= 0: - swa_value = self._translate_full_to_swa( - node.component_data[BASE_COMPONENT_TYPE].value - ) - node.component_data[self.component_type].value = swa_value - self.cache.lru_lists[self.component_type].insert_mru(node) - self.cache.component_evictable_size_[self.component_type] += len(swa_value) - elif split_pos < len(node.key): - # Node straddles the SWA eviction boundary - # Split into parent (tombstone, no SWA) and child (with SWA) - # After _split_node, `node` becomes the child - self.cache._split_node(node.key, node, split_pos) - swa_value = self._translate_full_to_swa( - node.component_data[BASE_COMPONENT_TYPE].value - ) - node.component_data[self.component_type].value = swa_value - self.cache.lru_lists[self.component_type].insert_mru(node) - self.cache.component_evictable_size_[self.component_type] += len(swa_value) - else: + if split_pos >= len(node.key): # Entire leaf is outside the SWA window — left as a tombstone. return + if split_pos > 0: + # Node straddles the boundary: split into an out-of-window parent + # (tombstone) and an in-window child; `node` becomes the child. + _, action = self.tree_core._split_node(node.key, node, split_pos) + assert action is None, "new leaf cannot be write-through-pending" + # Cap the in-window leaf at one window for lock granularity, then rebuild SWA + # onto the in-window node(s) at apply time; rebuild the older prefix first so + # the in-window tail lands more-MRU. + capped_parent = self._maybe_split_leaf_for_swa_lock(node) + if capped_parent is not None: + cache_actions.append( + SWARebuild( + capped_parent.id, + capped_parent.component_data[BASE_COMPONENT_TYPE].value, + ) + ) + cache_actions.append( + SWARebuild( + node.id, + node.component_data[BASE_COMPONENT_TYPE].value, + ) + ) - self._maybe_split_leaf_for_swa_lock(node) - - def _maybe_split_leaf_for_swa_lock(self, leaf: UnifiedTreeNode) -> None: - """Cap a fresh SWA leaf at one page-aligned window so locking it pins - only one window of SWA pool, not the whole (long chunked-prefill) leaf. + def _maybe_split_leaf_for_swa_lock( + self, leaf: UnifiedTreeNode + ) -> Optional[UnifiedTreeNode]: + """Cap a fresh in-window SWA leaf at one page-aligned window so locking it pins + only one window of SWA pool, not the whole (long chunked-prefill) leaf; return + the split-off parent (older window) or None. The SWA value is stamped later, so + this runs on the tombstone leaf. """ ct = self.component_type cd = leaf.component_data[ct] - if leaf is self.cache.root_node or cd.value is None or cd.lock_ref > 0: - return + if leaf is self.tree_core.root_node or cd.lock_ref > 0: + return None - page_size = self.cache.page_size + page_size = self.tree_core.page_size # Smallest page-aligned size that still covers the sliding window. tail_size = (self.sliding_window_size + page_size - 1) // page_size * page_size leaf_len = len(leaf.key) if leaf_len <= tail_size: - return + return None split_at = leaf_len - tail_size if page_size > 1 and (split_at % page_size != 0 or leaf_len % page_size != 0): - return + return None - self.cache._split_node(leaf.key, leaf, split_at) + new_parent, action = self.tree_core._split_node(leaf.key, leaf, split_at) + assert action is None, "fresh SWA leaf cannot be write-through-pending" + return new_parent def redistribute_on_node_split( self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode @@ -354,7 +357,7 @@ class SWAComponent(TreeComponent): child.component_data[self.component_type].host_value = child_swa_host_value[ split_len: ].clone() - host_lru = self.cache.host_lru_lists[self.component_type] + host_lru = self.tree_core.host_lru_lists[self.component_type] if new_parent.component_data[self.component_type].value is None: host_lru.insert_mru(new_parent) if child.component_data[ @@ -371,6 +374,8 @@ class SWAComponent(TreeComponent): def evict_component( self, node: UnifiedTreeNode, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], target: EvictLayer = EvictLayer.DEVICE, ) -> tuple[int, int]: ct = self.component_type @@ -383,19 +388,18 @@ class SWAComponent(TreeComponent): # Pass full indices to free_swa so slots with no SWA pair are # skipped. Freeing swa_value directly would double free those # entries since they all map to the same sentinel slot. - self.cache.token_to_kv_pool_allocator.free_swa( + device_frees[self.component_type].append( node.component_data[BASE_COMPONENT_TYPE].value ) freed = len(cd.value) - self.cache.component_evictable_size_[ct] -= freed + self.tree_core.component_evictable_size_[ct] -= freed cd.value = None # Host layer - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] if EvictLayer.HOST in target and cd.host_value is not None: host_freed = len(cd.host_value) - if self._swa_kv_pool_host is not None: - self._swa_kv_pool_host.free(cd.host_value) + host_frees[ct].append(cd.host_value) cd.host_value = None if host_lru.in_list(node): host_lru.remove_node(node) @@ -414,30 +418,56 @@ class SWAComponent(TreeComponent): def eviction_priority(self, is_leaf: bool) -> int: return 0 if is_leaf else 1 - def drive_eviction( - self, params: EvictParams, tracker: dict[ComponentType, int] - ) -> None: - request = params.swa_num_tokens + def _evict_device_start(self, request_cnt: int) -> None: + """Begin the device-eviction walk from this component's LRU cursor.""" + self._evict_device_request_cnt = request_cnt + self._evict_device_cursor = self.tree_core.lru_lists[ + self.component_type + ].get_lru_no_lock() + + def _evict_device_next_node( + self, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> Optional[NodeId]: + """Return the next device-leaf node for the driver to evict, or None. + Internal nodes are tombstoned inline (no IO); the cursor is re-validated + (reset to LRU head) if the previous node's eviction removed it.""" ct = self.component_type - lru = self.cache.lru_lists[ct] - x = lru.get_lru_no_lock() - while tracker[ct] < request and x is not None and lru.in_list(x): + lru = self.tree_core.lru_lists[ct] + if self._evict_device_cursor is not None and not lru.in_list( + self._evict_device_cursor + ): + self._evict_device_cursor = lru.get_lru_no_lock() + while ( + tracker[ct] < self._evict_device_request_cnt + and self._evict_device_cursor is not None + and lru.in_list(self._evict_device_cursor) + ): + x = self._evict_device_cursor 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, target=EvictLayer.DEVICE, tracker=tracker - ) - self.cache._cascade_evict(x, self, tracker) - x = x_next + if x in self.tree_core.evictable_device_leaves: + self._evict_device_cursor = lru.get_prev_no_lock(x) + return x.id + x_next = lru.get_prev_no_lock(x) + self.tree_core._evict_component_and_detach_lru( + x, + self, + target=EvictLayer.DEVICE, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + self.tree_core._cascade_evict( + x, self, tracker, device_frees=device_frees, host_frees=host_frees + ) + self._evict_device_cursor = x_next + return None + + def _evict_device_end(self) -> None: + """Clear the device-eviction walk cursor state.""" + self._evict_device_cursor = None def acquire_component_lock( self, @@ -446,12 +476,16 @@ class SWAComponent(TreeComponent): lock_host: bool = False, ) -> IncLockRefResult: ct = self.component_type - root = self.cache.root_node + root = self.tree_core.root_node sliding_window_size = self.sliding_window_size swa_lock_size = 0 swa_uuid = None uuid_key = "host_uuid" if lock_host else "uuid" - lru = self.cache.host_lru_lists[ct] if lock_host else self.cache.lru_lists[ct] + lru = ( + self.tree_core.host_lru_lists[ct] + if lock_host + else self.tree_core.lru_lists[ct] + ) # Tombstoned nodes (cd.value is None) have no SWA chunk to protect # skip them and keep walking up. This path is hit when HiCache @@ -472,8 +506,8 @@ class SWAComponent(TreeComponent): lru.remove_node(cur) else: key_len = len(cur.key) - self.cache.component_evictable_size_[ct] -= key_len - self.cache.component_protected_size_[ct] += key_len + self.tree_core.component_evictable_size_[ct] -= key_len + self.tree_core.component_protected_size_[ct] += key_len if lock_host: comp.host_lock_ref = ref + 1 else: @@ -498,7 +532,7 @@ class SWAComponent(TreeComponent): lock_host: bool = False, ) -> None: ct = self.component_type - root = self.cache.root_node + root = self.tree_core.root_node swa_uuid_for_lock = ( (params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock) if params @@ -522,13 +556,13 @@ class SWAComponent(TreeComponent): if ref == 1: if lock_host: if comp.value is None and comp.host_value is not None: - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] if not host_lru.in_list(cur): host_lru.insert_mru(cur) else: key_len = len(comp.value) - self.cache.component_evictable_size_[ct] += key_len - self.cache.component_protected_size_[ct] -= key_len + self.tree_core.component_evictable_size_[ct] += key_len + self.tree_core.component_protected_size_[ct] -= key_len if lock_host: comp.host_lock_ref = ref - 1 else: @@ -540,7 +574,9 @@ class SWAComponent(TreeComponent): def release_window_lock( self, node: UnifiedTreeNode, - swa_uuid_for_lock: Optional[int] = None, + swa_uuid_for_lock: Optional[int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], ) -> None: """Early-release the SWA lock along [node, swa_uuid_for_lock] while leaving Full and Mamba locks intact. @@ -553,7 +589,7 @@ class SWAComponent(TreeComponent): invoked at most once per (node, swa_uuid_for_lock) pair. """ ct = self.component_type - root = self.cache.root_node + root = self.tree_core.root_node cur = node while cur is not root: @@ -569,11 +605,15 @@ class SWAComponent(TreeComponent): cd.lock_ref -= 1 if cd.lock_ref == 0: key_len = len(cur.key) - self.cache.component_protected_size_[ct] -= key_len - self.cache.component_evictable_size_[ct] += key_len - if self.cache._is_device_leaf(cur): - self.cache._evict_component_and_detach_lru( - cur, self, target=EvictLayer.DEVICE + self.tree_core.component_protected_size_[ct] -= key_len + self.tree_core.component_evictable_size_[ct] += key_len + if self.tree_core._is_device_leaf(cur): + self.tree_core._evict_component_and_detach_lru( + cur, + self, + target=EvictLayer.DEVICE, + device_frees=device_frees, + host_frees=host_frees, ) if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock: @@ -608,12 +648,36 @@ class SWAComponent(TreeComponent): # ---- HiCache Hooks ---- + def prepare_prefetch( + self, + node_id: NodeId, + *, + prefetch_tokens: int = 0, + ) -> PreparePrefetchResult: + # unified_kv keeps SWA as a device-only ring -- nothing to prefetch into. + if self._swa_kv_pool_host is None: + return PreparePrefetchResult() + sw_pages = ( + self.cache.sliding_window_size + self.cache.page_size - 1 + ) // self.cache.page_size + if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages: + return PreparePrefetchResult() + num_tokens = sw_pages * self.cache.page_size + host_indices = self._swa_kv_pool_host.alloc(num_tokens) + if host_indices is None: + self.cache.evict_host(num_tokens, ComponentType.SWA) + host_indices = self._swa_kv_pool_host.alloc(num_tokens) + if host_indices is None: + return PreparePrefetchResult(alloc_failed=True) + return PreparePrefetchResult(host_indices=host_indices) + def build_hicache_transfers( self, node: UnifiedTreeNode, phase: CacheTransferPhase, *, - req: Optional[Req] = None, + mamba_pool_idx: Optional[torch.Tensor] = None, + host_indices: Optional[torch.Tensor] = None, token_ids: Optional[Sequence[int]] = None, prefetch_tokens: int = 0, last_hash: Optional[str] = None, @@ -621,7 +685,7 @@ class SWAComponent(TreeComponent): ct = self.component_type # unified_kv keeps SWA as a device-only ring. - if self._swa_kv_pool_host is None and self.cache.cache_controller is not None: + if not self.tree_core.has_swa_host_pool and self.tree_core.enable_hicache: return None if phase == CacheTransferPhase.BACKUP_HOST: @@ -644,7 +708,9 @@ class SWAComponent(TreeComponent): backed_up: list[torch.Tensor] = [] nodes: list = [] cur = node - while cur is not self.cache.root_node and n_swa < self.sliding_window_size: + while ( + cur is not self.tree_core.root_node and n_swa < self.sliding_window_size + ): cd = cur.component_data[ct] assert cd.host_value is not None or cd.value is not None if cd.value is not None: @@ -668,7 +734,7 @@ class SWAComponent(TreeComponent): name=PoolName.SWA, host_indices=torch.cat(backed_up), device_indices=None, - nodes_to_load=nodes, + nodes_to_load=[n.id for n in nodes], ) ] @@ -676,32 +742,21 @@ class SWAComponent(TreeComponent): cd = node.component_data[ct] if cd.host_value is None or not node.hash_value: return None - num_pages = len(cd.host_value) // self.cache.page_size + num_pages = len(cd.host_value) // self.tree_core.page_size if num_pages == 0: return None return [ PoolTransfer( name=PoolName.SWA, - host_indices=cd.host_value[-num_pages * self.cache.page_size :], + host_indices=cd.host_value[-num_pages * self.tree_core.page_size :], keys=node.hash_value[-num_pages:], hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] if phase == CacheTransferPhase.PREFETCH: - # Require a full sliding window. - sw_pages = ( - self.sliding_window_size + self.cache.page_size - 1 - ) // self.cache.page_size - if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages: - return None - num_tokens = sw_pages * self.cache.page_size - host_indices = self._swa_kv_pool_host.alloc(num_tokens) - if host_indices is None: - self.cache.evict_host(num_tokens, ComponentType.SWA) - host_indices = self._swa_kv_pool_host.alloc(num_tokens) - if host_indices is None: - return [] + assert host_indices is not None + sw_pages = host_indices.numel() // self.tree_core.page_size return [ PoolTransfer( name=PoolName.SWA, @@ -719,6 +774,7 @@ class SWAComponent(TreeComponent): phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), *, + cache_actions: list[CacheAction | ComponentAction], insert_result: Optional[InsertResult] = None, pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: @@ -735,35 +791,47 @@ class SWAComponent(TreeComponent): assert transfers and transfers[0].device_indices is not None xfer = transfers[0] device_indices = xfer.device_indices - allocator = self.cache.token_to_kv_pool_allocator + full_chunks: list[torch.Tensor] = [] + swa_chunks: list[torch.Tensor] = [] offset = 0 - for n in xfer.nodes_to_load or []: + for nid in xfer.nodes_to_load or []: + n = self.tree_core.node_by_id(nid) cd_n = n.component_data[ct] cd_full_n = n.component_data[BASE_COMPONENT_TYPE] n_tokens = len(cd_n.host_value) swa_chunk = device_indices[offset : offset + n_tokens].clone() - self._restore_device_value(n, swa_chunk) + self.tree_core.set_component_device_value( + n.id, self.component_type, swa_chunk + ) assert cd_full_n.value is not None and len(cd_full_n.value) == n_tokens - # rebuild the mapping for the loaded SWA chunk - allocator.set_full_to_swa_mapping(cd_full_n.value, swa_chunk) + full_chunks.append(cd_full_n.value) + swa_chunks.append(swa_chunk) offset += n_tokens assert offset == len(xfer.host_indices) + # rebuild the mapping for the loaded SWA chunk, defer to orchestrator level + if full_chunks: + cache_actions.append(RebuildFullToSWAMapping(full_chunks, swa_chunks)) return if phase == CacheTransferPhase.PREFETCH: self._commit_prefetch( node, transfers, + cache_actions=cache_actions, insert_result=insert_result, pool_storage_result=pool_storage_result, ) return - def _release_swa_host(self, host_indices: torch.Tensor) -> None: + def _release_swa_host( + self, + host_indices: torch.Tensor, + cache_actions: list[CacheAction | ComponentAction], + ) -> None: if host_indices is not None and host_indices.numel() > 0: - self.cache.cache_controller.append_host_mem_release( - extra_pools=[PoolTransfer(name=PoolName.SWA, host_indices=host_indices)] + cache_actions.append( + FreeComponentHostSlot([host_indices], component_type=ComponentType.SWA) ) def _attach_swa_host_value( @@ -773,18 +841,19 @@ class SWAComponent(TreeComponent): ct = self.component_type cd = node.component_data[ct] cd.host_value = host_indices.clone() - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] if cd.value is None and not host_lru.in_list(node): host_lru.insert_mru(node) - self.cache._update_evictable_leaf_sets(node) + self.tree_core._update_evictable_leaf_sets(node) if node.parent: - self.cache._update_evictable_leaf_sets(node.parent) + self.tree_core._update_evictable_leaf_sets(node.parent) def _commit_prefetch( self, anchor, transfers: list[PoolTransfer], *, + cache_actions: list[CacheAction | ComponentAction], insert_result: Optional[InsertResult] = None, pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: @@ -799,7 +868,7 @@ class SWAComponent(TreeComponent): if not transfers: return ct = self.component_type - page_size = self.cache.page_size + page_size = self.tree_core.page_size host_indices = transfers[0].host_indices window_require_pages = ( host_indices.numel() // page_size if host_indices is not None else 0 @@ -809,13 +878,18 @@ class SWAComponent(TreeComponent): if pool_storage_result else 0 ) - target = insert_result.inserted_host_node if insert_result else None + target = ( + self.tree_core.node_by_id(insert_result.inserted_host_node) + if insert_result is not None + and insert_result.inserted_host_node is not None + else None + ) if ( target is None or window_require_pages == 0 or loaded_pages < window_require_pages ): - self._release_swa_host(host_indices) + self._release_swa_host(host_indices, cache_actions) return # Buffer covers token range [loaded_start, total_len). @@ -835,37 +909,105 @@ class SWAComponent(TreeComponent): if cd.host_value is None and fill_len > 0: # Tombstone: split off the in-buffer tail if needed, then fill. if fill_start > node_start: - self.cache._split_node(cur.key, cur, fill_start - node_start) + _, action = self.tree_core._split_node( + cur.key, cur, fill_start - node_start + ) + if action is not None: + cache_actions.append(action) self._attach_swa_host_value(cur, slice_) else: # Already has SWA (or empty overlap): drop this slice. - self._release_swa_host(slice_) + self._release_swa_host(slice_, cache_actions) pos = node_start cur = cur.parent # Buffer prefix that fell outside the anchor→leaf path. if pos > loaded_start: - self._release_swa_host(host_indices[: pos - loaded_start]) + self._release_swa_host(host_indices[: pos - loaded_start], cache_actions) def drive_host_eviction( - self, num_tokens: int, tracker: dict[ComponentType, int] + self, + num_tokens: int, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], ) -> None: """Evict SWA host resources. Internal nodes: private tombstone (free SWA host only). Host leaves: atomic eviction via _evict_host_leaf.""" ct = self.component_type - host_lru = self.cache.host_lru_lists[ct] + host_lru = self.tree_core.host_lru_lists[ct] x = host_lru.get_lru_no_host_lock() while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x): x_next = host_lru.get_prev_no_host_lock(x) cd = x.component_data[ct] - if x in self.cache.evictable_host_leaves: - self.cache._evict_host_leaf(x, tracker) + if x in self.tree_core.evictable_host_leaves: + self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees) else: assert cd.host_value is not None - self.cache._evict_component_and_detach_lru( - x, self, target=EvictLayer.HOST, tracker=tracker + self.tree_core._evict_component_and_detach_lru( + x, + self, + target=EvictLayer.HOST, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + self.tree_core._cascade_evict( + x, + self, + tracker, + device_frees=device_frees, + host_frees=host_frees, + target=EvictLayer.HOST, ) - self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST) x = x_next + + def free_host_values(self, host_values: list[torch.Tensor]) -> None: + if self._swa_kv_pool_host is None: + return + for host_value in host_values: + self._swa_kv_pool_host.free(host_value) + + def apply_component_action(self, action: ComponentAction) -> None: + alloc = self.cache.token_to_kv_pool_allocator + if isinstance(action, FreeComponentDeviceSlot): + for indices in action.indices: + alloc.free_swa(indices) + return + if isinstance(action, FreeComponentHostSlot): + for host_indices in action.host_indices: + if host_indices is not None and host_indices.numel() > 0: + self.cache.cache_controller.append_host_mem_release( + extra_pools=[ + PoolTransfer(name=PoolName.SWA, host_indices=host_indices) + ] + ) + return + if isinstance(action, RebuildFullToSWAMapping): + assert len(action.full_indices) == len(action.swa_indices) + for full, swa in zip(action.full_indices, action.swa_indices): + alloc.set_full_to_swa_mapping(full, swa) + return + if isinstance(action, RecoverSWAWithLockedFull): + # Keep the locked full; remap it onto the incoming full's SWA translation, + # freeing only the incoming full, then store the swa on the node. + swa_value = self._translate_full_to_swa(action.incoming_full) + alloc.set_full_to_swa_mapping(action.kept_full, swa_value) + alloc.full_to_swa_index_mapping[action.incoming_full.to(torch.int64)] = 0 + alloc.full_attn_allocator.free(action.incoming_full) + self.tree_core.set_component_device_value( + action.node_id, self.component_type, swa_value + ) + return + if isinstance(action, SWARebuild): + # Translate the node's source full value to SWA and store it on the node. + swa_value = self._translate_full_to_swa(action.source_value) + self.tree_core.set_component_device_value( + action.node_id, self.component_type, swa_value + ) + return + raise AssertionError( + f"SWAComponent: unhandled ComponentAction {type(action).__name__}" + ) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index 43b723ed0..fb8db4ec2 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -10,7 +10,6 @@ from numpy import float64 from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, - EvictParams, IncLockRefResult, InsertParams, InsertResult, @@ -18,40 +17,26 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchResult, ) from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult +from sglang.srt.mem_cache.unified_cache.component_type import ( # noqa: F401 + BASE_COMPONENT_TYPE, + ComponentType, +) if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) + from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_radix_cache import ( + NodeId, UnifiedRadixCache, UnifiedTreeNode, ) -class ComponentType(int, Enum): - """Integer enum so that per-node list/tuple storage can be indexed directly.""" - - FULL = 0 - SWA = 1 - MAMBA = 2 - - def __str__(self) -> str: # keep human-readable logging - return self.name.lower() - - @property - def is_full(self) -> bool: - return self == ComponentType.FULL - - @property - def is_swa(self) -> bool: - return self == ComponentType.SWA - - @property - def is_mamba(self) -> bool: - return self == ComponentType.MAMBA - - -BASE_COMPONENT_TYPE = ComponentType.FULL _NUM_COMPONENT_TYPES = len(ComponentType) _LAST_ACCESS_TIME_COUNTER_FLOAT = float64(1.0) @@ -83,6 +68,16 @@ class PrepareLoadBackResult: allocated_mamba_slot: Optional[torch.Tensor] = None +@dataclasses.dataclass(frozen=True) +class PreparePrefetchResult: + """Outcome of prepare_prefetch; default = nothing to prepare.""" + + # Host pool exhausted; the caller aborts the prefetch. + alloc_failed: bool = False + # The component's pre-allocated host buffer (None = skip the build). + host_indices: Optional[torch.Tensor] = None + + class CacheTransferPhase(str, Enum): BACKUP_HOST = "backup_host" # D→H @@ -114,6 +109,9 @@ def next_component_uuid() -> int: class TreeComponent(ABC): def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): self.cache = cache + # Populated when the component passed to TreeCore constructor. + self.tree_core: Optional[UnifiedTreeCore] = None + self.is_evict_device_ongoing = False # Subclasses MUST set this as a class attribute (not @property) component_type: ComponentType @@ -130,6 +128,11 @@ class TreeComponent(ABC): value = node.component_data[self.component_type].value return len(value) if value is not None else 0 + def has_host_value_only(self, node: UnifiedTreeNode) -> bool: + """Whether this component's data is evicted from device but host-backed.""" + cd = node.component_data[self.component_type] + return cd.value is None and cd.host_value is not None + def refresh_lru( self, phase: LRURefreshPhase, @@ -141,9 +144,9 @@ class TreeComponent(ABC): case LRURefreshPhase.WALKDOWN: if node.component_data[ct].value is None: return - self.cache.lru_lists[ct].reset_node_mru(node) + self.tree_core.lru_lists[ct].reset_node_mru(node) case LRURefreshPhase.MATCH_END: - self.cache.lru_lists[ct].reset_node_and_parents_mru( + self.tree_core.lru_lists[ct].reset_node_and_parents_mru( node, root_node, self.node_has_component_data ) case LRURefreshPhase.INSERT_END: @@ -168,18 +171,22 @@ class TreeComponent(ABC): - Mamba: returns True iff the node has mamba component data.""" ... - def finalize_match_result( + def finalize_match_result_in_tree_core( self, result: MatchResult, params: MatchPrefixParams, value_chunks: list[torch.Tensor], best_value_len: int, ) -> MatchResult: - """Post-process the match result after prefix matching completes. - - Full & SWA: pass through unchanged. - - Mamba: performs copy-on-write — allocates a new mamba slot, copies - the matched node's mamba state into the request pool, and records - branching_seqlen in result.""" + """Tree-side post-processing inside the match walk (no cache access).""" + return result + + def finalize_match_result_in_cache( + self, params: MatchPrefixParams, result: MatchResult + ) -> MatchResult: + """Cache-level finalize after the match walk, dispatched by + `UnifiedRadixCache.match_prefix`; receives the NodeId-based result. + - Mamba: performs the copy-on-write into a per-request slot.""" return result def update_component_on_insert_overlap( @@ -189,6 +196,7 @@ class TreeComponent(ABC): total_prefix_len: int, value_slice: torch.Tensor, params: InsertParams, + cache_actions: list[CacheAction | ComponentAction], ) -> int: """Called per-node when an insert's key overlaps an existing node. Returns the index within value_slice from which this component @@ -204,6 +212,7 @@ class TreeComponent(ABC): prefix_len: int, total_prefix_len: int, params: InsertParams, + cache_actions: list[CacheAction | ComponentAction], ) -> None: """Called after _unevict_node_on_insert restores the base (Full) value on an evicted node. Aux components (e.g. SWA) override this to rebuild @@ -217,6 +226,7 @@ class TreeComponent(ABC): is_new_leaf: bool, params: InsertParams, result: InsertResult, + cache_actions: list[CacheAction | ComponentAction], ) -> None: """Finalize component data on the target (leaf) node after the insert walk completes. Called once per insert. @@ -250,6 +260,8 @@ class TreeComponent(ABC): def evict_component( self, node: UnifiedTreeNode, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], target: EvictLayer = EvictLayer.DEVICE, ) -> tuple[int, int]: """Free this component's KV resources on a node being evicted. @@ -262,6 +274,9 @@ class TreeComponent(ABC): - ALL: free both device and host memory. No tombstone — caller will delete the node. + Freed indices are collected into *device_frees*/*host_frees* for the + Controller to drain, never freed inline. + Returns (device_freed, host_freed) token counts.""" ... @@ -292,17 +307,53 @@ class TreeComponent(ABC): - Full evict internal: cascades to SWA + Mamba.""" return 0 + def evict_device_start(self, request_cnt: int) -> None: + """Begin this component's device-eviction walk (build its cursor/heap).""" + assert ( + not self.is_evict_device_ongoing + ), f"{self.component_type} device eviction already in progress" + self._evict_device_start(request_cnt) + self.is_evict_device_ongoing = True + + def evict_device_next_node( + self, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> Optional[NodeId]: + """Return the next device-leaf node for the driver to evict, or None. + Internal nodes are tombstoned inline (no IO).""" + assert ( + self.is_evict_device_ongoing + ), f"{self.component_type} device eviction not started" + return self._evict_device_next_node(tracker, device_frees, host_frees) + + def evict_device_end(self) -> None: + """Clear this component's device-eviction walk state.""" + assert ( + self.is_evict_device_ongoing + ), f"{self.component_type} device eviction not started" + self._evict_device_end() + self.is_evict_device_ongoing = False + @abstractmethod - def drive_eviction( - self, params: EvictParams, tracker: dict[ComponentType, int] - ) -> None: - """Drive eviction from this component's LRU list. - Each component extracts its own request from params, walks its own - LRU, evicts, and calls cache._cascade_evict for priority cascade. - Updates the shared tracker with freed amounts for all components. - - Full: walks leaf LRU, evicts full then cascades entire leaf. - - Mamba: walks full LRU; tombstones internal nodes (with cascade - to equal-priority components like swa), cascades leaves to all.""" + def _evict_device_start(self, request_cnt: int) -> None: + """Build this component's eviction cursor/heap.""" + ... + + @abstractmethod + def _evict_device_next_node( + self, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> Optional[NodeId]: + """Advance the walk; return the next device leaf or None.""" + ... + + @abstractmethod + def _evict_device_end(self) -> None: + """Clear this component's eviction cursor/heap.""" ... @abstractmethod @@ -386,11 +437,15 @@ class TreeComponent(ABC): ) -> None: pass + def free_host_values(self, host_values: list[torch.Tensor]) -> None: + """Free evicted host-tier values back to this component's host pool.""" + raise NotImplementedError(f"{self.component_type} must free its host values") + # ---- HiCache Hooks ---- def prepare_load_back( self, - node: UnifiedTreeNode, + node_id: NodeId, *, req: Optional[Req] = None, ) -> PrepareLoadBackResult: @@ -404,12 +459,22 @@ class TreeComponent(ABC): not go through.""" pass + def prepare_prefetch( + self, + node_id: NodeId, + *, + prefetch_tokens: int = 0, + ) -> PreparePrefetchResult: + """Cache-level host pre-allocation before a prefetch builds its transfers.""" + return PreparePrefetchResult() + def build_hicache_transfers( self, node: UnifiedTreeNode, phase: CacheTransferPhase, *, - req: Optional[Req] = None, + mamba_pool_idx: Optional[torch.Tensor] = None, + host_indices: Optional[torch.Tensor] = None, token_ids: Optional[Sequence[int]] = None, prefetch_tokens: int = 0, last_hash: Optional[str] = None, @@ -424,6 +489,7 @@ class TreeComponent(ABC): phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), *, + cache_actions: list[CacheAction | ComponentAction], insert_result: Optional[InsertResult] = None, pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: @@ -431,9 +497,20 @@ class TreeComponent(ABC): pass def drive_host_eviction( - self, num_tokens: int, tracker: dict[ComponentType, int] + self, + num_tokens: int, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], ) -> None: - """Evict from this component's host-side resources. + """Evict from this component's host-side resources, collecting freed + values into *device_frees*/*host_frees* for the Controller to drain. Called by HostPoolGroup when the host pool is full. Default no-op for components without host storage.""" pass + + def apply_component_action(self, action: ComponentAction) -> None: + """Apply a component-routed cache action; dispatched by the cache.""" + raise NotImplementedError( + f"{self.component_type} cannot apply {type(action).__name__}" + ) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index a0c67820a..69477a27a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1,18 +1,13 @@ from __future__ import annotations import logging -import sys import threading import time -from array import array -from collections import defaultdict -from functools import partial from queue import Empty, Queue -from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, Optional, TypeVar +from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, TypeVar import torch -from sglang.srt.disaggregation.kv_events import StorageMedium from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import ( @@ -28,7 +23,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, PoolName, @@ -39,25 +33,33 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.unified_cache.cache_action import ( + BackupKV, + CacheAction, + ComponentAction, + FreeComponentDeviceSlot, + FreeDeviceKV, + ReplaceWriteThroughOnNodeSplit, +) +from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core +from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401 + NodeId, + UnifiedLRUList, + UnifiedTreeCore, + UnifiedTreeNode, +) + +# UnifiedTreeNode / UnifiedLRUList live on the tree core; re-exported here +# because other modules and tests import them from this module. from sglang.srt.mem_cache.unified_cache_components import ( - _NUM_COMPONENT_TYPES, BASE_COMPONENT_TYPE, CacheTransferPhase, - ComponentData, ComponentType, - EvictLayer, FullComponent, - LRURefreshPhase, MambaComponent, PrepareLoadBackResult, SWAComponent, TreeComponent, - get_and_increase_time_counter, -) -from sglang.srt.mem_cache.utils import ( - compute_node_hash_values, - get_eviction_strategy, - split_node_hash_value, ) from sglang.srt.observability.metrics_collector import ( STAT_LOGGER_ROLE_STORAGE, @@ -78,218 +80,28 @@ if TYPE_CHECKING: T = TypeVar("T") -class UnifiedTreeNode: - counter = 0 - - def __init__(self, tree_components: tuple[ComponentType, ...], priority: int = 0): - self.children = defaultdict(partial(UnifiedTreeNode, tree_components)) - self.parent: UnifiedTreeNode | None = None - self.key: Optional[RadixKey] = None - self.tree_components = tree_components - # list indexed by ComponentType (int enum 0..N-1) - self.component_data: list[ComponentData] = [ - ComponentData() for _ in range(_NUM_COMPONENT_TYPES) - ] - self.last_access_time = get_and_increase_time_counter() - self.creation_time = get_and_increase_time_counter() - self.hash_value = None - self.hit_count = 0 - self.priority = priority - self.lru_prev: list[UnifiedTreeNode | None] = [None] * ( - _NUM_COMPONENT_TYPES * 2 - ) - self.lru_next: list[UnifiedTreeNode | None] = [None] * ( - _NUM_COMPONENT_TYPES * 2 - ) - self.id = UnifiedTreeNode.counter - UnifiedTreeNode.counter += 1 - self.write_through_pending_id: Optional[int] = None - - def component(self, component_type: ComponentType) -> ComponentData: - return self.component_data[component_type] - - @property - def backuped(self) -> bool: - """Tree-level: Full KV present on host.""" - return self.component_data[ComponentType.FULL].host_value is not None - - @property - def evicted(self) -> bool: - """Tree-level: Full KV not on device (non-root with value=None).""" - return ( - self.parent is not None - and self.component_data[ComponentType.FULL].value is None - ) - - def __lt__(self, other: UnifiedTreeNode): - return self.last_access_time < other.last_access_time - - def get_last_hash_value(self) -> Optional[str]: - if self.hash_value is None or len(self.hash_value) == 0: - return None - return self.hash_value[-1] - - def get_prefix_hash_values(self, node: UnifiedTreeNode) -> list[str]: - if node is None or node.hash_value is None: - return [] - - return node.get_prefix_hash_values(node.parent) + node.hash_value - - -class UnifiedLRUList: - def __init__( - self, - component_type: ComponentType, - tree_components: tuple[ComponentType, ...], - use_host_ptr: bool = False, - ): - self.component_type = component_type - # Pointer slot: host LRU uses offset slots so device/host pointers - # never collide on the same node. - self._pt: int = component_type + (_NUM_COMPONENT_TYPES if use_host_ptr else 0) - self.head = UnifiedTreeNode(tree_components) - self.tail = UnifiedTreeNode(tree_components) - self.head.lru_next[self._pt] = self.tail - self.tail.lru_prev[self._pt] = self.head - self.cache: dict[int, UnifiedTreeNode] = {} - - def _add_node_after(self, prev_node: UnifiedTreeNode, new_node: UnifiedTreeNode): - pt = self._pt - new_node.lru_prev[pt] = prev_node - new_node.lru_next[pt] = prev_node.lru_next[pt] - prev_node.lru_next[pt].lru_prev[pt] = new_node - prev_node.lru_next[pt] = new_node - - def _add_node(self, node: UnifiedTreeNode): - self._add_node_after(self.head, node) - - def _remove_node(self, node: UnifiedTreeNode): - pt = self._pt - node.lru_prev[pt].lru_next[pt] = node.lru_next[pt] - node.lru_next[pt].lru_prev[pt] = node.lru_prev[pt] - # Clear self pointers to break reference cycles among evicted nodes. - node.lru_prev[pt] = None - node.lru_next[pt] = None - - def insert_mru(self, node: UnifiedTreeNode): - assert node.id not in self.cache - self.cache[node.id] = node - self._add_node(node) - - def remove_node(self, node: UnifiedTreeNode): - assert node.id in self.cache - del self.cache[node.id] - self._remove_node(node) - - def reset_node_mru(self, node: UnifiedTreeNode): - assert node.id in self.cache - self._remove_node(node) - self._add_node(node) - - def reset_node_and_parents_mru( - self, - node: UnifiedTreeNode, - root_node: UnifiedTreeNode, - should_include, - ): - prev_node = self.head - while node != root_node: - if should_include(node): - assert node.id in self.cache - self._remove_node(node) - self._add_node_after(prev_node, node) - prev_node = node - node = node.parent - - def reset_node_and_window_ancestors_mru( - self, - node: UnifiedTreeNode, - root_node: UnifiedTreeNode, - window_size: int, - should_include, - ): - prev_node = self.head - accumulated = 0 - while node != root_node and accumulated < window_size: - if should_include(node): - assert node.id in self.cache - self._remove_node(node) - self._add_node_after(prev_node, node) - prev_node = node - accumulated += len(node.key) - node = node.parent - - def in_list(self, node: Optional[UnifiedTreeNode]): - return node is not None and node.id in self.cache - - def get_prev_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): - if check_id: - assert node.id in self.cache - pt = self._pt - ct = self.component_type - x = node.lru_prev[pt] - while x.component_data[ct].lock_ref > 0: - x = x.lru_prev[pt] - if x == self.head: - return None - return x - - def get_prev_leaf_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): - if check_id: - assert node.id in self.cache - pt = self._pt - ct = self.component_type - x = node.lru_prev[pt] - while x.component_data[ct].lock_ref > 0 or len(x.children) > 0: - x = x.lru_prev[pt] - if x == self.head: - return None - return x - - def get_prev_no_host_lock(self, node: UnifiedTreeNode, check_id: bool = True): - """Host-LRU walker: skip nodes whose component host_lock_ref > 0.""" - if check_id: - assert node.id in self.cache - pt = self._pt - ct = self.component_type - x = node.lru_prev[pt] - while x.component_data[ct].host_lock_ref > 0: - x = x.lru_prev[pt] - if x == self.head: - return None - return x - - def get_lru_no_lock(self): - return self.get_prev_no_lock(self.tail, check_id=False) - - def get_leaf_lru_no_lock(self): - return self.get_prev_leaf_no_lock(self.tail, check_id=False) - - def get_lru_no_host_lock(self): - return self.get_prev_no_host_lock(self.tail, check_id=False) - - COMPONENT_REGISTRY: dict[ComponentType, type[TreeComponent]] = { ComponentType.FULL: FullComponent, ComponentType.MAMBA: MambaComponent, ComponentType.SWA: SWAComponent, } + logger = logging.getLogger(__name__) class _OngoingWriteThrough(NamedTuple): """Tracks an in-flight D→H write-through operation.""" - node: UnifiedTreeNode + node_id: NodeId lock_params: Optional[DecLockRefParams] - publish_nodes: list[UnifiedTreeNode] + publish_node_ids: list[NodeId] class _OngoingLoadBack(NamedTuple): """Tracks an in-flight H→D load-back operation.""" - node: UnifiedTreeNode + node_id: NodeId lock_params: DecLockRefParams host_lock_params: DecLockRefParams @@ -297,7 +109,7 @@ class _OngoingLoadBack(NamedTuple): class _OngoingPrefetch(NamedTuple): """Tracks an in-flight storage→host prefetch operation.""" - anchor_node: UnifiedTreeNode + anchor_node_id: NodeId prefetch_key: RadixKey host_indices: torch.Tensor operation: PrefetchOperation @@ -305,24 +117,14 @@ class _OngoingPrefetch(NamedTuple): comp_xfers: dict[ComponentType, list[PoolTransfer]] -class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): +class UnifiedRadixCache(BasePrefixCache): def __init__( self, params: CacheInitParams, ): self.req_to_token_pool = params.req_to_token_pool self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator - self.page_size = params.page_size self.disable = params.disable - self.enable_kv_cache_events = params.enable_kv_cache_events - self.kv_event_queue = [] - self.eviction_policy = params.eviction_policy.lower() - self.eviction_strategy = get_eviction_strategy(self.eviction_policy) - - if self.token_to_kv_pool_allocator: - self.device = self.token_to_kv_pool_allocator.device - else: - self.device = torch.device("cpu") if params.enable_metrics: self.init_metrics_collector() @@ -333,9 +135,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): assert params.tree_components is not None self.tree_components = tuple(params.tree_components) - self.is_eagle = ( - params.is_eagle and ComponentType.MAMBA not in self.tree_components - ) component_registry = COMPONENT_REGISTRY if params.component_registry_override: component_registry = { @@ -348,6 +147,29 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): self._components_tuple: tuple[TreeComponent, ...] = tuple( self.components.values() ) + # Whether SWA is enabled. + self.is_swa_enabled = ComponentType.SWA in params.tree_components + # Whether Mamba is enabled. + self.is_mamba_enabled = ComponentType.MAMBA in params.tree_components + # Whether the mamba extra (ping-pong) buffer is enabled. + self.enable_mamba_extra_buffer = ( + params.enable_mamba_extra_buffer if self.is_mamba_enabled else False + ) + # SWA window size (None when SWA is not enabled). + self._sliding_window_size = ( + params.sliding_window_size if self.is_swa_enabled else None + ) + # The TreeCore owns the tree member-var state (structure, LRUs, sizes, + # evictable leaves) and drives the components' tree-level hooks. + self.tree_core = create_tree_core( + name=envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get(), + params=params, + components=self.components, + ) + # Components execute boundary actions through the tree core. + for component in self.components.values(): + component.tree_core = self.tree_core + self.sidecar_pool_specs: list[SidecarPoolSpec] = [] # Streaming session: embedded StreamingSession with self as inner. @@ -373,7 +195,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): # HiCache D↔H defaults (overridden by init_hicache) self.cache_controller: Optional[HybridCacheController] = None self.host_pool_group = None # set by attach_hybrid_pool_to_unified_cache - self.write_through_threshold = 256 self.prefetch_stop_policy = "best_effort" self.prefetch_threshold = 256 self.prefetch_timeout_base = 1.0 @@ -454,50 +275,23 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): def _reset_full(self) -> None: """Full reset: destroy entire tree and all state.""" - self.root_node = UnifiedTreeNode(self.tree_components) - self.root_node.priority = -sys.maxsize - self.root_node.key = RadixKey(array("q"), None) - self.root_node.component_data[BASE_COMPONENT_TYPE].value = [] - self.root_node.hash_value = [] - for ct in self.tree_components: - self.root_node.component_data[ct].lock_ref = 1 - self.component_evictable_size_ = {ct: 0 for ct in self.tree_components} - self.component_protected_size_ = {ct: 0 for ct in self.tree_components} + self.tree_core.reset() - self.lru_lists = { - ct: UnifiedLRUList(ct, self.tree_components) for ct in self.tree_components - } + # Reset Controller. self.session.slots.clear() - - self.evictable_device_leaves: set[UnifiedTreeNode] = set() - self.evictable_host_leaves: set[UnifiedTreeNode] = set() - self.host_lru_lists = { - ct: UnifiedLRUList(ct, self.tree_components, use_host_ptr=True) - for ct in self.tree_components - } self.ongoing_write_through: dict[int, _OngoingWriteThrough] = {} self.ongoing_load_back: dict[int, _OngoingLoadBack] = {} self.enable_storage = False self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {} self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {} - self.ongoing_backup: dict[int, tuple[UnifiedTreeNode, DecLockRefParams]] = {} + self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {} if self.cache_controller is not None: self.cache_controller.reset() self.cache_controller.mem_pool_host.clear() self.enable_storage = self.cache_controller.enable_storage - self._empty_match_result = MatchResult( - device_indices=torch.empty( - (0,), - dtype=torch.int64, - device=self.device, - ), - last_device_node=self.root_node, - last_host_node=self.root_node, - best_match_node=self.root_node, - ) - self._record_all_cleared_event() + self.tree_core._record_all_cleared_event() def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None: """Initialize HiCache infrastructure.""" @@ -547,11 +341,21 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): storage_extra_config=storage_extra_config, storage_prefetch_threshold=storage_prefetch_threshold, ) + # Tag HiCache enablement on the TreeCore. + if self.cache_controller is not None: + self.tree_core.set_hicache_enabled() + if self.supports_swa(): + swa = self.components[ComponentType.SWA] + self.tree_core.has_swa_host_pool = swa._swa_kv_pool_host is not None # State initialization self.write_through_threshold = ( 1 if server_args.hicache_write_policy == "write_through" else 2 ) + self.is_write_back = ( + self.cache_controller is not None + and self.cache_controller.write_policy == "write_back" + ) self.load_back_threshold = 10 self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy @@ -578,53 +382,36 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): result = self.session.try_match_prefix(params) if result is not None: return result - - key = params.key - key, _ = key.maybe_to_bigram_view(self.is_eagle) - if self.disable or len(key) == 0: - return self._empty_match_result - key = key.page_aligned(self.page_size) - if len(key) == 0: - return self._empty_match_result - - ( - value, - best_match_node, - best_match_device_node, - best_match_device_value_len, - full_kv_hit_length, - ) = self._match_prefix_helper(key) - return self._match_post_processor( - params, - value, - best_match_node, - best_match_device_node, - best_match_device_value_len, - full_kv_hit_length, - ) + if self.disable: + return self.tree_core.empty_match_result + result = self.tree_core.match_prefix(params) + # Apply the walk's actions (e.g. a pending write-through relocation on + # a split) before the finalizers, which can evict or raise. + self._apply_cache_actions(result.cache_actions) + for component in self._components_tuple: + result = component.finalize_match_result_in_cache(params, result) + # Finalizers must not emit actions; the walk's were applied above. + assert not result.cache_actions + return result def insert(self, params: InsertParams) -> InsertResult: if self.disable: return InsertResult(prefix_len=0) - - key = params.key - value = params.value - key, value = key.maybe_to_bigram_view(self.is_eagle, value) - key = key.page_aligned(self.page_size) - if value is not None: - value = value[: len(key)] - else: - value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) - - result = self._insert_helper(self.root_node, key, value, params) - return result - - @property - def is_write_back(self) -> bool: - return ( - self.cache_controller is not None - and self.cache_controller.write_policy == "write_back" - ) + # Fail fast on re-entrancy without touching the in-flight walk. + assert not self.tree_core.has_ongoing_insert(), "re-entrant insert" + # Pump the resumable insert, applying each step's actions at its barrier. + try: + step = self.tree_core.begin_insert(params) + while True: + self._apply_cache_actions(step.actions) + if step.result is not None: + # Walk actions flow through the steps; the result is action-free. + assert not step.result.cache_actions + return step.result + step = self.tree_core.resume_insert() + finally: + # Drain still-pending actions so frees reach the allocator on abort. + self._apply_cache_actions(self.tree_core.end_insert()) def evict(self, params: EvictParams) -> EvictResult: if self.disable: @@ -632,8 +419,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): start_time = time.perf_counter() tracker = {ct: 0 for ct in self.tree_components} - for component in self._components_tuple: - component.drive_eviction(params=params, tracker=tracker) + request_by_type = { + ComponentType.FULL: params.num_tokens, + ComponentType.SWA: params.swa_num_tokens, + ComponentType.MAMBA: params.mamba_num, + } + self._evict_components(request_by_type, tracker) if ( self.cache_controller is not None @@ -648,83 +439,144 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): mamba_num_evicted=tracker.get(ComponentType.MAMBA, 0), ) - def inc_lock_ref(self, node: Any) -> IncLockRefResult: - result = self.session.try_inc_lock_ref(node) + def _free_values( + self, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Free a tree-side step's returned device and host values right away.""" + # Both drains must run even if one raises. + try: + self._drain_device_frees(device_frees) + finally: + self._drain_host_frees(host_frees) + + def _accumulate_tracker( + self, + tracker: dict[ComponentType, int], + delta: dict[ComponentType, int], + ) -> None: + """Fold a step result's evicted delta into the running totals.""" + for ct, n in delta.items(): + tracker[ct] += n + + def _evict_device_next_node( + self, component_type: ComponentType, tracker: dict[ComponentType, int] + ) -> Optional[NodeId]: + """Advance the eviction walk one node, consuming its step result.""" + result = self.tree_core.evict_device_next_node(component_type, tracker) + self._free_values(result.device_frees, result.host_frees) + self._accumulate_tracker(tracker, result.tracker) + return result.node_id + + def _evict_device_leaf( + self, node_id: NodeId, tracker: dict[ComponentType, int] + ) -> Optional[BackupKV]: + """Evict one device leaf, consuming its step result; returns the + deferred write-back BackupKV when one must run before the demote.""" + result = self.tree_core.evict_device_leaf(node_id, self.is_write_back) + self._free_values(result.device_frees, result.host_frees) + self._accumulate_tracker(tracker, result.tracker) + return result.backup_kv + + def _demote(self, node_id: NodeId, tracker: dict[ComponentType, int]) -> None: + """Demote a backed-up node, consuming its step result.""" + result = self.tree_core.demote(node_id) + self._free_values(result.device_frees, result.host_frees) + self._accumulate_tracker(tracker, result.tracker) + + def _drop_subtree_no_host( + self, node_id: NodeId, tracker: dict[ComponentType, int] + ) -> bool: + """Run the write-back drop fallback, consuming its step result.""" + result = self.tree_core.drop_subtree_no_host(node_id) + self._free_values(result.device_frees, result.host_frees) + self._accumulate_tracker(tracker, result.tracker) + return result.is_dropped + + def _evict_components( + self, + request_by_type: dict[ComponentType, int], + tracker: dict[ComponentType, int], + ) -> None: + for ct in self.tree_components: + request_cnt = request_by_type[ct] + # Skip eviction walk if request is already met + if tracker[ct] >= request_cnt: + continue + self.tree_core.evict_device_start(ct, request_cnt) + try: + while ( + node_id := self._evict_device_next_node(ct, tracker) + ) is not None: + backup_kv = self._evict_device_leaf(node_id, tracker) + if backup_kv is not None: + # Deferred demote: run the D->H backup, demote only on success. + written = self._execute_and_commit_kv_backup( + backup_kv, write_back=True + ) + if written > 0: + self.writing_check(write_back=True) + self._demote(node_id, tracker) + elif self._drop_subtree_no_host(node_id, tracker): + logger.warning( + "write_back: KV subtree dropped without backup " + "due to host memory pressure, root node %d", + node_id, + ) + else: + logger.warning( + "write_back: backup failed under host memory " + "pressure but subtree drop declined (node " + "locked); root node %d stays device-resident " + "until host space frees", + node_id, + ) + finally: + self.tree_core.evict_device_end(ct) + + def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult: + result = self.session.try_inc_lock_ref(node_id) if result is not None: return result if self.disable: return IncLockRefResult() - result = IncLockRefResult() - for component in self._components_tuple: - result = component.acquire_component_lock(node=node, result=result) - - self._update_evictable_leaf_sets(node) - return result + return self.tree_core.inc_lock_ref(node_id) def dec_lock_ref( self, - node: Any, + node_id: NodeId, params: Optional[DecLockRefParams] = None, skip_swa: bool = False, ) -> DecLockRefResult: - result = self.session.try_dec_lock_ref(node, params) + result = self.session.try_dec_lock_ref(node_id, params) if result is not None: return result if self.disable: return DecLockRefResult() - for component in self._components_tuple: - if skip_swa and component.component_type == ComponentType.SWA: - continue - component.release_component_lock(node=node, params=params) - - self._update_evictable_leaf_sets(node) - # TODO: delta is not aggregated from components; no caller uses it yet. - return DecLockRefResult() + return self.tree_core.dec_lock_ref(node_id, params, skip_swa) def dec_swa_lock_only( self, - node: UnifiedTreeNode, + node_id: NodeId, swa_uuid_for_lock: Optional[int] = None, ) -> None: - """Early-release the SWA portion of a request's tree lock, plus any - strictly-lower-priority locks (e.g. Mamba) co-located on `node`. - """ if self.disable: return - swa_component = self.components.get(ComponentType.SWA) - if swa_component is None: - return - swa_component.release_window_lock(node, swa_uuid_for_lock) + result = self.tree_core.dec_swa_lock_only(node_id, swa_uuid_for_lock) + self._free_values(result.device_frees, result.host_frees) - # Drop strictly-lower-priority locks (e.g. Mamba) co-located on `node`. - swa_priority = swa_component.eviction_priority(is_leaf=False) - dec_params = DecLockRefParams(swa_uuid_for_lock=swa_uuid_for_lock) - for comp in self._components_tuple: - if comp.eviction_priority(is_leaf=False) < swa_priority: - comp.release_component_lock(node, dec_params) - - def inc_host_lock_ref(self, node: Any) -> IncLockRefResult: + def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: if self.disable: return IncLockRefResult() - result = IncLockRefResult() - for component in self._components_tuple: - result = component.acquire_component_lock( - node=node, result=result, lock_host=True - ) - - self._update_evictable_leaf_sets(node) - return result + return self.tree_core.inc_host_lock_ref(node_id) def dec_host_lock_ref( - self, node: Any, params: Optional[DecLockRefParams] = None + self, node_id: NodeId, params: Optional[DecLockRefParams] = None ) -> DecLockRefResult: if self.disable: return DecLockRefResult() - for component in self._components_tuple: - component.release_component_lock(node=node, params=params, lock_host=True) - - self._update_evictable_leaf_sets(node) - return DecLockRefResult() + return self.tree_core.dec_host_lock_ref(node_id, params) def cache_finished_req( self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int, **kwargs @@ -775,7 +627,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): kv_indices = kv_indices[:effective_cache_len] radix_key = RadixKey( - token_ids, req.extra_key, is_bigram=self.is_eagle + token_ids, req.extra_key, is_bigram=self.tree_core.is_eagle ).page_aligned(self.page_size) page_aligned_len = len(radix_key) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) @@ -854,7 +706,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): radix_key = RadixKey( token_ids[:effective_cache_len], req.extra_key, - is_bigram=self.is_eagle, + is_bigram=self.tree_core.is_eagle, ).page_aligned(self.page_size) page_aligned_len = len(radix_key) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) @@ -909,845 +761,159 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): # ---- Internal Helpers ---- - def _match_prefix_helper( - self, key: RadixKey - ) -> tuple[list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, int, int]: - # Non-HiCache mode has only device-resident matches, so the scheduler - # device anchor follows the best match. In HiCache mode, host-backed - # nodes can also match, so we separately track the best device-resident - # match for scheduler prefix indices and locking. - node = self.root_node - child_key = key.child_key(self.page_size) - value: list[torch.Tensor] = [] - best_match_node = node - best_match_device_node = node - best_match_device_value_len = 0 - full_kv_hit_length = 0 - - separate_device_match = self.cache_controller is not None - if separate_device_match: - validators = tuple( - comp.create_match_validator() for comp in self._components_tuple - ) - device_validators = tuple( - comp.create_match_validator(match_device_only=True) - for comp in self._components_tuple - ) - else: - validators = tuple( - comp.create_match_validator(match_device_only=True) - for comp in self._components_tuple - ) - - def _all_valid(validators, node): - return all([v(node) for v in validators]) - - def _update_best_if_valid(node): - nonlocal best_match_node - nonlocal best_match_device_value_len, best_match_device_node - matched = _all_valid(validators, node) - if matched: - best_match_node = node - - if not separate_device_match: - if matched: - best_match_device_value_len = len(value) - best_match_device_node = node - return - if _all_valid(device_validators, node): - best_match_device_value_len = len(value) - best_match_device_node = node - - while len(key) > 0 and child_key in node.children: - child = node.children[child_key] - - # HiCache: dead node (evicted + not backuped) — stop traversal - if child.evicted and not child.backuped: - break - - prefix_len = child.key.match(key, page_size=self.page_size) - full_kv_hit_length += prefix_len - if prefix_len < len(child.key): - node = self._split_node(child.key, child, prefix_len) - if not node.evicted: - value.append(node.component_data[BASE_COMPONENT_TYPE].value) - _update_best_if_valid(node) - break - - if not child.evicted: - value.append(child.component_data[BASE_COMPONENT_TYPE].value) - node = child - _update_best_if_valid(node) - key = key[prefix_len:] - if len(key): - child_key = key.child_key(self.page_size) - - return ( - value, - best_match_node, - best_match_device_node, - best_match_device_value_len, - full_kv_hit_length, - ) - - def _match_post_processor( - self, - params: MatchPrefixParams, - value: list[torch.Tensor], - best_match_node: UnifiedTreeNode, - best_match_device_node: UnifiedTreeNode, - best_match_device_value_len: int, - full_kv_hit_length: int, - ) -> MatchResult: - node_update = best_match_node - for comp in self._components_tuple: - if comp.component_type == BASE_COMPONENT_TYPE: - continue # Full uses last_access_time, not LRU - comp.refresh_lru(LRURefreshPhase.MATCH_END, node_update, self.root_node) - - cur_time = get_and_increase_time_counter() - while node_update: - node_update.last_access_time = cur_time - cur_time -= 0.00001 - node_update = node_update.parent - - # last_host_node will be used as the starting node for the subsequent - # `prefetch_from_storage` flow. We directly use best_match_node here, - # because best_match_node represents the node where all components - # have reached consensus on both device & host availability. - last_host_node = ( - best_match_node - if self.cache_controller is not None - else best_match_device_node - ) - - if best_match_device_value_len > 0: - device_indices = torch.cat(value[:best_match_device_value_len]) - else: - device_indices = self._empty_match_result.device_indices - result = MatchResult( - device_indices=device_indices, - last_device_node=best_match_device_node, - last_host_node=last_host_node, - best_match_node=best_match_node, - host_hit_length=0, - full_kv_hit_length=full_kv_hit_length, - ) - - for component in self._components_tuple: - result = component.finalize_match_result( - result=result, - params=params, - value_chunks=value, - best_value_len=best_match_device_value_len, - ) - return result - - def _split_node( - self, key: RadixKey, child: UnifiedTreeNode, split_len: int - ) -> UnifiedTreeNode: - new_node = UnifiedTreeNode(self.tree_components, priority=child.priority) - new_node.children = {key[split_len:].child_key(self.page_size): child} - new_node.parent = child.parent - new_node.key = child.key[:split_len] - new_node.hit_count = child.hit_count - new_node.creation_time = child.creation_time - - self._for_each_component_lru(child, UnifiedLRUList.remove_node) - - child.parent = new_node - child.key = child.key[split_len:] - new_node.hash_value, child.hash_value = split_node_hash_value( - child.hash_value, split_len, self.page_size - ) - - for component in self._components_tuple: - component.redistribute_on_node_split(new_parent=new_node, child=child) - new_node.parent.children[key.child_key(self.page_size)] = new_node - - if child.backuped: - self._replace_pending_write_through_node(child, [new_node, child]) - - self._for_each_component_lru( - new_node, UnifiedLRUList.insert_mru, skip_existing=True - ) - self._for_each_component_lru( - child, UnifiedLRUList.insert_mru, skip_existing=True - ) - child.last_access_time = get_and_increase_time_counter() - - self._update_evictable_leaf_sets(new_node) - self._update_evictable_leaf_sets(child) - return new_node - - def _touch_node(self, node: UnifiedTreeNode): - node.last_access_time = get_and_increase_time_counter() - if node != self.root_node: - for comp in self._components_tuple: - if comp.component_type == BASE_COMPONENT_TYPE: - continue - comp.refresh_lru(LRURefreshPhase.WALKDOWN, node, self.root_node) - - def _add_new_node( - self, - parent: UnifiedTreeNode, - key: RadixKey, - value: torch.Tensor, - priority: int = 0, - ) -> UnifiedTreeNode: - new_node = UnifiedTreeNode(self.tree_components, priority=priority) - new_node.parent = parent - new_node.key = key - new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone() - parent.children[key.child_key(self.page_size)] = new_node - self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value) - if self.enable_storage: - new_node.hash_value = compute_node_hash_values(new_node, self.page_size) - - self._update_evictable_leaf_sets(new_node) - self._update_evictable_leaf_sets(parent) - self._record_store_event(new_node) - return new_node - - def _unevict_node_on_insert( - self, node: UnifiedTreeNode, fresh_value: torch.Tensor + def _apply_cache_actions( + self, actions: list[CacheAction | ComponentAction] ) -> None: - """Restore an evicted node's Full device value from fresh KV indices - during insert.""" - ct = BASE_COMPONENT_TYPE - cd = node.component_data[ct] - assert cd.value is None - n = len(fresh_value) - cd.value = fresh_value.clone() - self.component_evictable_size_[ct] += n - self._update_evictable_leaf_sets(node) - if node.parent is not None: - self._update_evictable_leaf_sets(node.parent) - self._record_store_event(node, medium=StorageMedium.GPU) + # Apply and consume one at a time: a spent list cannot be double-applied. + actions.reverse() + try: + while actions: + self._apply_cache_action(actions.pop()) + finally: + actions.reverse() - def _insert_helper( - self, - node: UnifiedTreeNode, - key: RadixKey, - value: torch.Tensor, - params: InsertParams, - ) -> InsertResult: - priority = params.priority - if priority is None: - priority = 0 - self._touch_node(node) - node.priority = max(node.priority, priority) - if len(key) == 0: - return InsertResult(prefix_len=0, mamba_exist=True) - - child_key = key.child_key(self.page_size) - total_prefix_length = 0 - while len(key) > 0 and child_key in node.children: - node = node.children[child_key] - self._touch_node(node) - prefix_len = node.key.match(key, page_size=self.page_size) - if prefix_len < len(node.key): - node = self._split_node(node.key, node, prefix_len) - node.priority = max(node.priority, priority) - - if node.evicted: - self._unevict_node_on_insert(node, value[:prefix_len]) - # FULL was restored from the request's fresh KV. Aux - # components (e.g. SWA) may still hold tombstones and need - # to rebuild their value from the same slice. - for component in self._components_tuple: - if component.component_type == BASE_COMPONENT_TYPE: - continue - component.recover_after_unevict( - node=node, - prefix_len=prefix_len, - total_prefix_len=total_prefix_length, - params=params, - ) - else: - value_slice = value[:prefix_len] - consumed_from = prefix_len - # Let each component claim ownership of overlapping KV slots - for component in self._components_tuple: - comp_consumed_from = component.update_component_on_insert_overlap( - node=node, - prefix_len=prefix_len, - total_prefix_len=total_prefix_length, - value_slice=value_slice, - params=params, - ) - consumed_from = min(consumed_from, comp_consumed_from) - - dup_start = max(0, params.prev_prefix_len - total_prefix_length) - if dup_start < consumed_from: - self.token_to_kv_pool_allocator.free( - value_slice[dup_start:consumed_from] - ) - - self._inc_hit_count(node, params.chunked) - total_prefix_length += prefix_len - key = key[prefix_len:] - value = value[prefix_len:] - if len(key): - child_key = key.child_key(self.page_size) - - is_new_leaf = False - # Create new leaf for remaining suffix. A leaf survives on its Full - # value alone; auxiliary components (SWA, Mamba) may legitimately hold - # only a tombstone for this span (e.g. the whole leaf is outside the SWA - # window). Materialize it anyway so the Full KV stays cacheable. - if len(key): - target_node = self._add_new_node(node, key, value, priority=priority) - is_new_leaf = True + def _apply_cache_action(self, action: CacheAction | ComponentAction) -> None: + # Component actions route to their component class; the rest are + # cache-owned and handled here by type. + if isinstance(action, ComponentAction): + self.components[action.component_type].apply_component_action(action) + elif isinstance(action, ReplaceWriteThroughOnNodeSplit): + self._replace_pending_write_through_node( + action.ack_id, + action.old_node_id, + [action.new_node_id, action.new_child_node_id], + ) + elif isinstance(action, FreeDeviceKV): + for indices in action.indices: + self.token_to_kv_pool_allocator.free(indices) + elif isinstance(action, BackupKV): + self._execute_and_commit_kv_backup(action) else: - target_node = node + raise AssertionError(f"unhandled CacheAction: {type(action).__name__}") - # Finalize: let each component attach its data to the target node. - # e.g. Mamba attaches mamba_value to the leaf node - result = InsertResult(prefix_len=total_prefix_length) - for component in self._components_tuple: - component.commit_insert_component_data( - node=target_node, - is_new_leaf=is_new_leaf, - params=params, - result=result, + def _drain_device_frees( + self, device_frees: dict[ComponentType, list[torch.Tensor]] + ) -> None: + # Free per component device slots, consuming each entry as it frees. + for ct in list(device_frees): + self._apply_cache_action( + FreeComponentDeviceSlot(device_frees.pop(ct), component_type=ct) ) - if target_node is not self.root_node: - for component in self._components_tuple: - if component.component_type == BASE_COMPONENT_TYPE: - continue - component.refresh_lru( - LRURefreshPhase.INSERT_END, target_node, self.root_node - ) - - if is_new_leaf: - self._inc_hit_count(target_node, params.chunked) - return result - - def _insert_helper_host( - self, - node: UnifiedTreeNode, - key: RadixKey, - host_value: torch.Tensor, - hash_value: list[str], - ) -> InsertResult: - total_len = len(key) - self._touch_node(node) - if total_len == 0: - return InsertResult(prefix_len=0, mamba_exist=True) - - child_key = key.child_key(self.page_size) - matched_length = 0 - while len(key) > 0 and child_key in node.children: - node = node.children[child_key] - self._touch_node(node) - prefix_len = node.key.match(key, page_size=self.page_size) - - key = key[prefix_len:] - host_value = host_value[prefix_len:] - hash_value = hash_value[prefix_len // self.page_size :] - matched_length += prefix_len - - if prefix_len < len(node.key): - node = self._split_node(node.key, node, prefix_len) - - if len(key): - child_key = key.child_key(self.page_size) - - result = InsertResult(prefix_len=matched_length, total_len=total_len) - if len(key) == 0: - if ( - node is not self.root_node - and node.component_data[BASE_COMPONENT_TYPE].host_value is not None - ): - result.inserted_host_node = node - return result - - new_node = UnifiedTreeNode(self.tree_components, priority=node.priority) - new_node.parent = node - new_node.key = key - new_node.hash_value = hash_value - new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone() - node.children[child_key] = new_node - self._update_evictable_leaf_sets(new_node) - self._update_evictable_leaf_sets(node) - result.inserted_host_node = new_node - return result - - # ---- Evict Helpers ---- - - def _cascade_evict( - self, - node: UnifiedTreeNode, - trigger: TreeComponent, - tracker: dict[ComponentType, int], - target: EvictLayer = EvictLayer.DEVICE, - ): - """Cascade eviction from trigger to lower-or-equal priority components.""" - - is_leaf = False - if target == EvictLayer.DEVICE: - is_leaf = node in self.evictable_device_leaves - elif target == EvictLayer.HOST: - is_leaf = node in self.evictable_host_leaves - - trigger_priority = trigger.eviction_priority(is_leaf) - - for comp in self._components_tuple: - if comp.eviction_priority(is_leaf) <= trigger_priority: - if comp is not trigger and comp.node_has_component_data(node, target): - cd = node.component_data[comp.component_type] - # A comp whose TRUE internal priority outranks the trigger - # is only in this loop because leaf-collapse flattened - # priorities; a lock on it is a legit pin and must be - # spared. A lock on a strictly-lower-priority tier is a - # real strand — fall through to the assert below. - if comp.eviction_priority( - is_leaf=False - ) >= trigger.eviction_priority(is_leaf=False): - if EvictLayer.DEVICE in target and cd.lock_ref != 0: - continue - if EvictLayer.HOST in target and cd.host_lock_ref != 0: - continue - if EvictLayer.DEVICE in target: - assert cd.lock_ref == 0 - if EvictLayer.HOST in target: - assert cd.host_lock_ref == 0 - self._evict_component_and_detach_lru( - node, comp, target=target, tracker=tracker - ) - - # Now that all components (including SWA which depends on Full.value) - # have been freed, we can safely tombstone Full.value. - # This is deferred from evict_component because free_swa needs it. - if ( - target is EvictLayer.DEVICE - and trigger.component_type == BASE_COMPONENT_TYPE - ): - node.component_data[trigger.component_type].value = None - - self._update_evictable_leaf_sets(node) - - def _remove_leaf_from_parent(self, node: UnifiedTreeNode): - key = node.key.child_key(self.page_size) - v = node.parent.children.pop(key, None) - assert v == node - - def _evict_component_and_detach_lru( - self, - node: UnifiedTreeNode, - comp: TreeComponent, - target: EvictLayer = EvictLayer.DEVICE, - tracker: Optional[dict[ComponentType, int]] = None, - ) -> tuple[int, int]: - device_freed, host_freed = comp.evict_component(node, target=target) - if tracker is not None: - if EvictLayer.DEVICE in target: - tracker[comp.component_type] += device_freed - elif EvictLayer.HOST in target: - tracker[comp.component_type] += host_freed - - # Detach from the appropriate LRU list(s) - ct = comp.component_type - for layer, lru_lists in ( - (EvictLayer.DEVICE, self.lru_lists), - (EvictLayer.HOST, self.host_lru_lists), - ): - if layer in target: - lru = lru_lists[ct] - if lru.in_list(node): - lru.remove_node(node) - return device_freed, host_freed - - def _iteratively_delete_tombstone_leaf( - self, deleted_node: UnifiedTreeNode, tracker: dict[ComponentType, int] - ): - """Walk up from *deleted_node* and cascade-delete childless ancestors. - - Only the Full (base) component decides whether a node survives: - - Full device present → keep as D-leaf - - Full host present → keep as H-leaf - - neither → evict all remaining data, delete, continue up - """ - ct = BASE_COMPONENT_TYPE - cur = deleted_node.parent - while cur != self.root_node and len(cur.children) == 0: - if any( - cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data - ): - break - - has_device = cur.component_data[ct].value is not None - has_host = cur.component_data[ct].host_value is not None - - if has_device: - self._update_evictable_leaf_sets(cur) - break - - # Full device absent — clean up orphaned aux device data. - for comp in self.components.values(): - if comp.node_has_component_data(cur): - self._evict_component_and_detach_lru( - cur, comp, target=EvictLayer.DEVICE, tracker=tracker - ) - - if has_host: - self._update_evictable_leaf_sets(cur) - break - - # Full absent on both layers — evict remaining host data, delete. - for comp in self.components.values(): - if comp.node_has_component_data(cur, target=EvictLayer.HOST): - self._evict_component_and_detach_lru( - cur, comp, target=EvictLayer.HOST, tracker=tracker - ) - - self.evictable_host_leaves.discard(cur) - self._remove_leaf_from_parent(cur) - parent = cur.parent - self._update_evictable_leaf_sets(parent) - cur = parent - - def _for_each_component_lru( - self, - node: UnifiedTreeNode, - lru_op, - target: EvictLayer = EvictLayer.DEVICE, - skip_existing: bool = False, - ): - """Apply lru_op to each aux component's LRU that has data on this node. - If skip_existing=True, skip components already in the target LRU list.""" - lru_dict = self.host_lru_lists if target is EvictLayer.HOST else self.lru_lists - for ct in self.tree_components: - if ct == BASE_COMPONENT_TYPE: - continue # Full uses leaf sets, not LRU - cd = node.component_data[ct] - if (cd.host_value if target is EvictLayer.HOST else cd.value) is not None: - lru = lru_dict[ct] - if skip_existing and lru.in_list(node): - continue - lru_op(lru, node) + def _drain_host_frees( + self, host_frees: dict[ComponentType, list[torch.Tensor]] + ) -> None: + # Free per component host-pool slots, consuming each entry as it frees. + for ct in list(host_frees): + self.components[ct].free_host_values(host_frees.pop(ct)) def evict_host( self, num_tokens: int, component_type: ComponentType = BASE_COMPONENT_TYPE ) -> int: """Evict host resources for a specific component to free host pool space.""" - tracker: dict[ComponentType, int] = {ct: 0 for ct in self.tree_components} - comp = self.components.get(component_type) - if comp is not None: - comp.drive_host_eviction(num_tokens, tracker) - return tracker[component_type] - - def _is_device_leaf(self, node: UnifiedTreeNode) -> bool: - """D-leaf: Full device value present, no child with Full KV on device, - unlocked, not root. - - Only the Full (base) component is required; auxiliary components - (Mamba, SWA) are not mandatory for D-leaf membership.""" - ct = BASE_COMPONENT_TYPE - if node is self.root_node or node.evicted: - return False - if any(cd.lock_ref > 0 for cd in node.component_data): - return False - if any( - child.component_data[ct].value is not None - for child in node.children.values() - ): - return False - return True - - def _is_host_leaf(self, node: UnifiedTreeNode) -> bool: - """H-leaf: evicted, Full host value present, no children, unlocked, not root. - - Only the Full (base) component host_value is required; auxiliary - components are not mandatory for H-leaf membership.""" - if node is self.root_node or not node.evicted: - return False - if not node.backuped: - return False - if any(cd.host_lock_ref > 0 for cd in node.component_data): - return False - if len(node.children) > 0: - return False - return True - - def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: - """Update both device and host leaf sets for a node.""" - if self._is_device_leaf(node): - self.evictable_device_leaves.add(node) - else: - self.evictable_device_leaves.discard(node) - - if self._is_host_leaf(node): - self.evictable_host_leaves.add(node) - else: - self.evictable_host_leaves.discard(node) - - def _evict_to_host( - self, node: UnifiedTreeNode, tracker: Optional[dict[ComponentType, int]] = None - ) -> None: - """GPU→CPU demotion: release all device resources, node stays in tree.""" - assert not node.evicted and node.backuped - trigger = self.components[BASE_COMPONENT_TYPE] - self._evict_component_and_detach_lru( - node, trigger, target=EvictLayer.DEVICE, tracker=tracker - ) - self._cascade_evict(node, trigger, tracker) - self._record_remove_event(node, medium=StorageMedium.GPU) - - # after device eviction, insert aux components into host LRU. - self._for_each_component_lru( - node, UnifiedLRUList.insert_mru, target=EvictLayer.HOST, skip_existing=True - ) - self._update_evictable_leaf_sets(node.parent) - - def _evict_device_leaf( - self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] - ) -> None: - """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: Cascade evict all components - - All freed device tokens are accumulated into *tracker*. - """ - assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf" - if not node.backuped: - if ( - self.cache_controller is not None - and self.cache_controller.write_policy == "write_back" - ): - written = self.write_backup(node, write_back=True) - if written == 0: - if self._drop_subtree_no_host(node, tracker): - logger.warning( - "write_back: KV subtree dropped without backup " - "due to host memory pressure, root node %d", - node.id, - ) - else: - logger.warning( - "write_back: backup failed under host memory " - "pressure but subtree drop declined (node " - "locked); root node %d stays device-resident " - "until host space frees", - node.id, - ) - return - self.writing_check(write_back=True) - self._evict_to_host(node, tracker) - return - else: - # Write-through: node has no backup, delete entirely. - self._delete_unbacked_device_leaf(node, tracker) - return - self._evict_to_host(node, tracker) - - def _drop_subtree_no_host( - self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] - ) -> bool: - """Write-back fallback when a D-leaf's D->H backup fails under host - memory pressure: drop the subtree rooted at the unbacked leaf so - device eviction keeps making progress instead of leaving its KV - unevictable until host space frees up.""" - - assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf" - # A failed backup never issues the D->H copy, so the subtree root has - # no host state and no in-flight DMA reading its device slots. - assert not node.backuped and node.write_through_pending_id is None - if any(cd.host_lock_ref > 0 for cd in node.component_data): - return False - descendants: list[UnifiedTreeNode] = [] - stack = list(node.children.values()) - while stack: - cur = stack.pop() - if any( - cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data - ): - return False - descendants.append(cur) - stack.extend(cur.children.values()) - for desc in reversed(descendants): - # Host-only by construction: a device descendant would contradict - # this node being a D-leaf, and D-leaves evict before ancestors. - assert desc.evicted and desc.backuped, f"node {desc.id} not host-only" - assert desc.write_through_pending_id is None - self._release_all_component_layers(desc, StorageMedium.CPU, tracker) - self._remove_leaf_from_parent(desc) - self._delete_unbacked_device_leaf(node, tracker) - return True - - def _release_all_component_layers( - self, - node: UnifiedTreeNode, - medium: StorageMedium, - tracker: dict[ComponentType, int], - ) -> None: - """Free every component layer on the node and detach it from the LRU - lists and evictable leaf sets.""" - self._record_remove_event(node, medium=medium) - for comp in self._components_tuple: - self._evict_component_and_detach_lru( - node, comp, target=EvictLayer.ALL, tracker=tracker - ) - self.evictable_device_leaves.discard(node) - self.evictable_host_leaves.discard(node) - - def _delete_unbacked_device_leaf( - self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] - ) -> None: - """Delete a device leaf that has no host backup, freeing all layers.""" - self._release_all_component_layers(node, StorageMedium.GPU, tracker) - parent = node.parent - self._remove_leaf_from_parent(node) - self._update_evictable_leaf_sets(parent) - self._iteratively_delete_tombstone_leaf(node, tracker) - - def _evict_host_leaf( - self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] - ) -> None: - """Atomically evict all components on a host leaf. - - All freed tokens are accumulated into *tracker*.""" - assert self._is_host_leaf(node), f"node {node.id} is not an H-leaf" - - self._record_remove_event(node, medium=StorageMedium.CPU) - for comp in self._components_tuple: - _, hf = self._evict_component_and_detach_lru( - node, comp, target=EvictLayer.ALL, tracker=None - ) - tracker[comp.component_type] += hf - self.evictable_host_leaves.discard(node) - self._remove_leaf_from_parent(node) - self._iteratively_delete_tombstone_leaf(node, tracker) + result = self.tree_core.drive_host_eviction(component_type, num_tokens) + self._free_values(result.device_frees, result.host_frees) + return result.tracker.get(component_type, 0) # ---- HiCache: Backup / LoadBack ---- - def write_backup(self, node: UnifiedTreeNode, write_back: bool = False) -> int: - """Backup a node's data from device to host (D->H).""" - if self.cache_controller is None: - return 0 - - # Backup invariant (write-through): parent must be backuped first - if not write_back and ( - node.parent is not self.root_node and not node.parent.backuped - ): - if self.write_backup(node.parent) <= 0: - return 0 - - device_value = node.component_data[BASE_COMPONENT_TYPE].value - kv_xfer = PoolTransfer(name=PoolName.KV, device_indices=device_value) - - # Build aux transfers, keyed per component. - comp_xfers: dict[ComponentType, list] = {} - for comp in self._components_tuple: - if comp.component_type == BASE_COMPONENT_TYPE: + def _execute_and_commit_kv_backup( + self, action: BackupKV, write_back: bool = False + ) -> int: + """Run a backup action top-down, stopping at the first failed backup.""" + written = 0 + for node_id in action.node_ids: + # Overlapping chain actions: skip already-backed nodes. + if self.tree_core.is_backuped(node_id): continue - t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST) - if t: - comp_xfers[comp.component_type] = t - sidecar_xfers = self._build_sidecar_transfers( + device_value, comp_xfers = self.tree_core.build_backup_spec(node_id) + sidecar_xfers = self._build_backup_sidecar(device_value, comp_xfers) + host_indices = self._execute_kv_backup( + node_id, device_value, comp_xfers, sidecar_xfers + ) + if host_indices is None: + return 0 + self.tree_core.commit_backup(node_id, host_indices, comp_xfers) + lock_params = None + if not write_back: + lock_params = self.inc_lock_ref(node_id).to_dec_params() + self._track_write_through_node(node_id, lock_params) + written = len(host_indices) + return written + + def _build_backup_sidecar(self, device_value, comp_xfers): + """Gather sidecar transfer spec.""" + kv_xfer = PoolTransfer(name=PoolName.KV, device_indices=device_value) + return self._build_sidecar_transfers( CacheTransferPhase.BACKUP_HOST, kv_xfer, comp_xfers ) - # Pre-evict host if insufficient + def _execute_kv_backup(self, node_id, device_value, comp_xfers, sidecar_xfers): + """Execute Backup action.""" kv_tokens = len(device_value) host_avail = self.cache_controller.mem_pool_host.available_size() if host_avail < kv_tokens: needed = kv_tokens - host_avail - evicted = self.evict_host(needed) - if evicted < needed: - return 0 - + if self.evict_host(needed) < needed: + return None aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] aux_xfers.extend(sidecar_xfers) - host_indices = self.cache_controller.write( - device_value, node_id=node.id, extra_pools=aux_xfers or None + return self.cache_controller.write( + device_value, node_id=node_id, extra_pools=aux_xfers or None ) - if host_indices is None: - return 0 - - # Commit - kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) - self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer( - node, - CacheTransferPhase.BACKUP_HOST, - transfers=[kv_xfer], - ) - for ct, xfers in comp_xfers.items(): - self.components[ct].commit_hicache_transfer( - node, - CacheTransferPhase.BACKUP_HOST, - transfers=xfers, - ) - - lock_params = None - if not write_back: - lock_params = self.inc_lock_ref(node).to_dec_params() - self._track_write_through_node(node, lock_params) - return len(host_indices) def _track_write_through_node( self, - node: UnifiedTreeNode, + node_id: NodeId, lock_params: Optional[DecLockRefParams], ) -> None: - node.write_through_pending_id = node.id - self.ongoing_write_through[node.id] = _OngoingWriteThrough( - node, lock_params, [node] + self.tree_core.mark_write_through_pending(node_id) + self.ongoing_write_through[node_id] = _OngoingWriteThrough( + node_id, lock_params, [node_id] ) def _replace_pending_write_through_node( - self, old_node: UnifiedTreeNode, new_nodes: list[UnifiedTreeNode] + self, ack_id: int, old_node_id: NodeId, new_node_ids: list[NodeId] ) -> None: - ack_id = old_node.write_through_pending_id - if ack_id is None: - return - pending = self.ongoing_write_through.get(ack_id) if pending is None: return - lock_node, lock_params, publish_nodes = pending - updated_nodes = [] + lock_node_id, lock_params, publish_node_ids = pending + updated_node_ids = [] replaced = False - for node in publish_nodes: - if node is old_node: - updated_nodes.extend(new_nodes) + for node_id in publish_node_ids: + if node_id == old_node_id: + updated_node_ids.extend(new_node_ids) replaced = True else: - updated_nodes.append(node) + updated_node_ids.append(node_id) if not replaced: return - for node in new_nodes: - node.write_through_pending_id = ack_id self.ongoing_write_through[ack_id] = _OngoingWriteThrough( - lock_node, + lock_node_id, lock_params, - updated_nodes, + updated_node_ids, ) def _finish_write_through_ack(self, ack_id: int) -> None: - lock_node, lock_params, publish_nodes = self.ongoing_write_through.pop(ack_id) - for node in publish_nodes: - if node.write_through_pending_id == ack_id: - node.write_through_pending_id = None - self._record_store_event(node, medium=StorageMedium.CPU) + lock_node_id, lock_params, publish_node_ids = self.ongoing_write_through.pop( + ack_id + ) + self.tree_core.finish_write_through(publish_node_ids, ack_id) if lock_params is not None: - self.dec_lock_ref(lock_node, lock_params) + self.dec_lock_ref(lock_node_id, lock_params) if self.enable_storage: # Back up each fragment: after a split, lock_node only holds the # suffix; the prefix fragment must be persisted as well. - for node in publish_nodes: - self.write_backup_storage(node) + for node_id in publish_node_ids: + self.write_backup_storage(node_id) def load_back( self, - best_match_node: UnifiedTreeNode, + node_id: NodeId, mem_quota: Optional[int] = None, req=None, ) -> bool: @@ -1755,29 +921,24 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if self.cache_controller is None: return False - host_anchor_params = self.inc_host_lock_ref(best_match_node).to_dec_params() - # Build KV transfer - kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers( - best_match_node, CacheTransferPhase.LOAD_BACK - )[0] + host_anchor_params = self.inc_host_lock_ref(node_id).to_dec_params() - # Lock path & pre-evict if device pool is insufficient - result = self.inc_lock_ref(best_match_node) + # Lock the path before building transfers (the aux build can evict). + result = self.inc_lock_ref(node_id) ancestor_lock_params = result.to_dec_params() # Let each component pre-allocate per-request state for the load-back; # the finally below lets components recover it unless the load succeeds. preps: dict[ComponentType, PrepareLoadBackResult] = { - comp.component_type: comp.prepare_load_back(best_match_node, req=req) + comp.component_type: comp.prepare_load_back(node_id, req=req) for comp in self._components_tuple } success = False try: success = self._load_back_transfers( - best_match_node=best_match_node, + node_id=node_id, mem_quota=mem_quota, req=req, - kv_xfer=kv_xfer, result=result, ancestor_lock_params=ancestor_lock_params, host_anchor_params=host_anchor_params, @@ -1790,25 +951,16 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): def _load_back_transfers( self, *, - best_match_node: UnifiedTreeNode, + node_id: NodeId, mem_quota: Optional[int], req, - kv_xfer: PoolTransfer, result: IncLockRefResult, - ancestor_lock_params: Optional[DecLockRefParams], - host_anchor_params: Optional[DecLockRefParams], + ancestor_lock_params: DecLockRefParams, + host_anchor_params: DecLockRefParams, ) -> bool: + # Build the KV + per-component aux transfers. + kv_xfer, comp_xfers = self.tree_core.build_load_back_spec(node_id, req=req) kv_tokens = len(kv_xfer.host_indices) - # Build aux transfers, keyed per component. - comp_xfers: dict[ComponentType, list] = {} - for comp in self._components_tuple: - if comp.component_type == BASE_COMPONENT_TYPE: - continue - t = comp.build_hicache_transfers( - best_match_node, CacheTransferPhase.LOAD_BACK, req=req - ) - if t: - comp_xfers[comp.component_type] = t sidecar_xfers = self._build_sidecar_transfers( CacheTransferPhase.LOAD_BACK, kv_xfer, comp_xfers ) @@ -1819,8 +971,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if (kv_tokens < self.load_back_threshold and not comp_xfers) or ( mem_quota is not None and kv_tokens > mem_quota + result.delta ): - self.dec_lock_ref(best_match_node, ancestor_lock_params) - self.dec_host_lock_ref(best_match_node, host_anchor_params) + self.dec_lock_ref(node_id, ancestor_lock_params) + self.dec_host_lock_ref(node_id, host_anchor_params) return False if self.supports_swa(): @@ -1831,8 +983,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): needed = kv_tokens - avail result = self.evict(EvictParams(num_tokens=needed)) if result.num_tokens_evicted < needed: - self.dec_lock_ref(best_match_node, ancestor_lock_params) - self.dec_host_lock_ref(best_match_node, host_anchor_params) + self.dec_lock_ref(node_id, ancestor_lock_params) + self.dec_host_lock_ref(node_id, host_anchor_params) return False # Load H→D @@ -1840,35 +992,25 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): aux_xfers.extend(sidecar_xfers) device_indices = self.cache_controller.load( host_indices=kv_xfer.host_indices, - node_id=best_match_node.id, + node_id=node_id, extra_pools=aux_xfers or None, ) - self.dec_lock_ref(best_match_node, ancestor_lock_params) + self.dec_lock_ref(node_id, ancestor_lock_params) if device_indices is None: - self.dec_host_lock_ref(best_match_node, host_anchor_params) + self.dec_host_lock_ref(node_id, host_anchor_params) return False - # Commit: each component gets only its own transfers - kv_xfer.device_indices = device_indices - self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer( - best_match_node, - CacheTransferPhase.LOAD_BACK, - [kv_xfer], - ) - for node in kv_xfer.nodes_to_load or (): - self._record_store_event(node, medium=StorageMedium.GPU) - for ct, xfers in comp_xfers.items(): - self.components[ct].commit_hicache_transfer( - best_match_node, - CacheTransferPhase.LOAD_BACK, - xfers, + # Commit the loaded KV back onto the node + apply its emitted actions. + self._apply_cache_actions( + self.tree_core.commit_load_back( + node_id, device_indices, kv_xfer, comp_xfers ) + ) - self._update_evictable_leaf_sets(best_match_node) - self.ongoing_load_back[best_match_node.id] = _OngoingLoadBack( - best_match_node, - self.inc_lock_ref(best_match_node).to_dec_params(), + self.ongoing_load_back[node_id] = _OngoingLoadBack( + node_id, + self.inc_lock_ref(node_id).to_dec_params(), host_anchor_params, ) @@ -1925,73 +1067,54 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): ) return transfers - def _inc_hit_count(self, node: UnifiedTreeNode, chunked: bool = False) -> None: - """Increment hit count; trigger write_backup when threshold reached.""" - if node.evicted or chunked: + def write_backup_storage(self, node_id: NodeId) -> None: + if not self.enable_storage or self.cache_controller is None: return - if ( - self.cache_controller is not None - and self.cache_controller.write_policy == "write_back" - ): + spec = self.tree_core.build_storage_backup_spec( + node_id, self.hicache_storage_pass_prefix_keys + ) + if spec is None: return - node.hit_count += 1 - if ( - self.cache_controller is not None - and not node.backuped - and node.hit_count >= self.write_through_threshold - ): - self.write_backup(node) - - def write_backup_storage(self, node: UnifiedTreeNode) -> None: - if ( - not self.enable_storage - or self.cache_controller is None - or not node.backuped - ): - return - - prefix_keys = None - if self.hicache_storage_pass_prefix_keys: - prefix_keys = node.get_prefix_hash_values(node.parent) - - comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} - for comp in self._components_tuple: - if comp.component_type == BASE_COMPONENT_TYPE: - continue - transfers = comp.build_hicache_transfers( - node, - CacheTransferPhase.BACKUP_STORAGE, - ) - if transfers: - comp_xfers[comp.component_type] = transfers kv_xfer = PoolTransfer( name=PoolName.KV, - host_indices=node.component_data[BASE_COMPONENT_TYPE].host_value, - keys=node.hash_value, + host_indices=spec.host_value, + keys=spec.hash_value, ) sidecar_xfers = self._build_sidecar_transfers( - CacheTransferPhase.BACKUP_STORAGE, kv_xfer, comp_xfers + CacheTransferPhase.BACKUP_STORAGE, kv_xfer, spec.comp_xfers ) - aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + aux_xfers = [x for xfers in spec.comp_xfers.values() for x in xfers] aux_xfers.extend(sidecar_xfers) operation_id = self.cache_controller.write_storage( - node.component_data[BASE_COMPONENT_TYPE].host_value, - node.key.token_ids, - node.hash_value, - prefix_keys, + spec.host_value, + spec.token_ids, + spec.hash_value, + spec.prefix_keys, extra_pools=aux_xfers or None, ) self.ongoing_backup[operation_id] = ( - node, - self.inc_host_lock_ref(node).to_dec_params(), + node_id, + self.inc_host_lock_ref(node_id).to_dec_params(), ) + def is_backuped(self, node_id: NodeId) -> bool: + return self.tree_core.is_backuped(node_id) + + def is_root(self, node_id: NodeId) -> bool: + return self.tree_core.is_root(node_id) + + def get_last_hash_value(self, node_id: NodeId) -> Optional[str]: + return self.tree_core.get_last_hash_value(node_id) + + def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: + return self.tree_core.get_prefix_hash_values(node_id) + def prefetch_from_storage( self, req_id: str, - last_host_node: UnifiedTreeNode, + last_host_node_id: NodeId, new_input_tokens: list[int], last_hash: Optional[str] = None, prefix_keys: Optional[list[str]] = None, @@ -1999,11 +1122,11 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if not self.enable_storage or self.cache_controller is None: return - extra_key = last_host_node.key.extra_key if last_host_node.key else None + extra_key = self.tree_core.prefetch_anchor_info(last_host_node_id) prefetch_key = RadixKey( new_input_tokens, extra_key=extra_key, - is_bigram=self.is_eagle, + is_bigram=self.tree_core.is_eagle, ).page_aligned(self.page_size) prefetch_length = len(prefetch_key) if ( @@ -2012,24 +1135,32 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): ): return - anchor_lock_params = self.inc_host_lock_ref(last_host_node).to_dec_params() + anchor_lock_params = self.inc_host_lock_ref(last_host_node_id).to_dec_params() comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} alloc_failed = False - for comp in self._components_tuple: - if comp.component_type == BASE_COMPONENT_TYPE: + for ct in self.tree_components: + if ct == BASE_COMPONENT_TYPE: continue - transfers = comp.build_hicache_transfers( - last_host_node, + # Pre-allocate the component's prefetch host buffer so the build stays pure. + prep = self.components[ct].prepare_prefetch( + last_host_node_id, prefetch_tokens=len(prefetch_key) + ) + if prep.alloc_failed: + alloc_failed = True + break + if prep.host_indices is None: + continue + transfers = self.tree_core.build_hicache_transfers( + ct, + last_host_node_id, CacheTransferPhase.PREFETCH, token_ids=prefetch_key.token_ids, prefetch_tokens=len(prefetch_key), last_hash=last_hash, + host_indices=prep.host_indices, ) - if transfers == []: - alloc_failed = True - break if transfers: - comp_xfers[comp.component_type] = transfers + comp_xfers[ct] = transfers kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=None) sidecar_xfers = self._build_sidecar_transfers( CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers @@ -2038,7 +1169,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): self.cache_controller.append_host_mem_release( extra_pools=[x for xfers in comp_xfers.values() for x in xfers], ) - self.dec_host_lock_ref(last_host_node, anchor_lock_params) + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) return aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] @@ -2051,7 +1182,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): extra_pools=aux_xfers or None, ) self.ongoing_prefetch[req_id] = _OngoingPrefetch( - last_host_node, + last_host_node_id, prefetch_key, None, operation, @@ -2108,7 +1239,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return True ( - last_host_node, + last_host_node_id, prefetch_key, host_indices, operation, @@ -2132,7 +1263,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): completed_tokens, hash_value, host_indices, - last_host_node, + last_host_node_id, anchor_lock_params, prefetch_key, ) @@ -2141,21 +1272,27 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return True fetched_key = prefetch_key[:min_completed_tokens] - insert_result = self._insert_helper_host( - last_host_node, + insert_result = self.tree_core.insert_host( + last_host_node_id, fetched_key, host_indices[:min_completed_tokens], hash_value[: min_completed_tokens // self.page_size], ) - for ct, xfers in comp_xfers.items(): - self.components[ct].commit_hicache_transfer( - last_host_node, - CacheTransferPhase.PREFETCH, - xfers, - insert_result=insert_result, - pool_storage_result=operation.pool_storage_result, - ) + # Apply the host-insert walk's actions before the transfer commit. + self._apply_cache_actions(insert_result.cache_actions) + commit_actions: list[CacheAction | ComponentAction] = [] + self.tree_core.commit_hicache_transfers( + last_host_node_id, + CacheTransferPhase.PREFETCH, + comp_xfers, + cache_actions=commit_actions, + insert_result=insert_result, + pool_storage_result=operation.pool_storage_result, + ) + self._apply_cache_actions(commit_actions) + # The commit emits via commit_actions only; the walk's were applied above. + assert not insert_result.cache_actions self.cache_controller.mem_pool_host.free( host_indices[: insert_result.prefix_len] @@ -2163,7 +1300,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): self.cache_controller.append_host_mem_release( host_indices[min_completed_tokens:completed_tokens] ) - self.dec_host_lock_ref(last_host_node, anchor_lock_params) + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) @@ -2190,7 +1327,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): completed_tokens: int, hash_value: list[str], host_indices: torch.Tensor, - last_host_node: UnifiedTreeNode, + last_host_node_id: NodeId, anchor_lock_params: DecLockRefParams, prefetch_key: RadixKey, ) -> Optional[int]: @@ -2251,7 +1388,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): host_indices=host_indices[:completed_tokens], extra_pools=pool_transfers, ) - self.dec_host_lock_ref(last_host_node, anchor_lock_params) + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) self.prefetch_loaded_tokens_by_reqid[req_id] = 0 @@ -2279,7 +1416,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return ( - last_host_node, + last_host_node_id, prefetch_key, host_indices, operation, @@ -2293,7 +1430,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) self._barrier_attn_groups() - self.dec_host_lock_ref(last_host_node, anchor_lock_params) + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[rid] self.cache_controller.append_host_mem_release( host_indices=host_indices[:completed_tokens], @@ -2306,7 +1443,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if info is None: return ( - last_host_node, + last_host_node_id, prefetch_key, _host_indices, _operation, @@ -2317,7 +1454,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): cc.append_host_mem_release( extra_pools=[x for xfers in comp_xfers.values() for x in xfers] ) - self.dec_host_lock_ref(last_host_node, anchor_lock_params) + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) cc.prefetch_tokens_occupied = max( 0, cc.prefetch_tokens_occupied - len(prefetch_key) ) @@ -2391,8 +1528,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): drained += 1 entry = self.ongoing_backup.pop(operation.id, None) if entry is not None: - node, lock_params = entry - self.dec_host_lock_ref(node, lock_params) + node_id, lock_params = entry + self.dec_host_lock_ref(node_id, lock_params) if ( log_metrics and self.enable_storage_metrics @@ -2582,7 +1719,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return # Every rank must enter the all_reduce below; ongoing_write_through can - # diverge across ranks (e.g. write_backup returning 0 on a subset). + # diverge across ranks (e.g. a backup returning 0 on a subset). finish_count = 0 if self.pp_rank == 0: for ack in cc.ack_write_queue: @@ -2641,54 +1778,43 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): def init_load_back( self, params: InitLoadBackParams, - ) -> tuple[torch.Tensor, UnifiedTreeNode]: + ) -> tuple[torch.Tensor, NodeId]: """Prepare KV cache loading from host to device. Returns (device_indices, last_node) tuple.""" - best_match_node = params.best_match_node + best_match_node_id = params.best_match_node mem_quota = params.mem_quota req = params.req assert req is not None - last_best_match_device_node = req.last_node - - def _collect_new_prefix_indices() -> torch.Tensor: - prefix_chunks: list[torch.Tensor] = [] - node = best_match_node - while node is not last_best_match_device_node: - value = node.component_data[BASE_COMPONENT_TYPE].value - assert value is not None - prefix_chunks.append(value) - node = node.parent - if not prefix_chunks: - return self._empty_match_result.device_indices - prefix_chunks.reverse() - return torch.cat(prefix_chunks) + last_best_match_device_node_id = req.last_node if ( - best_match_node.evicted + self.tree_core.is_full_device_evicted(best_match_node_id) or params.host_hit_length > 0 or ( req is not None and (req.swa_host_hit_length > 0 or req.mamba_host_hit_length > 0) ) ): - if self.load_back(best_match_node, mem_quota, req=req): - new_indices = _collect_new_prefix_indices() + if self.load_back(best_match_node_id, mem_quota, req=req): + new_indices = self.tree_core.collect_full_device_indices( + best_match_node_id, last_best_match_device_node_id + ) if new_indices.numel() == 0: return ( - self._empty_match_result.device_indices, - last_best_match_device_node, + self.tree_core.empty_match_result.device_indices, + last_best_match_device_node_id, ) logger.debug( "init_load_back success: loaded %d tokens for node %d", len(new_indices), - best_match_node.id, + best_match_node_id, ) - return new_indices, best_match_node + return new_indices, best_match_node_id return ( - self._empty_match_result.device_indices, - last_best_match_device_node, + self.tree_core.empty_match_result.device_indices, + last_best_match_device_node_id, ) def check_hicache_events(self) -> None: @@ -2720,8 +1846,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): @property def sliding_window_size(self): - swa = self.components.get(ComponentType.SWA) - return swa.sliding_window_size if swa else None + return self._sliding_window_size def swa_reprefill_tail_tokens(self) -> int: """ @@ -2736,15 +1861,15 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): unified_compress_only_hicache = ( self.cache_controller is not None and swa is not None - and swa._swa_kv_pool_host is None + and not self.tree_core.has_swa_host_pool ) return swa.sliding_window_size if unified_compress_only_hicache else 0 def supports_swa(self) -> bool: - return ComponentType.SWA in self.components + return self.is_swa_enabled def supports_mamba(self) -> bool: - return ComponentType.MAMBA in self.components + return self.is_mamba_enabled # ---- Streaming session API (delegates to composed StreamingSession) ---- @@ -2770,92 +1895,44 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return self.session.session_held_mamba_slots(active_pool_idxs) def evictable_size(self) -> int: - return self.component_evictable_size_.get(BASE_COMPONENT_TYPE, 0) + return self.tree_core.evictable_size() def protected_size(self) -> int: - return self.component_protected_size_.get(BASE_COMPONENT_TYPE, 0) + return self.tree_core.protected_size() def full_evictable_size(self) -> int: - return self.evictable_size() + return self.tree_core.full_evictable_size() def full_protected_size(self) -> int: - return self.protected_size() + return self.tree_core.full_protected_size() def swa_evictable_size(self) -> int: - return self.component_evictable_size_.get(ComponentType.SWA, 0) + return self.tree_core.swa_evictable_size() def mamba_evictable_size(self) -> int: - return self.component_evictable_size_.get(ComponentType.MAMBA, 0) + return self.tree_core.mamba_evictable_size() def swa_protected_size(self) -> int: - return self.component_protected_size_.get(ComponentType.SWA, 0) + return self.tree_core.swa_protected_size() def mamba_protected_size(self) -> int: - return self.component_protected_size_.get(ComponentType.MAMBA, 0) + return self.tree_core.mamba_protected_size() - def total_size(self): - total_size = 0 - total_aux_size = 0 - stack = [self.root_node] - while stack: - node = stack.pop() - full_value = node.component_data[BASE_COMPONENT_TYPE].value - if full_value is not None: - total_size += len(full_value) - for ct in self.tree_components: - if ct == BASE_COMPONENT_TYPE: - continue - value = node.component_data[ct].value - if value is not None: - total_aux_size += len(value) - for child in node.children.values(): - stack.append(child) - return total_size, total_aux_size + def total_size(self) -> tuple[int, int]: + return self.tree_core.total_size() def all_values_flatten(self) -> torch.Tensor: - values = [] - - def _dfs(node: UnifiedTreeNode): - for child in node.children.values(): - v = child.component_data[BASE_COMPONENT_TYPE].value - if v is not None: - values.append(v) - _dfs(child) - - _dfs(self.root_node) - if values: - return torch.cat(values) - return torch.tensor([], dtype=torch.int64, device=self.device) - - def _all_component_values_flatten( - self, component_type: ComponentType - ) -> torch.Tensor: - if component_type not in self.components: - return torch.tensor([], dtype=torch.int64, device=self.device) - - values = [] - - def _dfs(node: UnifiedTreeNode): - value = node.component_data[component_type].value - if value is not None: - values.append(value) - for child in node.children.values(): - _dfs(child) - - _dfs(self.root_node) - if values: - return torch.cat(values) - return torch.tensor([], dtype=torch.int64, device=self.device) + return self.tree_core.all_values_flatten() def all_mamba_values_flatten(self) -> torch.Tensor: - return self._all_component_values_flatten(ComponentType.MAMBA) + return self.tree_core.all_mamba_values_flatten() def available_and_evictable_str(self) -> str: if self.supports_swa(): full_available_size = self.token_to_kv_pool_allocator.full_available_size() else: full_available_size = self.token_to_kv_pool_allocator.available_size() - full_evictable = self.component_evictable_size_[BASE_COMPONENT_TYPE] + full_evictable = self.tree_core.component_evictable_size(BASE_COMPONENT_TYPE) lines = [ f"Available full tokens: {full_available_size + full_evictable} " f"(full_available_size={full_available_size} + full_evictable_size_={full_evictable})" @@ -2871,20 +1948,11 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): continue lines.append( - f"Available {ct}: {available_size + self.component_evictable_size_[ct]} " - f"(available_size={available_size} + component_evictable_size_={self.component_evictable_size_[ct]})" + f"Available {ct}: {available_size + self.tree_core.component_evictable_size(ct)} " + f"(available_size={available_size} + component_evictable_size_={self.tree_core.component_evictable_size(ct)})" ) return "\n".join(lines) + "\n" - def _collect_all_nodes(self) -> list[UnifiedTreeNode]: - nodes = [] - stack = [self.root_node] - while stack: - node = stack.pop() - nodes.append(node) - stack.extend(node.children.values()) - return nodes - def sanity_check(self): """Verify tree invariants. @@ -2897,273 +1965,79 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if self.session.any_holding_kv(): return - write_back = ( - self.cache_controller is not None - and self.cache_controller.write_policy == "write_back" - ) - - errors: list[str] = [] - E = errors.append - all_nodes = self._collect_all_nodes() - all_node_set = set(all_nodes) - FCT = BASE_COMPONENT_TYPE - - # ── PART 1: Tree Structure ── - # Root state - if self.root_node.component_data[FCT].value is None: - E("[Root] root missing Full device value") - if self.root_node.component_data[FCT].lock_ref <= 0: - E( - f"[Root] root Full lock_ref={self.root_node.component_data[FCT].lock_ref}" - ) - if self.root_node.parent is not None: - E("[Root] root has a parent pointer") - # Parent ↔ child bidirectional consistency - for node in all_nodes: - for child in node.children.values(): - if child.parent is not node: - pid = child.parent.id if child.parent else None - E(f"[Tree] child {child.id} parent={pid}, expected {node.id}") - if child.key is None: - E(f"[Tree] node {child.id} has no key") - - # ── PART 2: Per-node state machine and leaf qualification ── - expected_dev_leaves: set[UnifiedTreeNode] = set() - expected_hst_leaves: set[UnifiedTreeNode] = set() - - for node in all_nodes: - if node is self.root_node: - continue - nid = node.id - full_dev = node.component_data[FCT].value is not None - full_hst = node.component_data[FCT].host_value is not None - - # Full is the tree backbone, so aux data requires Full data. - for ct in self.tree_components: - if ct == FCT: - continue - cd = node.component_data[ct] - if cd.value is not None and not full_dev: - E(f"node {nid} {ct} device present but Full.value=None") - if cd.host_value is not None and not full_hst: - E(f"node {nid} {ct} host present but Full.host_value=None") - - # Every node must keep Full data on at least one layer. - if not full_dev and not full_hst: - E(f"node {nid} dead: no Full device and no Full host") - - # Parent prefixes must keep data whenever the child does. - if node.parent is not None and node.parent is not self.root_node: - p_dev = node.parent.component_data[FCT].value is not None - p_hst = node.parent.component_data[FCT].host_value is not None - if full_dev and not p_dev: - E(f"node {nid} device present but parent {node.parent.id} evicted") - if full_hst and not p_hst and not write_back: - E(f"node {nid} backed up but parent {node.parent.id} not backed up") - - # Lock hierarchy and counters must stay sane. - fl = node.component_data[FCT].lock_ref - for ct in self.tree_components: - cd = node.component_data[ct] - if cd.lock_ref < 0: - E(f"node {nid} {ct} lock_ref={cd.lock_ref}") - if cd.host_lock_ref < 0: - E(f"node {nid} {ct} host_lock_ref={cd.host_lock_ref}") - if ct != FCT and fl < cd.lock_ref: - E(f"node {nid} full_lock={fl} < {ct}_lock={cd.lock_ref}") - if cd.value is None and cd.lock_ref > 0: - E(f"node {nid} {ct} evicted but lock_ref={cd.lock_ref}") - - # Collect expected leaf qualification (single pass) - if self._is_device_leaf(node): - expected_dev_leaves.add(node) - if self._is_host_leaf(node): - expected_hst_leaves.add(node) - - # ── PART 3: Tracking structures ── - - # Device leaf set must match the expected leaves. - if self.evictable_device_leaves != expected_dev_leaves: - extra = self.evictable_device_leaves - expected_dev_leaves - missing = expected_dev_leaves - self.evictable_device_leaves - if extra: - E(f"D-leaf extra: {[n.id for n in list(extra)[:5]]}") - if missing: - E(f"D-leaf missing: {[n.id for n in list(missing)[:5]]}") - - # Host leaf set must match the expected leaves. - if self.evictable_host_leaves != expected_hst_leaves: - extra = self.evictable_host_leaves - expected_hst_leaves - missing = expected_hst_leaves - self.evictable_host_leaves - if extra: - E(f"H-leaf extra: {[n.id for n in list(extra)[:5]]}") - if missing: - E(f"H-leaf missing: {[n.id for n in list(missing)[:5]]}") - - # D-leaf ∩ H-leaf = ∅ - overlap = self.evictable_device_leaves & self.evictable_host_leaves - if overlap: - E( - f"[Leaf] {len(overlap)} in both sets: {[n.id for n in list(overlap)[:5]]}" - ) - - # Stale nodes: leaf sets must only contain tree-reachable nodes - stale = self.evictable_device_leaves - all_node_set - if stale: - E( - f"{len(stale)} stale nodes in device_leaves: {[n.id for n in list(stale)[:5]]}" - ) - stale = self.evictable_host_leaves - all_node_set - if stale: - E( - f"{len(stale)} stale nodes in host_leaves: {[n.id for n in list(stale)[:5]]}" - ) - - # Per-component LRU tracking - for ct in self.tree_components: - lru = self.lru_lists[ct] - if ct == FCT: - # Full uses leaf sets, not LRU - if len(lru.cache) > 0: - E(f"Full device LRU not empty: {len(lru.cache)}") - if len(self.host_lru_lists[ct].cache) > 0: - E(f"Full host LRU not empty: {len(self.host_lru_lists[ct].cache)}") - else: - # Aux device values must match the device LRU. - tree_ids = { - n.id - for n in all_nodes - if n is not self.root_node - and n.component_data[ct].value is not None - } - lru_ids = set(lru.cache.keys()) - if tree_ids != lru_ids: - E( - f"{ct} device LRU: " - f"+tree={tree_ids - lru_ids}, +lru={lru_ids - tree_ids}" - ) - # Aux host-only states must match the host LRU. - host_lru = self.host_lru_lists[ct] - s3_ids = { - n.id - for n in all_nodes - if n is not self.root_node - and n.component_data[ct].value is None - and n.component_data[ct].host_value is not None - } - host_lru_ids = set(host_lru.cache.keys()) - if s3_ids != host_lru_ids: - E( - f"{ct} host LRU: " - f"+S3={s3_ids - host_lru_ids}, +lru={host_lru_ids - s3_ids}" - ) - # The same aux node must not appear in both device and host LRU. - inv5_overlap = lru_ids & host_lru_ids - if inv5_overlap: - E(f"{ct} in both device and host LRU: {inv5_overlap}") - # Linked-list integrity - self._check_lru_linked_list(lru, ct, "device", errors) - self._check_lru_linked_list(host_lru, ct, "host", errors) - - # ── PART 4: Size Accounting ── - for ct in self.tree_components: - evictable = 0 - protected = 0 - for n in all_nodes: - if n is self.root_node: - continue - cd = n.component_data[ct] - if cd.value is not None: - toks = len(cd.value) - if cd.lock_ref > 0: - protected += toks - else: - evictable += toks - if self.component_evictable_size_[ct] != evictable: - E( - f"[Size] {ct} evictable={self.component_evictable_size_[ct]} " - f"!= recomputed={evictable}" - ) - if self.component_protected_size_[ct] != protected: - E( - f"[Size] {ct} protected={self.component_protected_size_[ct]} " - f"!= recomputed={protected}" - ) - - # ── PART 5: Ongoing Operations ── - for nid, (n, _, _) in self.ongoing_write_through.items(): - if n not in all_node_set: - E(f"[Ongoing] write_through node {nid} not in tree") - elif n.component_data[FCT].lock_ref <= 0: - E( - f"[Ongoing] write_through node {nid} lock_ref={n.component_data[FCT].lock_ref}" - ) - for nid, (n, _, _) in self.ongoing_load_back.items(): - if n not in all_node_set: - E(f"[Ongoing] load_back node {nid} not in tree") - elif n.component_data[FCT].lock_ref <= 0: - E( - f"[Ongoing] load_back node {nid} lock_ref={n.component_data[FCT].lock_ref}" - ) - - # ── Result ── - if errors: - msg = ( - f"Sanity check FAILED ({len(errors)} violations " - f"across {len(all_nodes)} nodes):\n" - + "\n".join(f" {e}" for e in errors) - ) - logger.error(msg) - self.pretty_print() - raise AssertionError(msg) - - def _check_lru_linked_list( - self, - lru: UnifiedLRUList, - ct: ComponentType, - label: str, - errors: list[str], - ) -> None: - """Walk a LRU doubly-linked list, collect integrity errors.""" - pt = lru._pt # use LRU's own pointer slot - visited: set[int] = set() - x = lru.head.lru_next[pt] - prev = lru.head - while x is not None and x != lru.tail: - if x.lru_prev[pt] != prev: - errors.append(f"[{label}][{ct}] broken prev at node {x.id}") - if x.id not in lru.cache: - errors.append(f"[{label}][{ct}] node {x.id} in list not cache") - if x.id in visited: - errors.append(f"[{label}][{ct}] cycle at node {x.id}") - break - visited.add(x.id) - prev = x - x = x.lru_next[pt] - if x is None: - errors.append( - f"[{label}][{ct}] broken chain: lru_next is None " - f"after node {prev.id if hasattr(prev, 'id') else 'head'}" - ) - if len(visited) != len(lru.cache): - errors.append( - f"[{label}][{ct}] list={len(visited)} != cache={len(lru.cache)}" - ) + # Pass ongoing ops as lightweight (id, node_id) pairs so the tree core + # can resolve + validate them without reaching into Controller state. + ongoing_write_through = [ + (nid, wt.node_id) for nid, wt in self.ongoing_write_through.items() + ] + ongoing_load_back = [ + (nid, lb.node_id) for nid, lb in self.ongoing_load_back.items() + ] + self.tree_core.sanity_check(ongoing_write_through, ongoing_load_back) def pretty_print(self) -> None: - stack = [(self.root_node, 0)] - while stack: - node, indent = stack.pop() - component_str = " ".join( - f"{ct}={'yes' if node.component_data[ct].value is not None else 'no'}" - for ct in self.tree_components - ) - print( - " " * indent, - f"[{node.id}]", - len(node.key), - f"full_lock={node.component_data[BASE_COMPONENT_TYPE].lock_ref}", - component_str, - ) - for child in node.children.values(): - stack.append((child, indent + 2)) + self.tree_core.pretty_print() + + # ---- TreeCore state delegation ---- + # The facade re-exposes tree-owned config (page_size, enable_storage, ...) so its + # own coordination methods and external callers read them off the cache. + + # ``page_size`` keeps a setter: StreamingSession forwards assignment onto its + # inner cache (the PrefixCacheTrait surface). + @property + def page_size(self): + return self.tree_core.page_size + + @page_size.setter + def page_size(self, value) -> None: + self.tree_core.page_size = value + + @property + def enable_storage(self): + return self.tree_core.enable_storage + + @enable_storage.setter + def enable_storage(self, value) -> None: + self.tree_core.enable_storage = value + + @property + def write_through_threshold(self): + return self.tree_core.write_through_threshold + + @write_through_threshold.setter + def write_through_threshold(self, value) -> None: + self.tree_core.write_through_threshold = value + + @property + def is_write_back(self): + return self.tree_core.is_write_back + + @is_write_back.setter + def is_write_back(self, value) -> None: + self.tree_core.is_write_back = value + + @property + def device(self): + return self.tree_core.device + + @property + def root_node(self): + return self.tree_core.root_node + + def take_events(self): + # Drain the KV event queue from the TreeCore. + return self.tree_core.take_events() + + def resolve_node_handle(self, node_handle): + """Look up the node object from its NodeId. + + TODO(Jialin): Remove after the Unified Radix Cache split. + """ + if isinstance(node_handle, int): + return self.tree_core.node_by_id(node_handle) + # Internal callers (and the session sentinel / None) pass a non-int through. + return node_handle + + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: + """The root's NodeId -- URC match results carry NodeIds.""" + return self.tree_core.root_node.id diff --git a/test/registered/kv_canary/test_self_unit_radix_walker.py b/test/registered/kv_canary/test_self_unit_radix_walker.py index 395b7a627..6651bfd49 100644 --- a/test/registered/kv_canary/test_self_unit_radix_walker.py +++ b/test/registered/kv_canary/test_self_unit_radix_walker.py @@ -1,11 +1,16 @@ from __future__ import annotations import unittest +from array import array +from unittest import mock import torch from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode +from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_cache_components import ( BASE_COMPONENT_TYPE, ComponentType, @@ -129,17 +134,37 @@ class TestSelfUnitRadixWalker(CustomTestCase): ) self.assertEqual(result.slot_indices.tolist(), [3, 4]) + def test_unified_swa_sweep_gates_on_swa_lock_not_full_lock(self): + """With unlocked_only + swa_resident_only, the sweep filters on the SWA + component lock: a FULL-locked node whose SWA lock was already released + (early dec_swa_lock_only) must still be swept, and a node whose SWA + lock is still held must not.""" + cache = self._make_unified_cache((ComponentType.FULL, ComponentType.SWA)) + self._add_unified_child(cache, [1, 2], lock_ref=1, swa_value=[1, 2]) + held = self._add_unified_child(cache, [3, 4], swa_value=[3, 4]) + held.component_data[ComponentType.SWA].lock_ref = 1 + + result = walk_radix_cache_for_canary( + radix_cache=cache, unlocked_only=True, swa_resident_only=True + ) + self.assertEqual(result.slot_indices.tolist(), [1, 2]) + def _make_unified_cache( self, tree_components: tuple[ComponentType, ...] ) -> UnifiedRadixCache: cache = UnifiedRadixCache.__new__(UnifiedRadixCache) cache.tree_components = tree_components cache.components = {ct: None for ct in tree_components} - root = UnifiedTreeNode(tree_components) - root.component_data[BASE_COMPONENT_TYPE].value = torch.tensor( - [], dtype=torch.int32, device=self.device + cache.is_swa_enabled = ComponentType.SWA in tree_components + cache.tree_core = UnifiedTreeCore( + CacheInitParams( + disable=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=None, + page_size=1, + ), + {ct: mock.MagicMock() for ct in tree_components}, ) - cache.root_node = root return cache def _add_unified_child( @@ -190,6 +215,23 @@ class TestSelfUnitRadixWalker(CustomTestCase): ) self.assertEqual(result.slot_indices.tolist(), [3, 4]) + def test_unified_walk_spans_device_evicted_nodes_without_emitting_them(self): + """Verify device-evicted (host-only) nodes emit no slots but still advance + positions by their key length and pass the prev-slot chain through.""" + cache = self._make_unified_cache((ComponentType.FULL,)) + evicted = UnifiedTreeNode(cache.tree_components) + evicted.parent = cache.root_node + evicted.key = RadixKey(array("q", [7, 8]), None) + cache.root_node.children[evicted.id] = evicted + grandchild = self._add_unified_child(cache, [5, 6]) + cache.root_node.children.pop(grandchild.id) + grandchild.parent = evicted + evicted.children[grandchild.id] = grandchild + result = walk_radix_cache_for_canary(radix_cache=cache) + self.assertEqual(result.slot_indices.tolist(), [5, 6]) + self.assertEqual(result.positions.tolist(), [2, 3]) + self.assertEqual(result.prev_slot_indices.tolist(), [-1, 5]) + def test_unified_swa_resident_only_noop_without_swa_component(self): """Verify swa_resident_only is a no-op when SWA is not enabled.""" cache = self._make_unified_cache((ComponentType.FULL,)) diff --git a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py index be0eddc1f..283960f18 100644 --- a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py +++ b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py @@ -1,44 +1,37 @@ """CPU-only unit tests for the per-path Mamba checkpoint cap.""" -from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci register_cpu_ci(est_time=2, suite="base-a-test-cpu") +register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small") import argparse import unittest -from types import SimpleNamespace +from collections import defaultdict +from unittest import mock import torch +from test_unified_radix_cache_unittest import CacheConfig, UnifiedRadixCacheSuite +from sglang.srt.mem_cache.unified_cache.cache_action import MambaEvictExcessPathStates +from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_cache_components.mamba_component import ( MambaComponent, ) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( ComponentType, ) -from sglang.srt.mem_cache.unified_radix_cache import ( - UnifiedLRUList, - UnifiedRadixCache, - UnifiedTreeNode, -) +from sglang.srt.mem_cache.unified_radix_cache import UnifiedLRUList, UnifiedTreeNode from sglang.srt.server_args import ServerArgs +from sglang.test.test_utils import CustomTestCase -class _RecordingAllocator: - def __init__(self): - self.freed = [] - - def free(self, value): - self.freed.extend(value.tolist()) - - -class _FakeUnifiedCache: +class _FakeTreeCore: tree_components = (ComponentType.FULL, ComponentType.MAMBA) def __init__(self): self.root_node = UnifiedTreeNode(self.tree_components) self.evictable_device_leaves = set() - self.req_to_token_pool = SimpleNamespace(mamba_allocator=_RecordingAllocator()) self.component_evictable_size_ = {ComponentType.MAMBA: 0} self.component_protected_size_ = {ComponentType.MAMBA: 0} self.lru_lists = { @@ -48,41 +41,47 @@ class _FakeUnifiedCache: } self.host_lru_lists = { ComponentType.MAMBA: UnifiedLRUList( - ComponentType.MAMBA, self.tree_components + ComponentType.MAMBA, self.tree_components, use_host_ptr=True ) } self.evicted = [] self.cascaded = [] - def _evict_component_and_detach_lru(self, node, component, **kwargs): + def _evict_component_and_detach_lru(self, node, component, *args, **kwargs): self.evicted.append(node) - return UnifiedRadixCache._evict_component_and_detach_lru( - self, node, component, **kwargs + return UnifiedTreeCore._evict_component_and_detach_lru( + self, node, component, *args, **kwargs ) - def _cascade_evict(self, node, component, tracker): + def _cascade_evict(self, node, component, tracker, device_frees, host_frees): self.cascaded.append(node) +class _FakeUnifiedCache: + tree_components = _FakeTreeCore.tree_components + + def _build_unified_chain(cap, length=3): cache = _FakeUnifiedCache() + core = _FakeTreeCore() component = object.__new__(MambaComponent) component.cache = cache + component.tree_core = core component.mamba_max_states_per_path = cap nodes = [] - parent = cache.root_node + parent = core.root_node for index in range(length): - node = UnifiedTreeNode(cache.tree_components) + node = UnifiedTreeNode(core.tree_components) node.parent = parent node.component_data[ComponentType.FULL].value = torch.tensor([100 + index]) node.component_data[ComponentType.MAMBA].value = torch.tensor([index]) parent.children[index] = node - cache.component_evictable_size_[ComponentType.MAMBA] += 1 - cache.lru_lists[ComponentType.MAMBA].insert_mru(node) + core.component_evictable_size_[ComponentType.MAMBA] += 1 + core.lru_lists[ComponentType.MAMBA].insert_mru(node) nodes.append(node) parent = node - return component, nodes, cache + return component, nodes, core, cache class TestMambaPathStateCap(unittest.TestCase): @@ -114,20 +113,22 @@ class TestMambaPathStateCap(unittest.TestCase): ) def test_unified_cache_removes_only_shallow_mamba_state(self): - component, nodes, cache = _build_unified_chain(cap=2) + component, nodes, core, cache = _build_unified_chain(cap=2) - component._evict_excess_path_states(nodes[-1]) + device_frees = defaultdict(list) + host_frees = defaultdict(list) + component._evict_excess_path_states(nodes[-1], device_frees, host_frees) - self.assertEqual(cache.evicted, [nodes[0]]) - self.assertEqual(cache.cascaded, [nodes[0]]) + self.assertEqual(core.evicted, [nodes[0]]) + self.assertEqual(core.cascaded, [nodes[0]]) self.assertIsNone(nodes[0].component_data[ComponentType.MAMBA].value) self.assertIsNotNone(nodes[-1].component_data[ComponentType.MAMBA].value) self.assertEqual( - cache.req_to_token_pool.mamba_allocator.freed, + [v.item() for v in device_frees[ComponentType.MAMBA]], [0], ) - self.assertEqual(cache.component_evictable_size_[ComponentType.MAMBA], 2) - self.assertFalse(cache.lru_lists[ComponentType.MAMBA].in_list(nodes[0])) + self.assertEqual(core.component_evictable_size_[ComponentType.MAMBA], 2) + self.assertFalse(core.lru_lists[ComponentType.MAMBA].in_list(nodes[0])) self.assertTrue( all( node.component_data[ComponentType.FULL].value is not None @@ -136,37 +137,44 @@ class TestMambaPathStateCap(unittest.TestCase): ) def test_unified_cache_cap_is_soft_for_fork_and_locked_nodes(self): - component, nodes, cache = _build_unified_chain(cap=1, length=4) - fork_child = UnifiedTreeNode(cache.tree_components) + component, nodes, core, cache = _build_unified_chain(cap=1, length=4) + fork_child = UnifiedTreeNode(core.tree_components) fork_child.parent = nodes[0] nodes[0].children["fork"] = fork_child nodes[1].component_data[ComponentType.MAMBA].lock_ref = 1 - component._evict_excess_path_states(nodes[-1]) + device_frees = defaultdict(list) + host_frees = defaultdict(list) + component._evict_excess_path_states(nodes[-1], device_frees, host_frees) - self.assertEqual(cache.evicted, [nodes[2]]) + self.assertEqual(core.evicted, [nodes[2]]) self.assertIsNotNone(nodes[0].component_data[ComponentType.MAMBA].value) self.assertIsNotNone(nodes[1].component_data[ComponentType.MAMBA].value) self.assertIsNone(nodes[2].component_data[ComponentType.MAMBA].value) self.assertIsNotNone(nodes[3].component_data[ComponentType.MAMBA].value) def test_unified_cache_preserves_existing_host_backup(self): - component, nodes, cache = _build_unified_chain(cap=2) + component, nodes, core, cache = _build_unified_chain(cap=2) mamba_data = nodes[0].component_data[ComponentType.MAMBA] mamba_data.host_value = torch.tensor([10]) - component._evict_excess_path_states(nodes[-1]) + device_frees = defaultdict(list) + host_frees = defaultdict(list) + component._evict_excess_path_states(nodes[-1], device_frees, host_frees) self.assertIsNone(mamba_data.value) self.assertIsNotNone(mamba_data.host_value) - self.assertTrue(cache.host_lru_lists[ComponentType.MAMBA].in_list(nodes[0])) + self.assertTrue(core.host_lru_lists[ComponentType.MAMBA].in_list(nodes[0])) def test_unified_cache_negative_one_disables_cap(self): - component, nodes, cache = _build_unified_chain(cap=-1) + component, nodes, core, cache = _build_unified_chain(cap=-1) - component._evict_excess_path_states(nodes[-1]) + device_frees = defaultdict(list) + host_frees = defaultdict(list) + component._evict_excess_path_states(nodes[-1], device_frees, host_frees) - self.assertEqual(cache.evicted, []) + self.assertEqual(dict(device_frees), {}) + self.assertEqual(core.evicted, []) self.assertTrue( all( node.component_data[ComponentType.MAMBA].value is not None @@ -175,5 +183,162 @@ class TestMambaPathStateCap(unittest.TestCase): ) +@unittest.skipUnless(torch.cuda.is_available(), "mamba pool fixtures need CUDA") +class TestMambaPathCapWriteThroughOrdering(CustomTestCase): + """CI-active write-through/path-cap ordering regressions (the unified radix + cache unittest module is temporarily gated off on trunk).""" + + cfg = CacheConfig(components=(ComponentType.FULL, ComponentType.MAMBA)) + _rid = 0 + # Borrow the fixture helpers without inheriting the full gated suite. + _make_req = UnifiedRadixCacheSuite._make_req + _alloc = UnifiedRadixCacheSuite._alloc + _insert = UnifiedRadixCacheSuite._insert + _init_hicache = UnifiedRadixCacheSuite._init_hicache + _build_hicache_fixture = UnifiedRadixCacheSuite._build_hicache_fixture + + def test_write_through_backup_survives_mamba_path_cap(self): + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = 2 + cache.components[ComponentType.MAMBA].mamba_max_states_per_path = 1 + + # The first insert creates the ancestor with a mamba state (hit_count 1). + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + ancestor = next(iter(cache.root_node.children.values())) + self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value) + + # The extending insert crosses the ancestor's write-through threshold in + # the same walk whose commit runs the path-cap eviction; the cap must + # leave the pending-backup node's device state for the deferred BackupKV. + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + cache.writing_check(write_back=True) + + self.assertTrue(ancestor.backuped) + self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].host_value) + + def test_write_through_backup_chain_survives_mamba_path_cap(self): + """A failed backup leaves an unbacked ancestor inside a later deferred + backup chain; the cap walk must spare the whole chain, not just its tip.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = 3 + mamba_comp = cache.components[ComponentType.MAMBA] + + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + ancestor = next(iter(cache.root_node.children.values())) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + middle = next(iter(ancestor.children.values())) + + # The ancestor crosses the threshold here; a host-exhaustion failure + # leaves it unbacked with hit_count past the bar and its state intact. + with mock.patch.object(cache, "_execute_kv_backup", return_value=None): + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6]) + self.assertFalse(ancestor.backuped) + self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value) + + # The middle node crosses next, so the deferred chain is + # [ancestor, middle]; the extending insert adopts a new leaf state, + # firing the now-enabled cap walk before the chain executes — it must + # not evict either chain node's device state. + mamba_comp.mamba_max_states_per_path = 1 + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6, 7, 8]) + cache.writing_check(write_back=True) + + self.assertTrue(ancestor.backuped) + self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].host_value) + self.assertTrue(middle.backuped) + self.assertIsNotNone(middle.component_data[ComponentType.MAMBA].host_value) + + def test_backup_retry_after_mamba_cap_skips_tombstoned_state(self): + """A backup that fails before the cap and retries via the leaf action + rebuilds its spec post-cap: KV backs up, the tombstoned mamba arm stays gone.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = 1 + mamba_comp = cache.components[ComponentType.MAMBA] + + # A failed write-through leaves the ancestor unbacked with device state. + with mock.patch.object(cache, "_execute_kv_backup", return_value=None): + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + ancestor = next(iter(cache.root_node.children.values())) + self.assertFalse(ancestor.backuped) + self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value) + + # The walk backup fails again, the cap tombstones the unlocked + # ancestor's mamba state, then the leaf-action retry succeeds. + mamba_comp.mamba_max_states_per_path = 1 + real_backup = cache._execute_kv_backup + attempts = [] + + def fail_once(*args, **kwargs): + attempts.append(args) + if len(attempts) == 1: + return None + return real_backup(*args, **kwargs) + + with mock.patch.object(cache, "_execute_kv_backup", side_effect=fail_once): + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + leaf = next(iter(ancestor.children.values())) + cache.writing_check(write_back=True) + + # Post-cap spec rebuild: no resurrection of the tombstoned mamba state. + self.assertTrue(ancestor.backuped) + ancestor_cd = ancestor.component_data[ComponentType.MAMBA] + self.assertIsNone(ancestor_cd.value) + self.assertIsNone(ancestor_cd.host_value) + self.assertTrue(leaf.backuped) + self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].host_value) + cache.sanity_check() + + def test_walk_backup_excludes_same_insert_restamped_mamba(self): + """The walked target's backup executes before commit hooks, so a mamba + value re-stamped by the same insert stays out of the host backup.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + mamba_comp = cache.components[ComponentType.MAMBA] + + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + ancestor = next(iter(cache.root_node.children.values())) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + + # The cap walk tombstones the ancestor's mamba state. + mamba_comp.mamba_max_states_per_path = 1 + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6]) + self.assertIsNone(ancestor.component_data[ComponentType.MAMBA].value) + + # Re-inserting [1, 2] crosses the threshold and re-stamps the tombstone + # in the same insert; the backup must not carry the fresh mamba state. + cache.write_through_threshold = ancestor.hit_count + 1 + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + cache.writing_check(write_back=True) + + self.assertTrue(ancestor.backuped) + ancestor_cd = ancestor.component_data[ComponentType.MAMBA] + self.assertIsNotNone(ancestor_cd.value) + self.assertIsNone(ancestor_cd.host_value) + + def test_cap_walk_failure_still_drains_collected_frees(self): + """A cap walk that raises mid-eviction must still free the tombstoned + slots it already collected (the pre-split inline frees could not leak).""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + mamba_comp = cache.components[ComponentType.MAMBA] + + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + ancestor = next(iter(cache.root_node.children.values())) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + leaf = next(iter(ancestor.children.values())) + self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value) + + mamba_comp.mamba_max_states_per_path = 1 + available = req_to_token_pool.mamba_allocator.available_size() + with mock.patch.object( + cache.tree_core, "_cascade_evict", side_effect=RuntimeError("boom") + ): + with self.assertRaises(RuntimeError): + mamba_comp.apply_component_action(MambaEvictExcessPathStates(leaf.id)) + + self.assertIsNone(ancestor.component_data[ComponentType.MAMBA].value) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), available + 1 + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_tree_core_registry.py b/test/registered/unit/mem_cache/test_tree_core_registry.py new file mode 100644 index 000000000..28a37a51e --- /dev/null +++ b/test/registered/unit/mem_cache/test_tree_core_registry.py @@ -0,0 +1,189 @@ +"""Unit tests for the tree-core backend registry.""" + +import unittest +from unittest import mock + +from sglang.srt.environ import envs +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType +from sglang.srt.mem_cache.unified_cache.tree_core_registry import ( + _TREE_CORE_REGISTRY, + create_tree_core, + register_tree_core_backend, + registered_tree_core_backends, +) +from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore +from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + EvictLayer, + TreeComponent, +) +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _cache_init_params(**kwargs) -> CacheInitParams: + return CacheInitParams( + disable=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=None, + page_size=2, + tree_components=(ComponentType.FULL,), + **kwargs, + ) + + +class _StubFullComponent(TreeComponent): + component_type = ComponentType.FULL + + def create_match_validator(self, match_device_only: bool = False): + return lambda node: True + + def redistribute_on_node_split(self, new_parent, child): + return None + + def evict_component( + self, node, device_frees, host_frees, target: EvictLayer = EvictLayer.DEVICE + ) -> tuple[int, int]: + return 0, 0 + + def acquire_component_lock(self, node, result): + return result + + def release_component_lock(self, node, params): + return None + + def _evict_device_start(self, request_cnt) -> None: + pass + + def _evict_device_next_node(self, tracker, device_frees, host_frees): + return None + + def _evict_device_end(self) -> None: + pass + + +class _StubMambaComponent(_StubFullComponent): + component_type = ComponentType.MAMBA + + +class TreeCoreRegistryTest(CustomTestCase): + def setUp(self): + self._registry_snapshot = dict(_TREE_CORE_REGISTRY) + + def tearDown(self): + _TREE_CORE_REGISTRY.clear() + _TREE_CORE_REGISTRY.update(self._registry_snapshot) + + def test_registry_contains_python(self): + self.assertIn("python", registered_tree_core_backends()) + + def test_python_backend_builds_the_python_tree(self): + component = mock.MagicMock() + core = create_tree_core( + name="python", + params=_cache_init_params(), + components={ComponentType.FULL: component}, + ) + self.assertIsInstance(core, UnifiedTreeCore) + self.assertIs(component.tree_core, core) + + def test_unknown_backend_raises_naming_the_known_backends(self): + with self.assertRaisesRegex(ValueError, "not registered") as cm: + create_tree_core( + name="not_a_real_backend", + params=_cache_init_params(), + components={}, + ) + self.assertIn("'python'", str(cm.exception)) + + def test_register_rejects_empty_name(self): + with self.assertRaises(ValueError): + register_tree_core_backend(" ", mock.MagicMock()) + + def test_register_rejects_duplicate_name(self): + with self.assertRaises(ValueError): + register_tree_core_backend("python", mock.MagicMock()) + + def test_create_dispatches_to_a_registered_factory(self): + core = mock.MagicMock() + factory = mock.MagicMock(return_value=core) + register_tree_core_backend("custom", factory) + params = _cache_init_params() + components = {ComponentType.FULL: mock.MagicMock()} + result = create_tree_core(name="custom", params=params, components=components) + factory.assert_called_once_with(params, components) + self.assertIs(result, core) + + +class UnifiedRadixCacheTreeCoreSelectionTest(CustomTestCase): + def setUp(self): + self._registry_snapshot = dict(_TREE_CORE_REGISTRY) + + def tearDown(self): + _TREE_CORE_REGISTRY.clear() + _TREE_CORE_REGISTRY.update(self._registry_snapshot) + + def _cache_params( + self, + tree_components=(ComponentType.FULL,), + component_registry_override={ComponentType.FULL: _StubFullComponent}, + **kwargs, + ) -> CacheInitParams: + return CacheInitParams( + disable=True, + req_to_token_pool=ReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + ), + token_to_kv_pool_allocator=None, + page_size=1, + tree_components=tree_components, + component_registry_override=component_registry_override, + **kwargs, + ) + + def test_init_downgrades_is_eagle_when_mamba_is_enabled(self): + params = self._cache_params( + is_eagle=True, + tree_components=(ComponentType.FULL, ComponentType.MAMBA), + component_registry_override={ + ComponentType.FULL: _StubFullComponent, + ComponentType.MAMBA: _StubMambaComponent, + }, + ) + cache = UnifiedRadixCache(params) + self.assertFalse(cache.tree_core.is_eagle) + + def test_init_keeps_is_eagle_without_mamba(self): + params = self._cache_params(is_eagle=True) + cache = UnifiedRadixCache(params) + self.assertTrue(cache.tree_core.is_eagle) + + def test_default_backend_builds_the_python_tree_core(self): + cache = UnifiedRadixCache(params=self._cache_params()) + self.assertIsInstance(cache.tree_core, UnifiedTreeCore) + component = cache.components[ComponentType.FULL] + self.assertIs(component.tree_core, cache.tree_core) + + def test_env_var_routes_construction_to_the_selected_backend(self): + """SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND selects the registered factory + the cache constructs its tree through.""" + core = mock.MagicMock() + factory = mock.MagicMock(return_value=core) + register_tree_core_backend("custom_env_backend", factory) + with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("custom_env_backend"): + cache = UnifiedRadixCache(params=self._cache_params()) + factory.assert_called_once() + self.assertIs(cache.tree_core, core) + component = cache.components[ComponentType.FULL] + self.assertIs(component.tree_core, core) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 94e3c9403..bd721e671 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -566,7 +566,7 @@ def bench_lock_unlock( nodes = [] for seq in env.seqs[: num_seqs // 2]: r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - if r.last_device_node != env.tree.root_node: + if r.last_device_node != env.tree.root_node_handle(): nodes.append(r.last_device_node) if not nodes: return BenchResult("lock_unlock", 0, 0, 0, []) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 1794ae6f4..0830ea6d9 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -2,10 +2,12 @@ import json import shutil +import sys import tempfile import time import unittest from array import array +from collections import defaultdict from dataclasses import dataclass, replace from typing import Optional from unittest import mock @@ -31,6 +33,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( InsertParams, MatchPrefixParams, MatchResult, + zero_match_result, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import available_and_evictable_str @@ -43,6 +46,23 @@ from sglang.srt.mem_cache.memory_pool import ( ) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.mem_cache.unified_cache.cache_action import ( + FreeComponentDeviceSlot, + FreeComponentHostSlot, + FreeDeviceKV, + RebuildFullToSWAMapping, + RecoverSWAWithLockedFull, + ReplaceWriteThroughOnNodeSplit, + SWARebuild, +) +from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + DecSwaLockOnlyResult, + DemoteResult, + DriveHostEvictionResult, + DropSubtreeNoHostResult, + EvictDeviceLeafResult, + EvictDeviceNextNodeResult, +) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, @@ -54,6 +74,7 @@ from sglang.srt.mem_cache.unified_radix_cache import ( UnifiedLRUList, UnifiedRadixCache, UnifiedTreeNode, + _OngoingWriteThrough, ) from sglang.srt.runtime_context import get_server_args from sglang.srt.sampling.sampling_params import SamplingParams @@ -65,8 +86,8 @@ from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") -register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") +register_cuda_ci(est_time=16, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=16, suite="stage-b-test-1-gpu-small-amd") import pytest as _pytest_defer @@ -158,19 +179,25 @@ class _FakeFullComponent(TreeComponent): return None def evict_component( - self, node, target: EvictLayer = EvictLayer.DEVICE + self, node, device_frees, host_frees, target: EvictLayer = EvictLayer.DEVICE ) -> tuple[int, int]: return 0, 0 - def drive_eviction(self, params: EvictParams, tracker: dict[ComponentType, int]): - return None - def acquire_component_lock(self, node, result): return result def release_component_lock(self, node, params): return None + def _evict_device_start(self, request_cnt) -> None: + pass + + def _evict_device_next_node(self, tracker, device_frees, host_frees): + return None + + def _evict_device_end(self) -> None: + pass + class TestUnifiedRadixComponentRegistryOverride(CustomTestCase): def test_component_registry_override_is_instance_local(self): @@ -228,6 +255,13 @@ class TestUnifiedTreeNodeGetPrefixHashValues(CustomTestCase): self.assertEqual(n4.get_prefix_hash_values(n3), ["h1", "h2", "h3"]) +def _write_backup(cache, node, write_back: bool = False) -> int: + """Back up one node's KV D->H via the tree's build+execute primitives.""" + return cache._execute_and_commit_kv_backup( + cache.tree_core._build_backup_kv_action(node, write_back), write_back + ) + + def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False): """Create (tree, allocator, req_to_token_pool) from a CacheConfig.""" server_args = ServerArgs( @@ -375,7 +409,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): self.assertIsNotNone(value) cache.insert(InsertParams(key=RadixKey(tokens), value=value)) match = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) - leaf = match.last_device_node + leaf = cache.resolve_node_handle(match.last_device_node) self.assertTrue(leaf.key.is_bigram) self.assertEqual(len(leaf.hash_value), 2) @@ -412,7 +446,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): controller = FakeCacheController() cache.cache_controller = controller - cache.prefetch_from_storage("req", cache.root_node, tokens) + cache.prefetch_from_storage("req", cache.root_node.id, tokens) _, storage_key, _, _, _ = controller.prefetch_args self.assertIsInstance(storage_key, RadixKey) @@ -464,8 +498,9 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): def _leaf_for(self, cache, tokens): match = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIsNot(match.last_device_node, cache.root_node) - return match.last_device_node + leaf = cache.resolve_node_handle(match.last_device_node) + self.assertIsNot(leaf, cache.root_node) + return leaf def _init_hicache(self, cache, *, write_policy: str = "write_through"): import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler @@ -504,12 +539,12 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): cache.load_back_threshold = 0 def _backup_node(self, cache, node): - backed_up = cache.write_backup(node, write_back=True) + backed_up = _write_backup(cache, node, write_back=True) self.assertGreater(backed_up, 0) cache.writing_check(write_back=True) def _load_back_node(self, cache, node): - loaded = cache.load_back(node) + loaded = cache.load_back(node.id) self.assertTrue(loaded) producer_id = cache.ready_to_load_host_cache() self.assertNotEqual(producer_id, -1) @@ -593,7 +628,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): self._insert(cache, allocator, [1, 2, 3, 4]) node = self._leaf_for(cache, [1, 2, 3, 4]) - backed_up = cache.write_backup(node, write_back=True) + backed_up = _write_backup(cache, node, write_back=True) self.assertGreater(backed_up, 0) # Split the node while its write-through DMA is still pending. @@ -607,7 +642,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): cache.writing_check(write_back=True) self.assertEqual( [ - list(call.args[0].key.token_ids) + list(cache.resolve_node_handle(call.args[0]).key.token_ids) for call in backup_storage.call_args_list ], [[1, 2], [3, 4]], @@ -858,7 +893,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -899,7 +934,7 @@ class UnifiedRadixCacheSuite: req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len req.kv = ReqKvInfo(kv_allocated_len=kv_len, swa_evicted_seqlen=0) - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -945,7 +980,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.swa_prefix_lock_released = True @@ -980,7 +1015,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -1017,7 +1052,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, len(tokens)) req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) req.kv_committed_len = len(tokens) - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -1119,7 +1154,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -1200,7 +1235,11 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(m.device_indices), len(seq)) self.assertIsNotNone(req2.mamba_pool_idx) - src_value = m.last_device_node.component_data[ComponentType.MAMBA].value + src_value = ( + cache.resolve_node_handle(m.last_device_node) + .component_data[ComponentType.MAMBA] + .value + ) self.assertTrue( torch.all( mamba_pool.mamba_cache.conv[0][:, req2.mamba_pool_idx] @@ -1222,16 +1261,25 @@ class UnifiedRadixCacheSuite: cache, allocator, req_to_token_pool, tokens + self._make_seq(100, 1) ) - node = cache.match_prefix( + last_device_node = cache.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens))) ).last_device_node + node = cache.tree_core.node_by_id(last_device_node) old_full_value = node.component_data[ComponentType.FULL].value.clone() swa_component = cache.components[ComponentType.SWA] tracker = {ct: 0 for ct in cache.tree_components} - cache._evict_component_and_detach_lru(node, swa_component, tracker=tracker) + device_frees = defaultdict(list) + cache.tree_core._evict_component_and_detach_lru( + node, + swa_component, + tracker=tracker, + device_frees=device_frees, + host_frees=defaultdict(list), + ) + cache._drain_device_frees(device_frees) self.assertIsNone(node.component_data[ComponentType.SWA].value) - lock_result = cache.inc_lock_ref(node) + lock_result = cache.inc_lock_ref(last_device_node) req = self._make_req(req_to_token_pool) req.origin_input_ids = array("q", tokens) req.output_ids = [] @@ -1241,7 +1289,7 @@ class UnifiedRadixCacheSuite: fresh_value = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value) req.kv_committed_len = kv_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -1275,7 +1323,7 @@ class UnifiedRadixCacheSuite: req.last_node, DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), ) - cache.dec_lock_ref(node, lock_result.to_dec_params()) + cache.dec_lock_ref(last_device_node, lock_result.to_dec_params()) cache.sanity_check() def test_swa_insert_keeps_full_leaf_when_entire_span_is_outside_window(self): @@ -1325,7 +1373,7 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq_ab) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) - node_a = m.last_device_node + node_a = cache.resolve_node_handle(m.last_device_node) self.assertGreater(len(node_a.children), 0, "A must be internal") swa_cd = node_a.component_data[ComponentType.SWA] @@ -1337,15 +1385,15 @@ class UnifiedRadixCacheSuite: # strictly-lower-tier Mamba lock (the co-located Mamba is useless once SWA # is gone), leaving only the Full path-lock held. This is what guarantees # the later SWA-eviction cascade never meets a legitimately-locked Mamba. - lock_result = cache.inc_lock_ref(node_a) + lock_result = cache.inc_lock_ref(node_a.id) self.assertGreaterEqual(mamba_cd.lock_ref, 1, "Mamba locked before release") - cache.dec_swa_lock_only(node_a, lock_result.swa_uuid_for_lock) + cache.dec_swa_lock_only(node_a.id, lock_result.swa_uuid_for_lock) self.assertEqual(swa_cd.lock_ref, 0) self.assertEqual( mamba_cd.lock_ref, 0, "dec_swa_lock_only drops the lower-tier Mamba lock" ) self.assertGreaterEqual(full_cd.lock_ref, 1) - self.assertTrue(cache.lru_lists[ComponentType.SWA].in_list(node_a)) + self.assertTrue(cache.tree_core.lru_lists[ComponentType.SWA].in_list(node_a)) # Evict the child branch (Full/device eviction only) → A becomes a # Full-locked leaf with its now-unlocked SWA still in the SWA LRU. We do @@ -1355,7 +1403,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(node_a.children), 0, "A should now be a leaf") self.assertGreaterEqual(full_cd.lock_ref, 1, "Full must stay locked") self.assertTrue( - cache.lru_lists[ComponentType.SWA].in_list(node_a), + cache.tree_core.lru_lists[ComponentType.SWA].in_list(node_a), "A's unlocked SWA stays in the LRU (not tombstoned at transition)", ) @@ -1369,7 +1417,7 @@ class UnifiedRadixCacheSuite: self.assertIsNone(swa_cd.value, "A's SWA was freed by its own eviction") cache.sanity_check() - cache.dec_lock_ref(node_a, DecLockRefParams(swa_uuid_for_lock=None)) + cache.dec_lock_ref(node_a.id, DecLockRefParams(swa_uuid_for_lock=None)) cache.sanity_check() def test_swa_early_release_drops_co_located_mamba_lock(self): @@ -1380,9 +1428,11 @@ class UnifiedRadixCacheSuite: n_short = (self.cfg.sliding_window_size // self.cfg.page_size) + 4 seq_a = self._make_seq(1, n_short) self._insert(cache, allocator, req_to_token_pool, seq_a) - node_a = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq_a))) - ).last_device_node + node_a = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_a))) + ).last_device_node + ) self.assertEqual(len(node_a.children), 0, "A must be a leaf") swa_cd = node_a.component_data[ComponentType.SWA] @@ -1391,7 +1441,7 @@ class UnifiedRadixCacheSuite: self.assertIsNotNone(mamba_cd.value, "A must hold a Mamba checkpoint") # Natural lock acquisition — records inc_lock_ref in the lock trace. - lock_result = cache.inc_lock_ref(node_a) + lock_result = cache.inc_lock_ref(node_a.id) self.assertGreaterEqual(swa_cd.lock_ref, 1, "SWA locked") self.assertGreaterEqual(mamba_cd.lock_ref, 1, "Mamba locked") self.assertGreaterEqual(full_cd.lock_ref, 1, "Full locked") @@ -1399,7 +1449,7 @@ class UnifiedRadixCacheSuite: # Early SWA release (decode advanced past the window), via the public # path the scheduler calls. The leaf's SWA is tombstoned and the # co-located lower-tier Mamba lock must drop in the same release. - cache.dec_swa_lock_only(node_a, lock_result.swa_uuid_for_lock) + cache.dec_swa_lock_only(node_a.id, lock_result.swa_uuid_for_lock) self.assertEqual(swa_cd.lock_ref, 0, "SWA early-released") self.assertEqual( mamba_cd.lock_ref, @@ -1419,9 +1469,11 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq_a) self._insert(cache, allocator, req_to_token_pool, seq_ab) - node_a = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq_a))) - ).last_device_node + node_a = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_a))) + ).last_device_node + ) self.assertGreater(len(node_a.children), 0, "A must be internal") mamba_cd = node_a.component_data[ComponentType.MAMBA] @@ -1437,16 +1489,26 @@ class UnifiedRadixCacheSuite: self.assertEqual(full_cd.lock_ref, 0, "Full unlocked") tracker = {ct: 0 for ct in cache.tree_components} - cache._evict_component_and_detach_lru( + device_frees = defaultdict(list) + cache.tree_core._evict_component_and_detach_lru( node_a, cache.components[ComponentType.SWA], target=EvictLayer.DEVICE, tracker=tracker, + device_frees=device_frees, + host_frees=defaultdict(list), ) + cache._drain_device_frees(device_frees) # No higher-or-equal tier pins the node, so even with early-release on # the stranded Mamba lock must trip the hard-invariant assert. with self.assertRaises(AssertionError): - cache._cascade_evict(node_a, cache.components[ComponentType.SWA], tracker) + cache.tree_core._cascade_evict( + node_a, + cache.components[ComponentType.SWA], + tracker, + device_frees=defaultdict(list), + host_frees=defaultdict(list), + ) # Clean up the forced lock so teardown/sanity is consistent. cache.components[ComponentType.MAMBA].release_component_lock( @@ -1462,9 +1524,11 @@ class UnifiedRadixCacheSuite: seq_a = self._make_seq(1, n_short) self._insert(cache, allocator, req_to_token_pool, seq_a) - node_a = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq_a))) - ).last_device_node + node_a = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_a))) + ).last_device_node + ) self.assertEqual(len(node_a.children), 0, "A must be a leaf") mamba_cd = node_a.component_data[ComponentType.MAMBA] @@ -1479,16 +1543,26 @@ class UnifiedRadixCacheSuite: self.assertEqual(full_cd.lock_ref, 0, "Full unlocked") tracker = {ct: 0 for ct in cache.tree_components} - cache._evict_component_and_detach_lru( + device_frees = defaultdict(list) + cache.tree_core._evict_component_and_detach_lru( node_a, cache.components[ComponentType.SWA], target=EvictLayer.DEVICE, tracker=tracker, + device_frees=device_frees, + host_frees=defaultdict(list), ) + cache._drain_device_frees(device_frees) # No higher-or-equal tier pins the node, so even with early-release on # the stranded Mamba lock must trip the hard-invariant assert. with self.assertRaises(AssertionError): - cache._cascade_evict(node_a, cache.components[ComponentType.SWA], tracker) + cache.tree_core._cascade_evict( + node_a, + cache.components[ComponentType.SWA], + tracker, + device_frees=defaultdict(list), + host_frees=defaultdict(list), + ) # Clean up the forced lock so teardown/sanity is consistent. cache.components[ComponentType.MAMBA].release_component_lock( @@ -1507,13 +1581,15 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq_a) self._insert(cache, allocator, req_to_token_pool, seq_ab) - node_a = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq_a))) - ).last_device_node + node_a = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_a))) + ).last_device_node + ) self.assertGreater(len(node_a.children), 0, "A must have children") self.assertFalse( - cache._is_device_leaf(node_a), + cache.tree_core._is_device_leaf(node_a), "A is not a device-leaf while child holds Full on device", ) @@ -1533,19 +1609,23 @@ class UnifiedRadixCacheSuite: self._simulate_backup(cache, desc) self.assertTrue(desc.backuped, "desc must be backuped before demote") tracker = {ct: 0 for ct in cache.tree_components} - cache._evict_to_host(desc, tracker) + device_frees = defaultdict(list) + cache.tree_core._demote( + desc, tracker, device_frees=device_frees, host_frees=defaultdict(list) + ) + cache._drain_device_frees(device_frees) self.assertTrue(desc.evicted, "desc should be D->H demoted") self.assertIsNone(desc.component_data[ComponentType.FULL].value) self.assertTrue( - cache._is_device_leaf(node_a), + cache.tree_core._is_device_leaf(node_a), "A is a HiCache device-leaf (no child with Full on device)", ) self.assertGreater(len(node_a.children), 0, "A still has tree-children") - self.assertIn(node_a, cache.evictable_device_leaves) + self.assertIn(node_a, cache.tree_core.evictable_device_leaves) cache.sanity_check() - lock_result = cache.inc_lock_ref(node_a) + lock_result = cache.inc_lock_ref(node_a.id) swa_cd = node_a.component_data[ComponentType.SWA] mamba_cd = node_a.component_data[ComponentType.MAMBA] full_cd = node_a.component_data[ComponentType.FULL] @@ -1553,7 +1633,7 @@ class UnifiedRadixCacheSuite: self.assertGreaterEqual(mamba_cd.lock_ref, 1) self.assertGreaterEqual(full_cd.lock_ref, 1) - cache.dec_swa_lock_only(node_a, lock_result.swa_uuid_for_lock) + cache.dec_swa_lock_only(node_a.id, lock_result.swa_uuid_for_lock) self.assertEqual(swa_cd.lock_ref, 0, "SWA released") self.assertEqual(mamba_cd.lock_ref, 0, "Mamba dropped by dec_swa_lock_only") self.assertGreaterEqual(full_cd.lock_ref, 1, "Full kept by contract") @@ -1562,14 +1642,14 @@ class UnifiedRadixCacheSuite: "SWA slot stays under contract (lazy reclaim by drive_eviction)", ) self.assertTrue( - cache.lru_lists[ComponentType.SWA].in_list(node_a), + cache.tree_core.lru_lists[ComponentType.SWA].in_list(node_a), "SWA stays in LRU for drive_eviction to pick later", ) cache.dec_lock_ref( - node_a, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True + node_a.id, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True ) - self.assertTrue(cache._is_device_leaf(node_a)) + self.assertTrue(cache.tree_core._is_device_leaf(node_a)) cache.sanity_check() def test_swa_leaf_capped_to_window_on_insert(self): @@ -1591,9 +1671,11 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) cache.sanity_check() - leaf = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq))) - ).last_device_node + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).last_device_node + ) swa_val = leaf.component_data[ComponentType.SWA].value self.assertIsNotNone(swa_val) @@ -1606,19 +1688,19 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(swa_val), len(seq)) self.assertIs(leaf.parent, cache.root_node) - lock_result = cache.inc_lock_ref(leaf) + lock_result = cache.inc_lock_ref(leaf.id) # SWA pins one window; full attention pins everything. self.assertEqual(cache.swa_protected_size(), len(swa_val)) self.assertEqual(cache.full_protected_size(), len(seq)) cache.sanity_check() cache.dec_lock_ref( - leaf, + leaf.id, DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), ) cache.sanity_check() def _swa_lru_order(self, cache): - lru = cache.lru_lists[ComponentType.SWA] + lru = cache.tree_core.lru_lists[ComponentType.SWA] pt = lru._pt nodes: list = [] cur = lru.head.lru_next[pt] @@ -1782,7 +1864,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, pre_len) req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req.kv_committed_len = pre_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -1816,6 +1898,42 @@ class UnifiedRadixCacheSuite: ) cache.sanity_check() + def test_swa_lru_fresh_leaf_cap_rebuilds_both_nodes(self): + if not self._swa_pinning_cfg_supported(): + self.skipTest("requires SWA-only config with node size >= cushion") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + + # A single fresh in-window leaf longer than the cushion is cap-split into + # an older prefix + a one-window tail; both are rebuilt via SWARebuild. + seq = self._make_seq(1, 8) + self._insert(cache, allocator, req_to_token_pool, seq) + + order = self._swa_lru_order(cache) + self.assertEqual(len(order), 2) + tail, prefix = order[0], order[1] + + # Both nodes received an SWA value at apply time. + tail_swa = tail.component_data[ComponentType.SWA].value + prefix_swa = prefix.component_data[ComponentType.SWA].value + self.assertIsNotNone(tail_swa) + self.assertIsNotNone(prefix_swa) + + # The in-window tail leaf is more-MRU (order[0]) than its older prefix + # parent (order[1]); the tail is capped to one cushion. + cushion = self.cfg.sliding_window_size + self.cfg.page_size + self.assertEqual(len(tail.children), 0) + self.assertIs(prefix, tail.parent) + self.assertLess(len(tail.key), cushion) + self.assertEqual(len(tail.key) + len(prefix.key), len(seq)) + + # Rebuilt SWA spans the whole leaf, and evictable SWA size matches it. + self.assertEqual(len(tail_swa) + len(prefix_swa), len(seq)) + self.assertEqual( + cache.tree_core.component_evictable_size_[ComponentType.SWA], len(seq) + ) + + cache.sanity_check() + def test_swa_eager_eviction_noop_when_within_window(self): if not self.cfg.has_swa or self.cfg.has_mamba: self.skipTest("requires SWA without Mamba") @@ -1834,7 +1952,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, pre_len) req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req.kv_committed_len = pre_len - req.last_node = cache.root_node + req.last_node = cache.root_node.id req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -1891,7 +2009,11 @@ class UnifiedRadixCacheSuite: tracker = {ct: 0 for ct in cache.tree_components} - cache._iteratively_delete_tombstone_leaf(deleted, tracker) + device_frees = defaultdict(list) + cache.tree_core._iteratively_delete_tombstone_leaf( + deleted, tracker, device_frees, defaultdict(list) + ) + cache._drain_device_frees(device_frees) self.assertIn(parent_key, cache.root_node.children) self.assertIs(cache.root_node.children[parent_key], parent) @@ -1965,22 +2087,22 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) match = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = match.last_device_node + node = cache.resolve_node_handle(match.last_device_node) full_cd = node.component_data[ComponentType.FULL] aux_cd = node.component_data[aux] self.assertEqual(len(node.children), 0) self.assertIsNotNone(full_cd.value) self.assertIsNotNone(aux_cd.value) - lock_result = cache.inc_lock_ref(node) + lock_result = cache.inc_lock_ref(node.id) self.assertGreater(full_cd.lock_ref, 0) self.assertGreater(aux_cd.lock_ref, 0) aux_len = len(aux_cd.value) - cache.component_protected_size_[aux] -= aux_len - cache.component_evictable_size_[aux] += aux_len + cache.tree_core.component_protected_size_[aux] -= aux_len + cache.tree_core.component_evictable_size_[aux] += aux_len aux_cd.lock_ref = 0 - self.assertNotIn(node, cache.evictable_device_leaves) + self.assertNotIn(node, cache.tree_core.evictable_device_leaves) evict_params = EvictParams(num_tokens=0) if aux == ComponentType.SWA: @@ -1996,10 +2118,10 @@ class UnifiedRadixCacheSuite: self.assertEqual(result.mamba_num_evicted, aux_len) self.assertIsNotNone(full_cd.value) self.assertIsNone(aux_cd.value) - self.assertFalse(cache.lru_lists[aux].in_list(node)) + self.assertFalse(cache.tree_core.lru_lists[aux].in_list(node)) cache.dec_lock_ref( - node, + node.id, DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), ) cache.sanity_check() @@ -2047,7 +2169,10 @@ class UnifiedRadixCacheSuite: ), ) # After unlock, base should be in evictable_device_leaves - self.assertIn(m_base.last_device_node, cache.evictable_device_leaves) + self.assertIn( + cache.resolve_node_handle(m_base.last_device_node), + cache.tree_core.evictable_device_leaves, + ) cache.sanity_check() def test_evict_iterative_tombstone_cleanup(self): @@ -2252,7 +2377,7 @@ class UnifiedRadixCacheSuite: def _write_path_to_l3(self, cache, node): """Offload every node on root->node path from host to L3 storage.""" for n in self._path_chain(cache, node): - cache.write_backup_storage(n) + cache.write_backup_storage(n.id) def _flush_l3_backups(self, cache, timeout: float = 10.0): """Wait for backup threads to finish, then drain acks (release locks).""" @@ -2303,7 +2428,7 @@ class UnifiedRadixCacheSuite: seq = self._make_seq(1, 4) self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - leaf = m.last_device_node + leaf = cache.resolve_node_handle(m.last_device_node) # D->H first, then H->L3. self._backup_node(cache, leaf) @@ -2352,7 +2477,7 @@ class UnifiedRadixCacheSuite: ) self._insert(prod, prod_alloc, prod_rtp, seq) mp = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - prod_leaf = mp.last_device_node + prod_leaf = prod.resolve_node_handle(mp.last_device_node) self._fill_full_kv(prod_alloc, mp.device_indices, marker=7) expected_k, expected_v = self._snapshot_full_kv(prod_alloc, mp.device_indices) self._backup_node(prod, prod_leaf) @@ -2368,14 +2493,16 @@ class UnifiedRadixCacheSuite: prefetch_threshold=1, ) req_id = "l3-prefetch-req" - cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None) + cons.prefetch_from_storage( + req_id, cons.root_node.id, array("q", seq), None, None + ) self._run_prefetch_to_completion(cons, req_id) cons.drain_storage_control_queues() # The full prefix must now be a host hit (loaded from L3). mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(mc.host_hit_length, len(seq)) - host_node = mc.last_host_node + host_node = cons.resolve_node_handle(mc.last_host_node) self.assertIsNot(host_node, cons.root_node) self.assertTrue(host_node.evicted) @@ -2424,7 +2551,7 @@ class UnifiedRadixCacheSuite: def _swa_host_on_path(self, cache, seq): m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_host_node + node = cache.resolve_node_handle(m.last_host_node) while node is not cache.root_node: if node.component_data[ComponentType.SWA].host_value is not None: return True @@ -2437,9 +2564,11 @@ class UnifiedRadixCacheSuite: prod, storage_backend="file", storage_dir=storage_dir, prefetch_threshold=1 ) self._insert(prod, prod_alloc, prod_rtp, seq) - leaf = prod.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq))) - ).last_device_node + leaf = prod.resolve_node_handle( + prod.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).last_device_node + ) self._backup_node(prod, leaf) self._write_path_to_l3(prod, leaf) self._flush_l3_backups(prod) @@ -2452,7 +2581,9 @@ class UnifiedRadixCacheSuite: return cons def _consume_prefetch(self, cons, seq, req_id): - cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None) + cons.prefetch_from_storage( + req_id, cons.root_node.id, array("q", seq), None, None + ) self._run_prefetch_to_completion(cons, req_id) cons.drain_storage_control_queues() @@ -2529,6 +2660,115 @@ class UnifiedRadixCacheSuite: self.assertIn(2, min_sizes) cons.sanity_check() + def test_tp_swa_prefetch_drop_frees_host_pool(self): + """A dropped SWA prefetch must return its whole host buffer to the pool + (no leak, no over-free).""" + setup = self._setup_swa_tp_prefetch() + if setup is None: + return + storage_dir, seq = setup + + cons = self._l3_consumer(storage_dir) + cons.tp_world_size = 2 + self._patch_tp_all_reduce(cons, drop_swa=True) + avail_before = cons.swa_kv_pool_host.available_size() + self._consume_prefetch(cons, seq, "drop") + + self.assertEqual( + cons.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).host_hit_length, + 0, + ) + # Whole window dropped -> its host buffer is fully released back. + self.assertEqual(cons.swa_kv_pool_host.available_size(), avail_before) + + def test_hicache_write_back_evict_drops_unbacked_leaf_when_host_full(self): + """Write-back eviction will keep freeing device KV when the host pool + is exhausted and host eviction cannot free space to prevent OOM.""" + if self._skip_unsupported_hicache_test(): + return + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_hicache(cache, write_policy="write_back") + + # Two-node chain: the drop must cascade parent-ward across eviction + # iterations, not just delete a single leaf. + seq_parent = self._make_seq(1, 2) + self._insert(cache, allocator, req_to_token_pool, seq_parent) + seq = seq_parent + self._make_seq(1000, 1) + self._insert(cache, allocator, req_to_token_pool, seq) + + # Exhaust the KV host pool. The tree has no host leaves, so + # evict_host cannot free anything and every D->H backup fails. + host_pool = cache.cache_controller.mem_pool_host + self.assertIsNotNone(host_pool.alloc(host_pool.available_size())) + self.assertEqual(host_pool.available_size(), 0) + + result = cache.evict(EvictParams(num_tokens=len(seq))) + self.assertGreaterEqual(result.num_tokens_evicted, len(seq)) + + # The chain is gone entirely: no device hit, no host hit. + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(len(m.device_indices), 0) + self.assertEqual(m.host_hit_length, 0) + cache.sanity_check() + + def test_hicache_write_back_drop_respects_pins_then_frees_subtree(self): + """The host-pressure drop fallback must decline while any node in the + unbacked subtree is host-pinned, then reclaim the whole subtree -- + including a demoted child's host backup -- once unpinned.""" + if self._skip_unsupported_hicache_test(): + return + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_hicache(cache, write_policy="write_back") + host_pool = cache.cache_controller.mem_pool_host + baseline_host = host_pool.available_size() + + parent_seq = self._make_seq(1, 2) + self._insert(cache, allocator, req_to_token_pool, parent_seq) + child_seq = parent_seq + self._make_seq(1000, 1) + self._insert(cache, allocator, req_to_token_pool, child_seq) + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq)))) + child = cache.resolve_node_handle(m.last_device_node) + + # Evict only the child leaf -> real backup + demote, leaving it + # host-only under a still-unbacked device parent (write-back backs up + # single nodes leaf-first, so this is a normal intermediate state). + result = cache.evict(EvictParams(num_tokens=len(child.key))) + self.assertGreaterEqual(result.num_tokens_evicted, len(child.key)) + self.assertTrue(child.evicted and child.backuped) + parent = child.parent + self.assertFalse(parent.backuped) + self.assertGreater(baseline_host - host_pool.available_size(), 0) + + # From here every backup fails (controller.write returns None), so + # each evict() attempts the drop fallback on the parent. + with mock.patch.object(cache.cache_controller, "write", return_value=None): + # Pinned subtree root: drop declines, chain stays intact. + cache.inc_host_lock_ref(parent.id) + result = cache.evict(EvictParams(num_tokens=len(parent_seq))) + self.assertEqual(result.num_tokens_evicted, 0) + cache.dec_host_lock_ref(parent.id) + + # Pinned host-only descendant: drop declines as well. + cache.inc_host_lock_ref(child.id) + result = cache.evict(EvictParams(num_tokens=len(parent_seq))) + self.assertEqual(result.num_tokens_evicted, 0) + m = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", parent_seq))) + ) + self.assertEqual(len(m.device_indices), len(parent_seq)) + cache.dec_host_lock_ref(child.id) + + # Unpinned: the subtree drops and the child's host slots return. + result = cache.evict(EvictParams(num_tokens=len(parent_seq))) + self.assertGreaterEqual(result.num_tokens_evicted, len(parent_seq)) + self.assertEqual(host_pool.available_size(), baseline_host) + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq)))) + self.assertEqual(len(m.device_indices), 0) + self.assertEqual(m.host_hit_length, 0) + cache.sanity_check() + def _skip_unsupported_hicache_test(self): if self.cfg.has_swa and self.cfg.has_mamba: self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks") @@ -2673,7 +2913,7 @@ class UnifiedRadixCacheSuite: for ancestor in reversed(chain): if ancestor.backuped: continue - backed_up = cache.write_backup(ancestor, write_back=True) + backed_up = _write_backup(cache, ancestor, write_back=True) self.assertGreater(backed_up, 0) cache.writing_check(write_back=True) self.assertTrue(node.backuped) @@ -2689,7 +2929,7 @@ class UnifiedRadixCacheSuite: self._backup_node(cache, node) def _load_back_node(self, cache, node): - loaded = cache.load_back(node) + loaded = cache.load_back(node.id) self.assertTrue(loaded) producer_id = cache.ready_to_load_host_cache() self.assertNotEqual(producer_id, -1) @@ -2735,94 +2975,37 @@ class UnifiedRadixCacheSuite: [conv[:, mamba_indices].float().cpu().clone() for conv in mamba_cache.conv], ) - def test_hicache_write_back_evict_drops_unbacked_leaf_when_host_full(self): - """Write-back eviction will keep freeing device KV when the host pool - is exhausted and host eviction cannot free space to prevent OOM.""" + def test_hicache_evict_device_leaf_aborts_demote_when_backup_fails(self): + """The tree op defers a write-back victim's demote: it returns a BackupKV + without freeing device, and _demote requires a completed backup, so a + failed backup leaves the node recoverable on device.""" if self._skip_unsupported_hicache_test(): return cache, allocator, req_to_token_pool = build_fixture(self.cfg) self._init_hicache(cache, write_policy="write_back") + ct = ComponentType.FULL - # Two-node chain: the drop must cascade parent-ward across eviction - # iterations, not just delete a single leaf. - seq_parent = self._make_seq(1, 2) - self._insert(cache, allocator, req_to_token_pool, seq_parent) - seq = seq_parent + self._make_seq(1000, 1) + seq = self._make_seq(1, 2) self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node + node = cache.resolve_node_handle(m.last_device_node) self.assertIsNot(node, cache.root_node) self.assertFalse(node.backuped) + self.assertFalse(node.evicted) - # Exhaust the KV host pool. The tree has no host leaves, so - # evict_host cannot free anything and every D->H backup fails. - host_pool = cache.cache_controller.mem_pool_host - self.assertIsNotNone(host_pool.alloc(host_pool.available_size())) - self.assertEqual(host_pool.available_size(), 0) + # Tree op defers: returns a BackupKV and leaves the node on device (no demote). + leaf_result = cache.tree_core.evict_device_leaf(node.id, is_write_back=True) + self.assertIsNotNone(leaf_result.backup_kv) + cache._free_values(leaf_result.device_frees, leaf_result.host_frees) - result = cache.evict(EvictParams(num_tokens=len(seq))) - self.assertGreaterEqual(result.num_tokens_evicted, len(seq)) + self.assertFalse(node.evicted) + self.assertIsNotNone(node.component_data[ct].value) + self.assertIsNone(node.component_data[ct].host_value) - # The chain is gone entirely: no device hit, no host hit. - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - self.assertEqual(len(m.device_indices), 0) - self.assertEqual(m.host_hit_length, 0) - cache.sanity_check() + # Demote needs a completed backup: _demote asserts on the un-backed node. + with self.assertRaises(AssertionError): + cache.tree_core.demote(node.id) - def test_hicache_write_back_drop_respects_pins_then_frees_subtree(self): - """The host-pressure drop fallback must decline while any node in the - unbacked subtree is host-pinned, then reclaim the whole subtree -- - including a demoted child's host backup -- once unpinned.""" - if self._skip_unsupported_hicache_test(): - return - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - self._init_hicache(cache, write_policy="write_back") - host_pool = cache.cache_controller.mem_pool_host - baseline_host = host_pool.available_size() - - parent_seq = self._make_seq(1, 2) - self._insert(cache, allocator, req_to_token_pool, parent_seq) - child_seq = parent_seq + self._make_seq(1000, 1) - self._insert(cache, allocator, req_to_token_pool, child_seq) - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq)))) - child = m.last_device_node - - # Evict only the child leaf -> real backup + demote, leaving it - # host-only under a still-unbacked device parent (write-back backs up - # single nodes leaf-first, so this is a normal intermediate state). - result = cache.evict(EvictParams(num_tokens=len(child.key))) - self.assertGreaterEqual(result.num_tokens_evicted, len(child.key)) - self.assertTrue(child.evicted and child.backuped) - parent = child.parent - self.assertFalse(parent.backuped) - self.assertGreater(baseline_host - host_pool.available_size(), 0) - - # From here every backup fails (controller.write returns None), so - # each evict() attempts the drop fallback on the parent. - with mock.patch.object(cache.cache_controller, "write", return_value=None): - # Pinned subtree root: drop declines, chain stays intact. - cache.inc_host_lock_ref(parent) - result = cache.evict(EvictParams(num_tokens=len(parent_seq))) - self.assertEqual(result.num_tokens_evicted, 0) - cache.dec_host_lock_ref(parent) - - # Pinned host-only descendant: drop declines as well. - cache.inc_host_lock_ref(child) - result = cache.evict(EvictParams(num_tokens=len(parent_seq))) - self.assertEqual(result.num_tokens_evicted, 0) - m = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", parent_seq))) - ) - self.assertEqual(len(m.device_indices), len(parent_seq)) - cache.dec_host_lock_ref(child) - - # Unpinned: the subtree drops and the child's host slots return. - result = cache.evict(EvictParams(num_tokens=len(parent_seq))) - self.assertGreaterEqual(result.num_tokens_evicted, len(parent_seq)) - self.assertEqual(host_pool.available_size(), baseline_host) - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq)))) - self.assertEqual(len(m.device_indices), 0) - self.assertEqual(m.host_hit_length, 0) cache.sanity_check() def test_hicache_evict_to_host(self): @@ -2834,7 +3017,7 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node + node = cache.resolve_node_handle(m.last_device_node) self._backup_node(cache, node) self.assertTrue(node.backuped) @@ -2850,8 +3033,8 @@ class UnifiedRadixCacheSuite: self.assertIsNotNone(node.component_data[ComponentType.FULL].host_value) # Should be in host_leaves, not device_leaves - self.assertNotIn(node, cache.evictable_device_leaves) - self.assertIn(node, cache.evictable_host_leaves) + self.assertNotIn(node, cache.tree_core.evictable_device_leaves) + self.assertIn(node, cache.tree_core.evictable_host_leaves) cache.sanity_check() def test_hicache_match_through_evicted_node(self): @@ -2894,7 +3077,7 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node + node = cache.resolve_node_handle(m.last_device_node) self._backup_node(cache, node) cache.evict(EvictParams(num_tokens=len(seq))) @@ -2904,16 +3087,16 @@ class UnifiedRadixCacheSuite: m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", query)))) self.assertEqual(len(m.device_indices), 0) - self.assertIs(m.last_device_node, cache.root_node) + self.assertIs(cache.resolve_node_handle(m.last_device_node), cache.root_node) # Locate the host prefix via last_host_node and rebuild prefix/suffix # from path keys (a leaf may span several nodes). if self.cfg.has_mamba: self.assertEqual(m.host_hit_length, 0) - self.assertIs(m.last_host_node, cache.root_node) + self.assertIs(cache.resolve_node_handle(m.last_host_node), cache.root_node) else: self.assertEqual(m.host_hit_length, len(expected_prefix)) - split_parent = m.last_host_node + split_parent = cache.resolve_node_handle(m.last_host_node) self.assertIsNot(split_parent, cache.root_node) self.assertTrue(split_parent.evicted) self.assertTrue(split_parent.backuped) @@ -2947,22 +3130,48 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node + node = cache.resolve_node_handle(m.last_device_node) self._backup_node(cache, node) cache.evict(EvictParams(num_tokens=len(seq))) self.assertTrue(node.evicted) - self.assertIn(node, cache.evictable_host_leaves) + self.assertIn(node, cache.tree_core.evictable_host_leaves) # Now evict host cache.evict_host(len(seq)) # Node should be removed from tree - self.assertNotIn(node, cache.evictable_host_leaves) + self.assertNotIn(node, cache.tree_core.evictable_host_leaves) self.assertEqual(len(cache.root_node.children), 0) cache.sanity_check() + def test_hicache_evict_keeps_node_on_device_when_backup_fails(self): + """evict(): a failed write-back backup skips the demote, leaving the node + device-resident and recoverable.""" + if self._skip_unsupported_hicache_test(): + return + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_hicache(cache, write_policy="write_back") + ct = ComponentType.FULL + + seq = self._make_seq(1, 2) + self._insert(cache, allocator, req_to_token_pool, seq) + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + node = cache.resolve_node_handle(m.last_device_node) + self.assertIsNot(node, cache.root_node) + self.assertFalse(node.backuped) + self.assertFalse(node.evicted) + + # Backup IO fails (returns 0): evict() must skip the demote. + with mock.patch.object(cache, "_execute_and_commit_kv_backup", return_value=0): + cache.evict(EvictParams(num_tokens=len(seq))) + + self.assertFalse(node.evicted) + self.assertIsNotNone(node.component_data[ct].value) + self.assertIsNone(node.component_data[ct].host_value) + cache.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(): @@ -2972,7 +3181,7 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, base) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", base)))) - node = m.last_device_node + node = cache.resolve_node_handle(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( @@ -3031,7 +3240,7 @@ class UnifiedRadixCacheSuite: self._backup_tree(cache) # Verify: every backed-up node's parent is also backed-up (or root) - all_nodes = cache._collect_all_nodes() + all_nodes = cache.tree_core._collect_all_nodes() for node in all_nodes: if node is cache.root_node: continue @@ -3075,7 +3284,7 @@ class UnifiedRadixCacheSuite: cache.evict(EvictParams(num_tokens=len(seq))) self.assertTrue(split_leaf.evicted) self.assertTrue(split_leaf.backuped) - self.assertIn(split_leaf, cache.evictable_host_leaves) + self.assertIn(split_leaf, cache.tree_core.evictable_host_leaves) cache.sanity_check() def test_swa_deep_tree_backup_evict_loadback_stress(self): @@ -3120,7 +3329,7 @@ class UnifiedRadixCacheSuite: for i in range(2): # width: branches off the base prefix insert_swa(base[: 2 * ps] + self._make_seq(80000 + 1000 * i, 3), 0) - self.assertGreaterEqual(len(cache._collect_all_nodes()), 5) + self.assertGreaterEqual(len(cache.tree_core._collect_all_nodes()), 5) # Stepwise eviction -> demote to host, sanity after each round. for _ in range(4): @@ -3138,9 +3347,9 @@ class UnifiedRadixCacheSuite: # Load evicted prefixes back from host, sanity after each. for tokens in (base, base[: 2 * ps]): m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - anchor = m.best_match_node + anchor = cache.resolve_node_handle(m.best_match_node) if anchor is not cache.root_node and anchor.evicted: - if cache.load_back(anchor): + if cache.load_back(anchor.id): self._finish_pending_loads(cache) self._release_ongoing_load_back_locks(cache) cache.sanity_check() @@ -3162,19 +3371,19 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node + node = cache.resolve_node_handle(m.last_device_node) for aux in aux_types: - self.assertTrue(cache.lru_lists[aux].in_list(node)) - self.assertFalse(cache.host_lru_lists[aux].in_list(node)) + self.assertTrue(cache.tree_core.lru_lists[aux].in_list(node)) + self.assertFalse(cache.tree_core.host_lru_lists[aux].in_list(node)) self._simulate_backup(cache, node) cache.evict(EvictParams(num_tokens=len(seq))) for aux in aux_types: - self.assertFalse(cache.lru_lists[aux].in_list(node)) + self.assertFalse(cache.tree_core.lru_lists[aux].in_list(node)) if node.component_data[aux].host_value is not None: - self.assertTrue(cache.host_lru_lists[aux].in_list(node)) + self.assertTrue(cache.tree_core.host_lru_lists[aux].in_list(node)) cache.sanity_check() def _build_chain_pages(self, cache, allocator, req_to_token_pool, num_pages): @@ -3189,7 +3398,7 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) chain: list = [] - cur = m.last_device_node + cur = cache.resolve_node_handle(m.last_device_node) while cur is not cache.root_node: chain.append(cur) cur = cur.parent @@ -3200,8 +3409,8 @@ class UnifiedRadixCacheSuite: for node, lock_params, host_lock_params in list( cache.ongoing_load_back.values() ): - cache.dec_lock_ref(node, lock_params) - cache.dec_host_lock_ref(node, host_lock_params) + cache.dec_lock_ref(node.id, lock_params) + cache.dec_host_lock_ref(node.id, host_lock_params) cache.ongoing_load_back.clear() def _finish_pending_loads(self, cache): @@ -3224,12 +3433,12 @@ class UnifiedRadixCacheSuite: cd.host_value = cd.value.clone() old_value = cd.value cd.value = None - if component_type in cache.lru_lists and cache.lru_lists[ + if component_type in cache.tree_core.lru_lists and cache.tree_core.lru_lists[ component_type ].in_list(node): - cache.lru_lists[component_type].remove_node(node) - cache.host_lru_lists[component_type].insert_mru(node) - cache.component_evictable_size_[component_type] -= len(old_value) + cache.tree_core.lru_lists[component_type].remove_node(node) + cache.tree_core.host_lru_lists[component_type].insert_mru(node) + cache.tree_core.component_evictable_size_[component_type] -= len(old_value) def test_match_prefix_best_and_device_node_without_hicache(self): cache, allocator, req_to_token_pool = build_fixture(self.cfg) @@ -3243,10 +3452,96 @@ class UnifiedRadixCacheSuite: result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(len(result.device_indices), len(seq)) - self.assertIs(result.best_match_node, result.last_device_node) - self.assertIs(result.last_host_node, result.last_device_node) + self.assertIs( + cache.resolve_node_handle(result.best_match_node), + cache.resolve_node_handle(result.last_device_node), + ) + self.assertIs( + cache.resolve_node_handle(result.last_host_node), + cache.resolve_node_handle(result.last_device_node), + ) self.assertEqual(result.host_hit_length, 0) + def test_full_kv_hit_length_counts_the_split_fragment(self): + """A mid-node partial match splits the node; the matched fragment still + counts toward full_kv_hit_length (independent of component validators).""" + if self.cfg.page_size != 1: + self.skipTest("page_size=1 keeps the split boundary mid-node") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._insert(cache, allocator, req_to_token_pool, list(range(1, 9))) + result = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3, 4, 99]))) + ) + self.assertEqual(result.full_kv_hit_length, 4) + + def test_has_swa_host_pool_flag_matches_attached_pool(self): + """init_hicache caches SWA host-pool presence on the tree core after + the pool assembler runs; the flag must agree with the attached handle.""" + if not self.cfg.has_swa: + self.skipTest("requires SWA") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self.assertFalse(cache.tree_core.has_swa_host_pool) + if self._skip_unsupported_hicache_test(): + return + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + swa = cache.components[ComponentType.SWA] + self.assertEqual( + cache.tree_core.has_swa_host_pool, swa._swa_kv_pool_host is not None + ) + + def test_zero_match_result_carries_node_id_handles(self): + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + ps = self.cfg.page_size + min_tokens = 2 * ps + if self.cfg.has_swa: + min_tokens = max(min_tokens, self.cfg.sliding_window_size + ps) + seq = self._make_seq(1, (min_tokens + ps - 1) // ps) + self._insert(cache, allocator, req_to_token_pool, seq) + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + + zeroed = zero_match_result(cache, result) + + self.assertEqual(len(zeroed.device_indices), 0) + # zeroed results must carry the root's NodeId, not the raw node + self.assertEqual(zeroed.best_match_node, cache.root_node.id) + # The env-gated force-miss path passes the request's extra key. + salted = zero_match_result(cache, result, extra_key="salt") + self.assertEqual(salted.best_match_node, cache.root_node.id) + self.assertEqual(zeroed.last_device_node, cache.root_node.id) + self.assertEqual(zeroed.last_host_node, cache.root_node.id) + # and the handles must work with the NodeId-based lock APIs + lock = cache.inc_lock_ref(zeroed.best_match_node) + cache.dec_lock_ref(zeroed.best_match_node, lock.to_dec_params()) + + def test_evict_host_drains_freed_host_values_to_the_pools(self): + if self.cfg.has_swa and self.cfg.has_mamba: + self.skipTest("no hicache strategy covers FULL+SWA+MAMBA") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + + host_pool_attr = { + ComponentType.FULL: "full_kv_pool_host", + ComponentType.SWA: "swa_kv_pool_host", + ComponentType.MAMBA: "mamba_pool_host", + } + for ct in cache.tree_components: + # set by the hicache assembler for every enabled component + pool = getattr(cache, host_pool_attr[ct]) + if pool is None: + continue + available_before = pool.available_size() + evicted = cache.evict_host(pool.size, ct) + if evicted == 0: + continue + # the drain must return every freed host value to the pool + self.assertGreater(pool.available_size(), available_before) + def test_hicache_mamba_host_best_match_keeps_device_anchor(self): if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: self.skipTest("requires page_size=1 Full+Mamba") @@ -3264,11 +3559,126 @@ class UnifiedRadixCacheSuite: result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIs(result.best_match_node, leaf) - self.assertIs(result.last_device_node, parent) + self.assertIs(cache.resolve_node_handle(result.best_match_node), leaf) + self.assertIs(cache.resolve_node_handle(result.last_device_node), parent) self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key)) self.assertEqual(result.host_hit_length, len(leaf.key)) + def test_mamba_has_host_value_only_predicate(self): + """Needs a slot only when mamba is host-only (device evicted, host backed up).""" + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if not chain: + self.skipTest("chain too short") + node = chain[-1] + cd = node.component_data[ComponentType.MAMBA] + dev = torch.tensor([0], dtype=torch.int64) + host = torch.tensor([0], dtype=torch.int64) + + # host-only: device evicted, host backup present -> restore needs a slot + cd.value, cd.host_value = None, host + self.assertTrue( + cache.tree_core.component_has_host_value_only(node.id, ComponentType.MAMBA) + ) + # device + host: device value present -> no restore, no slot (the D+H bug case) + cd.value, cd.host_value = dev, host + self.assertFalse( + cache.tree_core.component_has_host_value_only(node.id, ComponentType.MAMBA) + ) + # device-only / neither: nothing to restore + cd.value, cd.host_value = dev, None + self.assertFalse( + cache.tree_core.component_has_host_value_only(node.id, ComponentType.MAMBA) + ) + cd.value, cd.host_value = None, None + self.assertFalse( + cache.tree_core.component_has_host_value_only(node.id, ComponentType.MAMBA) + ) + + def test_mamba_device_value_accessor(self): + """Returns the device mamba value to CoW from, or None when device-evicted.""" + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if not chain: + self.skipTest("chain too short") + node = chain[-1] + cd = node.component_data[ComponentType.MAMBA] + dev = torch.tensor([0], dtype=torch.int64) + host = torch.tensor([0], dtype=torch.int64) + + # device present -> CoW source is the device value (host backup irrelevant) + cd.value, cd.host_value = dev, None + self.assertIs( + cache.tree_core.get_component_device_value(node.id, ComponentType.MAMBA), + dev, + ) + cd.value, cd.host_value = dev, host + self.assertIs( + cache.tree_core.get_component_device_value(node.id, ComponentType.MAMBA), + dev, + ) + # device evicted -> nothing to CoW from + cd.value, cd.host_value = None, host + self.assertIsNone( + cache.tree_core.get_component_device_value(node.id, ComponentType.MAMBA) + ) + cd.value, cd.host_value = None, None + self.assertIsNone( + cache.tree_core.get_component_device_value(node.id, ComponentType.MAMBA) + ) + + def test_prepare_prefetch_swa(self): + if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+SWA") + cache, _, _ = self._build_hicache_fixture() + sw = cache.sliding_window_size + swa = cache.components[ComponentType.SWA] + # below a full window -> does not participate, no alloc + prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=sw - 1) + self.assertFalse(prep.alloc_failed) + self.assertIsNone(prep.host_indices) + # a full window available -> participates, allocs one window of host pages + prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=sw) + self.assertEqual(int(prep.host_indices.numel()), sw) + # a non-participating component never allocs + prep = cache.components[ComponentType.FULL].prepare_prefetch( + cache.root_node.id, prefetch_tokens=sw + ) + self.assertFalse(prep.alloc_failed) + self.assertIsNone(prep.host_indices) + + def test_prepare_prefetch_swa_pool_exhausted(self): + if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+SWA") + cache, _, _ = self._build_hicache_fixture() + sw = cache.sliding_window_size + swa = cache.components[ComponentType.SWA] + # pool can't satisfy even after evict -> participates but aborts (no buffer) + with mock.patch.object( + cache.swa_kv_pool_host, "alloc", return_value=None + ), mock.patch.object(cache, "evict_host", autospec=True) as evict_host: + self.assertTrue( + swa.prepare_prefetch( + cache.root_node.id, prefetch_tokens=sw + ).alloc_failed + ) + # the retry must evict the SWA host pool, not the default (FULL) one + evict_host.assert_called_once_with(sw, ComponentType.SWA) + + def test_prepare_prefetch_mamba(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, _, _ = self._build_hicache_fixture() + # mamba always participates and allocs exactly one page + prep = cache.components[ComponentType.MAMBA].prepare_prefetch( + cache.root_node.id, prefetch_tokens=0 + ) + self.assertEqual(int(prep.host_indices.numel()), 1) + def test_hicache_swa_host_best_match_keeps_device_anchor(self): if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1: self.skipTest("requires page_size=1 Full+SWA") @@ -3286,8 +3696,8 @@ class UnifiedRadixCacheSuite: result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIs(result.best_match_node, leaf) - self.assertIs(result.last_device_node, parent) + self.assertIs(cache.resolve_node_handle(result.best_match_node), leaf) + self.assertIs(cache.resolve_node_handle(result.last_device_node), parent) self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key)) self.assertEqual(result.host_hit_length, len(leaf.key)) self.assertEqual(result.swa_host_hit_length, len(leaf.key)) @@ -3304,8 +3714,8 @@ class UnifiedRadixCacheSuite: result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIs(result.best_match_node, leaf) - self.assertIs(result.last_device_node, parent) + self.assertIs(cache.resolve_node_handle(result.best_match_node), leaf) + self.assertIs(cache.resolve_node_handle(result.last_device_node), parent) self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key)) self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.swa_host_hit_length, len(leaf.key)) @@ -3317,31 +3727,41 @@ class UnifiedRadixCacheSuite: chunk_size = get_server_args().mamba_cache_chunk_size tokens = self._make_seq(1, chunk_size + 1) self._insert(cache, allocator, req_to_token_pool, tokens) - leaf = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", tokens))) - ).last_device_node + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + ) mamba_cd = leaf.component_data[ComponentType.MAMBA] mamba_cd.value = None no_hicache = cache.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens))) ) - self.assertIs(no_hicache.best_match_node, cache.root_node) - self.assertIs(no_hicache.last_device_node, cache.root_node) + self.assertIs( + cache.resolve_node_handle(no_hicache.best_match_node), cache.root_node + ) + self.assertIs( + cache.resolve_node_handle(no_hicache.last_device_node), cache.root_node + ) self.assertEqual(no_hicache.mamba_branching_seqlen, chunk_size) tree_h, allocator_h, req_to_token_pool_h = self._build_hicache_fixture() self._insert(tree_h, allocator_h, req_to_token_pool_h, tokens) - leaf_h = tree_h.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", tokens))) - ).last_device_node + leaf_h = tree_h.resolve_node_handle( + tree_h.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + ) self._backup_node(tree_h, leaf_h) tree_h.evict(EvictParams(num_tokens=len(tokens))) with_hicache = tree_h.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens))) ) - self.assertIs(with_hicache.best_match_node, leaf_h) - self.assertIs(with_hicache.last_device_node, tree_h.root_node) + self.assertIs(tree_h.resolve_node_handle(with_hicache.best_match_node), leaf_h) + self.assertIs( + tree_h.resolve_node_handle(with_hicache.last_device_node), tree_h.root_node + ) self.assertIsNone(with_hicache.mamba_branching_seqlen) def test_mamba_branching_seqlen_uses_device_full_hit_under_hicache(self): @@ -3354,16 +3774,18 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, prefix) self._insert(cache, allocator, req_to_token_pool, tokens) - leaf = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", tokens))) - ).last_device_node + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + ) parent = leaf.parent leaf.component_data[ComponentType.MAMBA].value = None result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIs(result.best_match_node, parent) - self.assertIs(result.last_device_node, parent) + self.assertIs(cache.resolve_node_handle(result.best_match_node), parent) + self.assertIs(cache.resolve_node_handle(result.last_device_node), parent) self.assertEqual(len(result.device_indices), chunk_size) self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.full_kv_hit_length, len(tokens)) @@ -3379,31 +3801,36 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, prefix) self._insert(cache, allocator, req_to_token_pool, tokens) - leaf = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", tokens))) - ).last_device_node + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + ) parent = leaf.parent self._backup_node(cache, leaf) - lock_result = cache.inc_lock_ref(parent) + lock_result = cache.inc_lock_ref(parent.id) try: cache.evict(EvictParams(num_tokens=len(leaf.key))) finally: cache.dec_lock_ref( - parent, + parent.id, DecLockRefParams( swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None) ), ) self.assertTrue(leaf.evicted) self.assertTrue(leaf.backuped) + device_frees = defaultdict(list) + host_frees = defaultdict(list) cache.components[ComponentType.MAMBA].evict_component( - leaf, target=EvictLayer.HOST + leaf, device_frees, host_frees, target=EvictLayer.HOST ) + cache._free_values(device_frees, host_frees) result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) - self.assertIs(result.best_match_node, parent) - self.assertIs(result.last_device_node, parent) + self.assertIs(cache.resolve_node_handle(result.best_match_node), parent) + self.assertIs(cache.resolve_node_handle(result.last_device_node), parent) self.assertEqual(len(result.device_indices), chunk_size) self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.full_kv_hit_length, len(tokens)) @@ -3451,262 +3878,12 @@ class UnifiedRadixCacheSuite: ) ) - self.assertIs(new_node, leaf) + self.assertIs(cache.resolve_node_handle(new_node), leaf) self.assertEqual(len(torch.cat([req.prefix_indices, new_indices])), len(tokens)) self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].value) self._finish_pending_loads(cache) self._release_ongoing_load_back_locks(cache) - def test_load_back_abort_frees_unpublished_mamba_slot(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - self.assertTrue(leaf.evicted) - self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].host_value) - - # A request whose mamba slot was released: load_back's CoW arm allocates one. - req = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - mamba_avail = req_to_token_pool.mamba_allocator.available_size() - - # Impossible quota -> load_back aborts after building the transfers. - loaded = cache.load_back(leaf, mem_quota=-(10**9), req=req) - - self.assertFalse(loaded) - # the aborted call must return its slot and not leave req pointing at it - self.assertIsNone(req.mamba_pool_idx) - self.assertEqual( - req_to_token_pool.mamba_allocator.available_size(), mamba_avail - ) - self._release_ongoing_load_back_locks(cache) - - def test_load_back_load_failure_frees_unpublished_mamba_slot(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - self.assertTrue(leaf.evicted) - - req = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - mamba_avail = req_to_token_pool.mamba_allocator.available_size() - - # cache_controller.load() failing (device alloc / transfer resolution) - # must also return the slot this call allocated. - with mock.patch.object(cache.cache_controller, "load", return_value=None): - loaded = cache.load_back(leaf, req=req) - - self.assertFalse(loaded) - self.assertIsNone(req.mamba_pool_idx) - self.assertEqual( - req_to_token_pool.mamba_allocator.available_size(), mamba_avail - ) - self._release_ongoing_load_back_locks(cache) - - def test_load_back_abort_keeps_preexisting_mamba_slot(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - self.assertTrue(leaf.evicted) - - # The request already owns its slot: an aborted load-back must not free it. - req = self._make_req(req_to_token_pool) - preexisting_slot = req.mamba_pool_idx - self.assertIsNotNone(preexisting_slot) - mamba_avail = req_to_token_pool.mamba_allocator.available_size() - - loaded = cache.load_back(leaf, mem_quota=-(10**9), req=req) - - self.assertFalse(loaded) - self.assertIs(req.mamba_pool_idx, preexisting_slot) - self.assertEqual( - req_to_token_pool.mamba_allocator.available_size(), mamba_avail - ) - self._release_ongoing_load_back_locks(cache) - - def test_load_back_success_publishes_fresh_mamba_slot(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - self.assertTrue(leaf.evicted) - - req = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - mamba_avail = req_to_token_pool.mamba_allocator.available_size() - - loaded = cache.load_back(leaf, req=req) - - self.assertTrue(loaded) - # the successful load must keep the freshly allocated slot published - self.assertIsNotNone(req.mamba_pool_idx) - self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].value) - # one slot restores the node's mamba value, one is the request's CoW slot - self.assertEqual( - req_to_token_pool.mamba_allocator.available_size(), mamba_avail - 2 - ) - self._finish_pending_loads(cache) - self._release_ongoing_load_back_locks(cache) - - def test_load_back_success_copies_mamba_state_into_request_slot(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - - # Stamp the node's mamba state so the host backup carries it. - node_mamba_indices = leaf.component_data[ComponentType.MAMBA].value.clone() - self._fill_mamba_state(req_to_token_pool, node_mamba_indices, marker=11) - expected_temporal, expected_conv = self._snapshot_mamba_state( - req_to_token_pool, node_mamba_indices - ) - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - self.assertTrue(leaf.evicted) - - req = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - - loaded = cache.load_back(leaf, req=req) - self.assertTrue(loaded) - self.assertIsNotNone(req.mamba_pool_idx) - self._finish_pending_loads(cache) - - # The CoW slot must actually hold the backed-up mamba state, not merely exist. - actual_temporal, actual_conv = self._snapshot_mamba_state( - req_to_token_pool, req.mamba_pool_idx.unsqueeze(0) - ) - self.assertTrue(torch.equal(actual_temporal, expected_temporal)) - self.assertEqual(len(actual_conv), len(expected_conv)) - for actual, expected in zip(actual_conv, expected_conv): - self.assertTrue(torch.equal(actual, expected)) - self._release_ongoing_load_back_locks(cache) - - def test_prepare_load_back_mamba(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - comp = cache.components[ComponentType.MAMBA] - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - self.assertTrue(leaf.evicted) - - # a request that already owns a slot -> nothing to prepare - req = self._make_req(req_to_token_pool) - self.assertIsNone(comp.prepare_load_back(leaf, req=req).allocated_mamba_slot) - - # no request -> nothing to prepare - self.assertIsNone(comp.prepare_load_back(leaf, req=None).allocated_mamba_slot) - - # fresh request + host-backed mamba -> allocates and publishes onto req - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - prep = comp.prepare_load_back(leaf, req=req) - self.assertIsNotNone(prep.allocated_mamba_slot) - self.assertEqual(int(req.mamba_pool_idx), int(prep.allocated_mamba_slot[0])) - - # node without host-backed mamba -> nothing to prepare - req2 = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req2.mamba_pool_idx.unsqueeze(0)) - req2.mamba_pool_idx = None - root = cache.root_node - self.assertIsNone(root.component_data[ComponentType.MAMBA].host_value) - self.assertIsNone(comp.prepare_load_back(root, req=req2).allocated_mamba_slot) - self.assertIsNone(req2.mamba_pool_idx) - - def test_prepare_load_back_skips_device_present_node(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - comp = cache.components[ComponentType.MAMBA] - - # Back up without evicting: device value stays and a host copy is added, so build_hicache_transfers no-ops and prepare must not allocate a dead slot. - self._backup_node(cache, leaf) - cd = leaf.component_data[ComponentType.MAMBA] - self.assertIsNotNone(cd.value) - self.assertIsNotNone(cd.host_value) - - req = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - mamba_avail = req_to_token_pool.mamba_allocator.available_size() - - self.assertIsNone(comp.prepare_load_back(leaf, req=req).allocated_mamba_slot) - self.assertIsNone(req.mamba_pool_idx) - self.assertEqual( - req_to_token_pool.mamba_allocator.available_size(), mamba_avail - ) - - def test_prepare_load_back_mamba_pool_exhausted(self): - if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: - self.skipTest("requires page_size=1 Full+Mamba") - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) - if len(chain) < 3: - self.skipTest("chain too short") - leaf = chain[-1] - comp = cache.components[ComponentType.MAMBA] - - self._backup_node(cache, leaf) - cache.evict(EvictParams(num_tokens=len(leaf.key))) - - req = self._make_req(req_to_token_pool) - req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) - req.mamba_pool_idx = None - retry_slot = req_to_token_pool.mamba_allocator.alloc(1) - - # first alloc fails -> prepare must evict a mamba slot and retry - with mock.patch.object( - req_to_token_pool.mamba_allocator, "alloc", side_effect=[None, retry_slot] - ), mock.patch.object(cache, "evict", autospec=True) as evict: - prep = comp.prepare_load_back(leaf, req=req) - evict.assert_called_once_with(EvictParams(num_tokens=0, mamba_num=1)) - self.assertIs(prep.allocated_mamba_slot, retry_slot) - self.assertEqual(int(req.mamba_pool_idx), int(retry_slot[0])) - def test_scheduler_hicache_aux_only_load_back_appends_full_device_indices(self): if self.cfg.page_size != 1: self.skipTest("page_size=1 keeps the expected suffix precise") @@ -3742,7 +3919,7 @@ class UnifiedRadixCacheSuite: ) ) - self.assertIs(new_node, leaf) + self.assertIs(cache.resolve_node_handle(new_node), leaf) self.assertEqual(new_indices.tolist(), leaf_full.tolist()) self.assertEqual(len(torch.cat([req.prefix_indices, new_indices])), len(tokens)) self.assertEqual( @@ -3772,6 +3949,10 @@ class UnifiedRadixCacheSuite: ) self._apply_match_to_req(req, match) + # Simulate a request without its own mamba slot so load-back allocates one + # (that allocation is what a called-off load-back must free + not publish). + req.mamba_pool_idx = None + avail_before = req_to_token_pool.mamba_allocator.available_size() new_indices, new_node = cache.init_load_back( InitLoadBackParams( best_match_node=req.best_match_node, @@ -3782,9 +3963,365 @@ class UnifiedRadixCacheSuite: ) self.assertEqual(len(new_indices), 0) - self.assertIs(new_node, match.last_device_node) + self.assertIs( + cache.resolve_node_handle(new_node), + cache.resolve_node_handle(match.last_device_node), + ) self.assertIsNone(leaf.component_data[ComponentType.FULL].value) self.assertIsNone(leaf.component_data[ComponentType.MAMBA].value) + # A failed load-back must roll back the pre-allocated mamba slot. + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), avail_before + ) + + def test_scheduler_hicache_load_back_rolls_back_mamba_on_load_failure(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + tokens = self._match_tokens_for_chain(chain) + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + + req = self._make_req(req_to_token_pool) + match = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req) + ) + self._apply_match_to_req(req, match) + + # Simulate a request without its own mamba slot so load-back allocates one. + req.mamba_pool_idx = None + avail_before = req_to_token_pool.mamba_allocator.available_size() + # H->D load fails after the mamba slot is pre-allocated -> must free it. + with mock.patch.object(cache.cache_controller, "load", return_value=None): + new_indices, _ = cache.init_load_back( + InitLoadBackParams( + best_match_node=req.best_match_node, + host_hit_length=req.host_hit_length, + req=req, + mem_quota=None, + ) + ) + + self.assertEqual(len(new_indices), 0) + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), avail_before + ) + + def test_scheduler_hicache_load_back_rolls_back_mamba_on_capacity_failure(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + tokens = self._match_tokens_for_chain(chain) + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + + req = self._make_req(req_to_token_pool) + match = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req) + ) + self._apply_match_to_req(req, match) + + # Simulate a request without its own mamba slot so load-back allocates one. + req.mamba_pool_idx = None + avail_before = req_to_token_pool.mamba_allocator.available_size() + # No device room and eviction frees nothing -> load-back bails after the + # mamba pre-alloc, which must still be freed. + with ( + mock.patch.object( + cache.token_to_kv_pool_allocator, "available_size", return_value=0 + ), + mock.patch.object( + cache, "evict", return_value=mock.Mock(num_tokens_evicted=0) + ), + ): + new_indices, _ = cache.init_load_back( + InitLoadBackParams( + best_match_node=req.best_match_node, + host_hit_length=req.host_hit_length, + req=req, + mem_quota=None, + ) + ) + + self.assertEqual(len(new_indices), 0) + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), avail_before + ) + + def test_load_back_abort_frees_unpublished_mamba_slot(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].host_value) + + # A request whose mamba slot was released: load_back's CoW arm allocates one. + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + mamba_avail = req_to_token_pool.mamba_allocator.available_size() + + # Impossible quota -> load_back aborts after building the transfers. + loaded = cache.load_back(leaf.id, mem_quota=-(10**9), req=req) + + self.assertFalse(loaded) + # the aborted call must return its slot and not leave req pointing at it + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), mamba_avail + ) + self._release_ongoing_load_back_locks(cache) + + def test_load_back_load_failure_frees_unpublished_mamba_slot(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + mamba_avail = req_to_token_pool.mamba_allocator.available_size() + + # cache_controller.load() failing (device alloc / transfer resolution) + # must also return the slot this call allocated. + with mock.patch.object(cache.cache_controller, "load", return_value=None): + loaded = cache.load_back(leaf.id, req=req) + + self.assertFalse(loaded) + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), mamba_avail + ) + self._release_ongoing_load_back_locks(cache) + + def test_load_back_abort_keeps_preexisting_mamba_slot(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + + # The request already owns its slot: an aborted load-back must not free it. + req = self._make_req(req_to_token_pool) + preexisting_slot = req.mamba_pool_idx + self.assertIsNotNone(preexisting_slot) + mamba_avail = req_to_token_pool.mamba_allocator.available_size() + + loaded = cache.load_back(leaf.id, mem_quota=-(10**9), req=req) + + self.assertFalse(loaded) + self.assertIs(req.mamba_pool_idx, preexisting_slot) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), mamba_avail + ) + self._release_ongoing_load_back_locks(cache) + + def test_load_back_success_publishes_fresh_mamba_slot(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + mamba_avail = req_to_token_pool.mamba_allocator.available_size() + + loaded = cache.load_back(leaf.id, req=req) + + self.assertTrue(loaded) + # the successful load must keep the freshly allocated slot published + self.assertIsNotNone(req.mamba_pool_idx) + self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].value) + # one slot restores the node's mamba value, one is the request's CoW slot + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), mamba_avail - 2 + ) + self._finish_pending_loads(cache) + self._release_ongoing_load_back_locks(cache) + + def test_load_back_success_copies_mamba_state_into_request_slot(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + + # Stamp the node's mamba state so the host backup carries it. + node_mamba_indices = leaf.component_data[ComponentType.MAMBA].value.clone() + self._fill_mamba_state(req_to_token_pool, node_mamba_indices, marker=11) + expected_temporal, expected_conv = self._snapshot_mamba_state( + req_to_token_pool, node_mamba_indices + ) + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + + loaded = cache.load_back(leaf.id, req=req) + self.assertTrue(loaded) + self.assertIsNotNone(req.mamba_pool_idx) + self._finish_pending_loads(cache) + + # The CoW slot must actually hold the backed-up mamba state, not merely exist. + actual_temporal, actual_conv = self._snapshot_mamba_state( + req_to_token_pool, req.mamba_pool_idx.unsqueeze(0) + ) + self.assertTrue(torch.equal(actual_temporal, expected_temporal)) + self.assertEqual(len(actual_conv), len(expected_conv)) + for actual, expected in zip(actual_conv, expected_conv): + self.assertTrue(torch.equal(actual, expected)) + self._release_ongoing_load_back_locks(cache) + + def test_prepare_load_back_mamba(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + comp = cache.components[ComponentType.MAMBA] + + self._backup_node(cache, leaf) + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + + # device value still present -> nothing to prepare even though host-backed + self.assertIsNone(comp.prepare_load_back(leaf.id, req=req).allocated_mamba_slot) + self.assertIsNone(req.mamba_pool_idx) + + cache.evict(EvictParams(num_tokens=len(leaf.key))) + self.assertTrue(leaf.evicted) + + # a request that already owns a slot -> nothing to prepare + req_owned = self._make_req(req_to_token_pool) + self.assertIsNone( + comp.prepare_load_back(leaf.id, req=req_owned).allocated_mamba_slot + ) + + # no request -> nothing to prepare + self.assertIsNone( + comp.prepare_load_back(leaf.id, req=None).allocated_mamba_slot + ) + + # fresh request + host-only mamba -> allocates and publishes onto req + prep = comp.prepare_load_back(leaf.id, req=req) + self.assertIsNotNone(prep.allocated_mamba_slot) + self.assertEqual(int(req.mamba_pool_idx), int(prep.allocated_mamba_slot[0])) + + # node without host-backed mamba -> nothing to prepare + req2 = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req2.mamba_pool_idx.unsqueeze(0)) + req2.mamba_pool_idx = None + root = cache.root_node + self.assertIsNone(root.component_data[ComponentType.MAMBA].host_value) + self.assertIsNone( + comp.prepare_load_back(root.id, req=req2).allocated_mamba_slot + ) + self.assertIsNone(req2.mamba_pool_idx) + + def test_prepare_load_back_skips_device_present_node(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + comp = cache.components[ComponentType.MAMBA] + + # Back up without evicting: device value stays and a host copy is added, so build_hicache_transfers no-ops and prepare must not allocate a dead slot. + self._backup_node(cache, leaf) + cd = leaf.component_data[ComponentType.MAMBA] + self.assertIsNotNone(cd.value) + self.assertIsNotNone(cd.host_value) + + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + mamba_avail = req_to_token_pool.mamba_allocator.available_size() + + self.assertIsNone(comp.prepare_load_back(leaf.id, req=req).allocated_mamba_slot) + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual( + req_to_token_pool.mamba_allocator.available_size(), mamba_avail + ) + + def test_prepare_load_back_mamba_pool_exhausted(self): + if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: + self.skipTest("requires page_size=1 Full+Mamba") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + comp = cache.components[ComponentType.MAMBA] + + self._backup_node(cache, leaf) + cache.evict(EvictParams(num_tokens=len(leaf.key))) + + req = self._make_req(req_to_token_pool) + req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + retry_slot = req_to_token_pool.mamba_allocator.alloc(1) + + # first alloc fails -> prepare must evict a mamba slot and retry + with mock.patch.object( + req_to_token_pool.mamba_allocator, "alloc", side_effect=[None, retry_slot] + ), mock.patch.object(cache, "evict", autospec=True) as evict: + prep = comp.prepare_load_back(leaf.id, req=req) + evict.assert_called_once_with(EvictParams(num_tokens=0, mamba_num=1)) + self.assertIs(prep.allocated_mamba_slot, retry_slot) + self.assertEqual(int(req.mamba_pool_idx), int(retry_slot[0])) def test_hicache_swa_load_back_min_suffix(self): """LOAD_BACK collects only the suffix nodes needed to cover sliding_window_size.""" @@ -3827,6 +4364,8 @@ class UnifiedRadixCacheSuite: # host_indices must cover exactly the expected suffix tokens (>= sw). self.assertEqual(int(xfer.host_indices.numel()), expected_pages * ps) self.assertGreaterEqual(int(xfer.host_indices.numel()), sw) + # nodes_to_load holds NodeIds; compare against the chain's ids. + chain = [n.id for n in chain] self.assertEqual(xfer.nodes_to_load, chain[-expected_pages:]) def _swa_finalize_setup(self): @@ -3889,7 +4428,7 @@ class UnifiedRadixCacheSuite: best_match_node=leaf, host_hit_length=0, ) - result = swa_comp.finalize_match_result( + result = swa_comp.finalize_match_result_in_tree_core( result=result, params=MatchPrefixParams( key=RadixKey(array("q", self._make_seq(1, 1))) @@ -3919,8 +4458,8 @@ class UnifiedRadixCacheSuite: n.component_data[ComponentType.SWA].value = None # SWA LRU bookkeeping must reflect tombstone state for the # _restore_device_value path to exercise the host->device move. - cache.lru_lists[ComponentType.SWA].remove_node(n) - cache.host_lru_lists[ComponentType.SWA].insert_mru(n) + cache.tree_core.lru_lists[ComponentType.SWA].remove_node(n) + cache.tree_core.host_lru_lists[ComponentType.SWA].insert_mru(n) # Build the LOAD_BACK transfer the same way load_back() would. swa_comp = cache.components[ComponentType.SWA] @@ -3929,7 +4468,7 @@ class UnifiedRadixCacheSuite: ) self.assertIsNotNone(transfers) xfer = transfers[0] - self.assertEqual(xfer.nodes_to_load, loaded_nodes) + self.assertEqual(xfer.nodes_to_load, [n.id for n in loaded_nodes]) # Allocate SWA device slots from the inner allocator (mirrors how # _resolve_pool_transfers_allocation routes via device_alloc_fn -> @@ -3940,11 +4479,16 @@ class UnifiedRadixCacheSuite: xfer.device_indices = new_swa # Snapshot pre-commit state for invariants checks. - pre_evictable = cache.component_evictable_size_[ComponentType.SWA] + pre_evictable = cache.tree_core.component_evictable_size_[ComponentType.SWA] + load_actions = [] swa_comp.commit_hicache_transfer( - chain[-1], CacheTransferPhase.LOAD_BACK, transfers=transfers + chain[-1], + CacheTransferPhase.LOAD_BACK, + transfers=transfers, + cache_actions=load_actions, ) + cache._apply_cache_actions(load_actions) # (1) cd.value restored, host LRU -> device LRU swap done. offset = 0 @@ -3957,8 +4501,10 @@ class UnifiedRadixCacheSuite: new_swa[offset : offset + chunk_len].tolist(), ) offset += chunk_len - self.assertTrue(cache.lru_lists[ComponentType.SWA].in_list(n)) - self.assertFalse(cache.host_lru_lists[ComponentType.SWA].in_list(n)) + self.assertTrue(cache.tree_core.lru_lists[ComponentType.SWA].in_list(n)) + self.assertFalse( + cache.tree_core.host_lru_lists[ComponentType.SWA].in_list(n) + ) self.assertEqual(offset, n_swa) # (2) full_to_swa_index_mapping rebuilt for every loaded chunk. @@ -3970,10 +4516,28 @@ class UnifiedRadixCacheSuite: # Evictable size moved up by the restored token count. self.assertEqual( - cache.component_evictable_size_[ComponentType.SWA] - pre_evictable, + cache.tree_core.component_evictable_size_[ComponentType.SWA] + - pre_evictable, n_swa, ) + def test_swa_remapping_rejects_length_mismatch(self): + """A RebuildFullToSWAMapping with mismatched full/swa lengths is rejected on apply.""" + if self._skip_unsupported_hicache_test(): + return + if ComponentType.SWA not in self.cfg.components: + self.skipTest("requires the SWA component") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_hicache(cache, write_policy="write_back") + full = torch.tensor([1, 2], dtype=torch.int64, device=cache.device) + swa = torch.tensor([1], dtype=torch.int64, device=cache.device) + # Per-pair token-count mismatch is rejected. + with self.assertRaises(AssertionError): + cache._apply_cache_actions([RebuildFullToSWAMapping([full], [swa])]) + # A mismatched number of full vs swa chunks is rejected. + with self.assertRaises(AssertionError): + cache._apply_cache_actions([RebuildFullToSWAMapping([full], [swa, swa])]) + def _swa_anchor_chain_tokens(self, num_pages: int) -> list[int]: """Reproduce the token sequence used by _build_chain_pages.""" tokens: list[int] = [] @@ -4040,7 +4604,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(transfers), 1) xfer = transfers[0] self.assertEqual(xfer.name, PoolName.SWA) - self.assertEqual(xfer.nodes_to_load, [y]) + self.assertEqual(xfer.nodes_to_load, [y.id]) self.assertEqual(int(xfer.host_indices.numel()), ps) with self.assertRaises(AssertionError): @@ -4058,7 +4622,7 @@ class UnifiedRadixCacheSuite: best_match_node=x, host_hit_length=0, ) - result = swa_comp.finalize_match_result( + result = swa_comp.finalize_match_result_in_tree_core( result=base, params=MatchPrefixParams(key=RadixKey(array("q", self._make_seq(1, 1)))), value_chunks=[], @@ -4085,11 +4649,11 @@ class UnifiedRadixCacheSuite: self.assertIsNotNone(old_swa) cd.value = None - cache.lru_lists[ComponentType.SWA].remove_node(tombstone) - cache.host_lru_lists[ComponentType.SWA].insert_mru(tombstone) - cache.component_evictable_size_[ComponentType.SWA] -= len(old_swa) + cache.tree_core.lru_lists[ComponentType.SWA].remove_node(tombstone) + cache.tree_core.host_lru_lists[ComponentType.SWA].insert_mru(tombstone) + cache.tree_core.component_evictable_size_[ComponentType.SWA] -= len(old_swa) - temp_lock = cache.inc_lock_ref(leaf) + temp_lock = cache.inc_lock_ref(leaf.id) self.assertEqual(cd.lock_ref, 0) xfer = cache.components[ComponentType.SWA].build_hicache_transfers( @@ -4098,19 +4662,24 @@ class UnifiedRadixCacheSuite: new_swa = allocator.swa_attn_allocator.alloc(int(xfer.host_indices.numel())) self.assertIsNotNone(new_swa) xfer.device_indices = new_swa + load_actions = [] cache.components[ComponentType.SWA].commit_hicache_transfer( - leaf, CacheTransferPhase.LOAD_BACK, transfers=[xfer] + leaf, + CacheTransferPhase.LOAD_BACK, + transfers=[xfer], + cache_actions=load_actions, ) + cache._apply_cache_actions(load_actions) - load_back_lock = cache.inc_lock_ref(leaf) - request_lock = cache.inc_lock_ref(leaf) + load_back_lock = cache.inc_lock_ref(leaf.id) + request_lock = cache.inc_lock_ref(leaf.id) self.assertEqual(cd.lock_ref, 2) - cache.dec_lock_ref(leaf, temp_lock.to_dec_params()) + cache.dec_lock_ref(leaf.id, temp_lock.to_dec_params()) self.assertEqual(cd.lock_ref, 2) - cache.dec_lock_ref(leaf, load_back_lock.to_dec_params()) - cache.dec_lock_ref(leaf, request_lock.to_dec_params()) + cache.dec_lock_ref(leaf.id, load_back_lock.to_dec_params()) + cache.dec_lock_ref(leaf.id, request_lock.to_dec_params()) self.assertEqual(cd.lock_ref, 0) def test_hicache_swa_load_back_uses_full_pool_capacity(self): @@ -4175,7 +4744,7 @@ class UnifiedRadixCacheSuite: ) with mock.patch.object(cache, "evict", wraps=cache.evict) as evict_mock: - self.assertTrue(cache.load_back(leaf)) + self.assertTrue(cache.load_back(leaf.id)) # Full pre-eviction must not be triggered by SWA pool pressure. full_pre_evict_calls = [ @@ -4231,7 +4800,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(cd_y.lock_ref, 0) self.assertEqual(cd_a.lock_ref, 0) - temp_lock = cache.inc_lock_ref(anchor) + temp_lock = cache.inc_lock_ref(anchor.id) self.assertEqual(cd_anchor.lock_ref, 0) self.assertEqual(cd_y.lock_ref, 1) self.assertEqual(cd_a.lock_ref, 1) @@ -4240,17 +4809,17 @@ class UnifiedRadixCacheSuite: cd_anchor.value = anchor_value - second_lock = cache.inc_lock_ref(anchor) + second_lock = cache.inc_lock_ref(anchor.id) self.assertEqual(cd_anchor.lock_ref, 1) self.assertEqual(cd_y.lock_ref, 2) self.assertEqual(cd_a.lock_ref, 2) - cache.dec_lock_ref(anchor, temp_lock.to_dec_params()) + cache.dec_lock_ref(anchor.id, temp_lock.to_dec_params()) self.assertEqual(cd_anchor.lock_ref, 1) self.assertEqual(cd_y.lock_ref, 1) self.assertEqual(cd_a.lock_ref, 1) - cache.dec_lock_ref(anchor, second_lock.to_dec_params()) + cache.dec_lock_ref(anchor.id, second_lock.to_dec_params()) self.assertEqual(cd_anchor.lock_ref, 0) self.assertEqual(cd_y.lock_ref, 0) self.assertEqual(cd_a.lock_ref, 0) @@ -4268,18 +4837,18 @@ class UnifiedRadixCacheSuite: seq = self._make_seq(1, 2) self._insert(cache, allocator, req_to_token_pool, seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node + node = cache.resolve_node_handle(m.last_device_node) cd = node.component_data[ComponentType.MAMBA] old_mamba = cd.value self.assertIsNotNone(old_mamba) self._simulate_backup(cache, node) cd.value = None - cache.lru_lists[ComponentType.MAMBA].remove_node(node) - cache.host_lru_lists[ComponentType.MAMBA].insert_mru(node) - cache.component_evictable_size_[ComponentType.MAMBA] -= len(old_mamba) + cache.tree_core.lru_lists[ComponentType.MAMBA].remove_node(node) + cache.tree_core.host_lru_lists[ComponentType.MAMBA].insert_mru(node) + cache.tree_core.component_evictable_size_[ComponentType.MAMBA] -= len(old_mamba) - temp_lock = cache.inc_lock_ref(node) + temp_lock = cache.inc_lock_ref(node.id) self.assertEqual(cd.lock_ref, 0) xfer = cache.components[ComponentType.MAMBA].build_hicache_transfers( @@ -4289,18 +4858,18 @@ class UnifiedRadixCacheSuite: self.assertIsNotNone(new_mamba) xfer.device_indices = new_mamba cache.components[ComponentType.MAMBA].commit_hicache_transfer( - node, CacheTransferPhase.LOAD_BACK, transfers=[xfer] + node, CacheTransferPhase.LOAD_BACK, transfers=[xfer], cache_actions=[] ) - load_back_lock = cache.inc_lock_ref(node) - request_lock = cache.inc_lock_ref(node) + load_back_lock = cache.inc_lock_ref(node.id) + request_lock = cache.inc_lock_ref(node.id) self.assertEqual(cd.lock_ref, 2) - cache.dec_lock_ref(node, temp_lock.to_dec_params()) + cache.dec_lock_ref(node.id, temp_lock.to_dec_params()) self.assertEqual(cd.lock_ref, 2) - cache.dec_lock_ref(node, load_back_lock.to_dec_params()) - cache.dec_lock_ref(node, request_lock.to_dec_params()) + cache.dec_lock_ref(node.id, load_back_lock.to_dec_params()) + cache.dec_lock_ref(node.id, request_lock.to_dec_params()) self.assertEqual(cd.lock_ref, 0) def test_hicache_mixed_backup_evict_insert(self): @@ -4317,7 +4886,7 @@ class UnifiedRadixCacheSuite: for i in range(3): m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i])))) - self._backup_node(cache, m.last_device_node) + self._backup_node(cache, cache.resolve_node_handle(m.last_device_node)) # Evict to free some tokens cache.evict(EvictParams(num_tokens=len(seqs[0]) * 2)) @@ -4330,7 +4899,10 @@ class UnifiedRadixCacheSuite: cache.sanity_check() # Verify D-leaf / H-leaf mutual exclusion - overlap = cache.evictable_device_leaves & cache.evictable_host_leaves + overlap = ( + cache.tree_core.evictable_device_leaves + & cache.tree_core.evictable_host_leaves + ) self.assertEqual(len(overlap), 0) def test_hicache_write_back_leaf_backup(self): @@ -4346,20 +4918,20 @@ class UnifiedRadixCacheSuite: self._insert(cache, allocator, req_to_token_pool, leaf_seq) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf_seq)))) - leaf = m.last_device_node + leaf = cache.resolve_node_handle(m.last_device_node) parent = leaf.parent self.assertIsNot(parent, cache.root_node) self.assertFalse(leaf.backuped) self.assertFalse(parent.backuped) - lr = cache.inc_lock_ref(parent) + lr = cache.inc_lock_ref(parent.id) try: evict_tokens = len(leaf_seq) - len(base) cache.evict(EvictParams(num_tokens=evict_tokens)) finally: cache.dec_lock_ref( - parent, + parent.id, DecLockRefParams( swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None) ), @@ -4373,6 +4945,38 @@ class UnifiedRadixCacheSuite: cache.sanity_check() + def test_build_backup_kv_action_orders_ancestors_first(self): + """write-through backs up unbacked ancestors parent-first; write-back only the leaf.""" + if self.cfg.has_swa: + self.skipTest("SWA boundary splits deepen the chain non-deterministically") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + + top_seq = self._make_seq(1, 1) + mid_seq = top_seq + self._make_seq(1000, 1) + leaf_seq = mid_seq + self._make_seq(2000, 1) + self._insert(cache, allocator, req_to_token_pool, top_seq) + self._insert(cache, allocator, req_to_token_pool, mid_seq) + self._insert(cache, allocator, req_to_token_pool, leaf_seq) + + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf_seq)))) + leaf = cache.resolve_node_handle(m.last_device_node) + mid = leaf.parent + top = mid.parent + self.assertIsNot(top, cache.root_node) + self.assertFalse(leaf.backuped) + self.assertFalse(mid.backuped) + self.assertFalse(top.backuped) + + # write-through: every unbacked ancestor, top-most first, child last + write_through = cache.tree_core._build_backup_kv_action(leaf, write_back=False) + self.assertEqual(write_through.node_ids, [top.id, mid.id, leaf.id]) + + # write-back: only the eviction victim, even with unbacked ancestors + write_back = cache.tree_core._build_backup_kv_action(leaf, write_back=True) + self.assertEqual(write_back.node_ids, [leaf.id]) + + cache.sanity_check() + class UnifiedLRUListBoundedRefreshTest(CustomTestCase): @@ -4472,7 +5076,7 @@ class TestUnifiedMambaLRUMatchRefresh(CustomTestCase): cfg = CacheConfig(page_size=1, components=(ComponentType.FULL, ComponentType.MAMBA)) def _mamba_lru_mru_to_lru(self, cache): - lru = cache.lru_lists[ComponentType.MAMBA] + lru = cache.tree_core.lru_lists[ComponentType.MAMBA] pt = lru._pt out, cur = [], lru.head.lru_next[pt] while cur is not lru.tail: @@ -4505,9 +5109,10 @@ class TestUnifiedMambaLRUMatchRefresh(CustomTestCase): ) def match_leaf(tokens): - return cache.match_prefix( + node_id = cache.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens))) ).best_match_node + return cache.tree_core.node_by_id(node_id) # Two independent sessions, each a 2-node mamba chain: # root -> a1 -> b1 and root -> a2 -> b2 @@ -4564,7 +5169,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): kv_indices = allocator.alloc(len(tokens)) self.assertIsNotNone(kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) - req.last_node = cache.root_node + req.last_node = cache.root_node.id cache.cache_finished_req( req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len() @@ -4723,5 +5328,855 @@ for _cfg in _CONFIGS: del _cfg, _name +def _component_with_cache(component_type, cache): + """A registry component instance bound to a (mock) cache and its tree_core.""" + component = object.__new__(COMPONENT_REGISTRY[component_type]) + component.cache = cache + component.tree_core = cache.tree_core + return component + + +class TestUnifiedRadixCacheActionRouting(CustomTestCase): + """CacheAction routing: each type forwards to the right Controller API.""" + + def test_apply_cache_action_routes_replace_write_through(self): + cache = mock.MagicMock() + action = ReplaceWriteThroughOnNodeSplit( + ack_id=7, old_node_id=2, new_node_id=3, new_child_node_id=2 + ) + UnifiedRadixCache._apply_cache_action(cache, action) + cache._replace_pending_write_through_node.assert_called_once_with(7, 2, [3, 2]) + + def test_apply_cache_action_routes_free_device_kv(self): + cache = mock.MagicMock() + first, second = torch.tensor([4, 5]), torch.tensor([6]) + action = FreeDeviceKV([first, second]) + UnifiedRadixCache._apply_cache_action(cache, action) + cache.token_to_kv_pool_allocator.free.assert_has_calls( + [mock.call(first), mock.call(second)] + ) + + def test_apply_cache_action_routes_free_component_device_kv(self): + cache = mock.MagicMock() + component = mock.MagicMock() + cache.components = {ComponentType.SWA: component} + action = FreeComponentDeviceSlot( + [torch.tensor([4, 5])], component_type=ComponentType.SWA + ) + UnifiedRadixCache._apply_cache_action(cache, action) + component.apply_component_action.assert_called_once_with(action) + + def test_apply_component_action_device_kv_full_swa_uses_full_attn(self): + cache = mock.MagicMock() + cache.is_swa_enabled = True + indices = torch.tensor([4, 5]) + _component_with_cache(ComponentType.FULL, cache).apply_component_action( + FreeComponentDeviceSlot([indices], component_type=ComponentType.FULL) + ) + cache.token_to_kv_pool_allocator.full_attn_allocator.free.assert_called_once_with( + indices + ) + + def test_apply_component_action_device_kv_swa_uses_free_swa(self): + cache = mock.MagicMock() + indices = torch.tensor([4, 5]) + _component_with_cache(ComponentType.SWA, cache).apply_component_action( + FreeComponentDeviceSlot([indices], component_type=ComponentType.SWA) + ) + cache.token_to_kv_pool_allocator.free_swa.assert_called_once_with(indices) + + def test_apply_component_action_device_kv_mamba_uses_mamba_allocator(self): + cache = mock.MagicMock() + cache.req_to_token_pool.mamba_ckpt_pool = None + indices = torch.tensor([4, 5]) + _component_with_cache(ComponentType.MAMBA, cache).apply_component_action( + FreeComponentDeviceSlot([indices], component_type=ComponentType.MAMBA) + ) + cache.req_to_token_pool.mamba_allocator.free.assert_called_once_with(indices) + + def test_apply_component_action_device_kv_mamba_routes_to_int8_ckpt_pool(self): + cache = mock.MagicMock() + indices = torch.tensor([4, 5]) + _component_with_cache(ComponentType.MAMBA, cache).apply_component_action( + FreeComponentDeviceSlot([indices], component_type=ComponentType.MAMBA) + ) + cache.req_to_token_pool.mamba_ckpt_pool.free.assert_called_once_with(indices) + cache.req_to_token_pool.mamba_allocator.free.assert_not_called() + + def test_apply_cache_action_routes_free_component_host_kv(self): + cache = mock.MagicMock() + component = mock.MagicMock() + cache.components = {ComponentType.SWA: component} + action = FreeComponentHostSlot( + [torch.tensor([4, 5])], component_type=ComponentType.SWA + ) + UnifiedRadixCache._apply_cache_action(cache, action) + component.apply_component_action.assert_called_once_with(action) + + def test_apply_component_action_host_kv_swa(self): + cache = mock.MagicMock() + first, second = torch.tensor([4, 5]), torch.tensor([6]) + empty = torch.empty((0,), dtype=torch.int64) + _component_with_cache(ComponentType.SWA, cache).apply_component_action( + FreeComponentHostSlot( + [first, empty, second], component_type=ComponentType.SWA + ), + ) + calls = cache.cache_controller.append_host_mem_release.call_args_list + # empty page skipped; each non-empty page released under the SWA pool + self.assertEqual(len(calls), 2) + self.assertEqual(calls[0].kwargs["extra_pools"][0].name, PoolName.SWA) + self.assertIs(calls[0].kwargs["extra_pools"][0].host_indices, first) + self.assertEqual(calls[1].kwargs["extra_pools"][0].name, PoolName.SWA) + self.assertIs(calls[1].kwargs["extra_pools"][0].host_indices, second) + + def test_apply_component_action_host_kv_mamba(self): + cache = mock.MagicMock() + indices = torch.tensor([4, 5]) + _component_with_cache(ComponentType.MAMBA, cache).apply_component_action( + FreeComponentHostSlot([indices], component_type=ComponentType.MAMBA) + ) + calls = cache.cache_controller.append_host_mem_release.call_args_list + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0].kwargs["extra_pools"][0].name, PoolName.MAMBA) + self.assertIs(calls[0].kwargs["extra_pools"][0].host_indices, indices) + + def test_apply_cache_action_routes_swa_rebuild(self): + cache = mock.MagicMock() + component = mock.MagicMock() + cache.components = {ComponentType.SWA: component} + action = SWARebuild(node_id=5, source_value=torch.tensor([3, 4])) + UnifiedRadixCache._apply_cache_action(cache, action) + component.apply_component_action.assert_called_once_with(action) + + def test_apply_component_action_swa_rebuild(self): + cache = mock.MagicMock() + alloc = cache.token_to_kv_pool_allocator + source_value = torch.tensor([3, 4], dtype=torch.int64) + swa_value = alloc.translate_loc_from_full_to_swa.return_value + _component_with_cache(ComponentType.SWA, cache).apply_component_action( + SWARebuild(node_id=5, source_value=source_value), + ) + # translate the source full to SWA and store it on the node (no free) + alloc.translate_loc_from_full_to_swa.assert_called_once_with(source_value) + alloc.free.assert_not_called() + cache.tree_core.set_component_device_value.assert_called_once_with( + 5, ComponentType.SWA, swa_value + ) + + def test_apply_cache_action_routes_swa_recover_on_full_locked(self): + cache = mock.MagicMock() + component = mock.MagicMock() + cache.components = {ComponentType.SWA: component} + action = RecoverSWAWithLockedFull( + node_id=5, + kept_full=torch.tensor([1, 2]), + incoming_full=torch.tensor([3, 4]), + ) + UnifiedRadixCache._apply_cache_action(cache, action) + component.apply_component_action.assert_called_once_with(action) + + def test_apply_component_action_swa_recover_on_full_locked(self): + cache = mock.MagicMock() + alloc = cache.token_to_kv_pool_allocator + kept_full = torch.tensor([1, 2], dtype=torch.int64) + incoming_full = torch.tensor([3, 4], dtype=torch.int64) + swa_value = alloc.translate_loc_from_full_to_swa.return_value + _component_with_cache(ComponentType.SWA, cache).apply_component_action( + RecoverSWAWithLockedFull( + node_id=5, + kept_full=kept_full, + incoming_full=incoming_full, + ), + ) + # keep the locked full, remap it onto the incoming full's SWA translation + alloc.translate_loc_from_full_to_swa.assert_called_once_with(incoming_full) + alloc.set_full_to_swa_mapping.assert_called_once_with(kept_full, swa_value) + # the incoming full's stale mapping is cleared, then its slot freed (full-only) + key, val = alloc.full_to_swa_index_mapping.__setitem__.call_args.args + self.assertTrue(torch.equal(key, incoming_full)) + self.assertEqual(val, 0) + alloc.full_attn_allocator.free.assert_called_once_with(incoming_full) + alloc.free.assert_not_called() + cache.tree_core.set_component_device_value.assert_called_once_with( + 5, ComponentType.SWA, swa_value + ) + + def test_apply_cache_action_unknown_type_raises(self): + cache = mock.MagicMock() + with self.assertRaises(AssertionError): + UnifiedRadixCache._apply_cache_action(cache, object()) + + def test_apply_cache_actions_applies_each_in_order(self): + cache = mock.MagicMock() + first = ReplaceWriteThroughOnNodeSplit( + ack_id=1, old_node_id=1, new_node_id=2, new_child_node_id=1 + ) + second = ReplaceWriteThroughOnNodeSplit( + ack_id=2, old_node_id=3, new_node_id=4, new_child_node_id=3 + ) + UnifiedRadixCache._apply_cache_actions(cache, [first, second]) + cache._apply_cache_action.assert_has_calls( + [mock.call(first), mock.call(second)] + ) + + def test_chained_replace_write_through_requires_list_order(self): + # A pending node split twice in one walk emits two chained Replaces: + # the second one's old_node_id only enters the publish list when the + # first is applied, so list order is a hard contract. + def make_cache(): + cache = mock.MagicMock() + cache.ongoing_write_through = {7: _OngoingWriteThrough(10, None, [5, 10])} + return cache + + def apply(cache, action): + UnifiedRadixCache._replace_pending_write_through_node( + cache, + action.ack_id, + action.old_node_id, + [action.new_node_id, action.new_child_node_id], + ) + + # split X(10) -> [A(11), X(10)], then fragment A(11) -> [B(12), A(11)] + first = ReplaceWriteThroughOnNodeSplit( + ack_id=7, old_node_id=10, new_node_id=11, new_child_node_id=10 + ) + second = ReplaceWriteThroughOnNodeSplit( + ack_id=7, old_node_id=11, new_node_id=12, new_child_node_id=11 + ) + + # in list order, the publish list threads through both replaces + cache = make_cache() + apply(cache, first) + apply(cache, second) + self.assertEqual( + cache.ongoing_write_through[7].publish_node_ids, [5, 12, 11, 10] + ) + + # reversed order silently drops the second fragment - documents why + # emission order must be preserved end to end + cache = make_cache() + apply(cache, second) + apply(cache, first) + self.assertEqual(cache.ongoing_write_through[7].publish_node_ids, [5, 11, 10]) + + +class _InsertWalkSuite(CustomTestCase): + """Fixture helpers from UnifiedRadixCacheSuite, without inheriting its tests.""" + + _rid = 0 + _make_req = UnifiedRadixCacheSuite._make_req + _alloc = UnifiedRadixCacheSuite._alloc + _insert = UnifiedRadixCacheSuite._insert + _init_hicache = UnifiedRadixCacheSuite._init_hicache + _build_hicache_fixture = UnifiedRadixCacheSuite._build_hicache_fixture + _make_seq = UnifiedRadixCacheSuite._make_seq + _skip_unsupported_hicache_test = ( + UnifiedRadixCacheSuite._skip_unsupported_hicache_test + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") +class TestResumableInsertWalk(_InsertWalkSuite): + cfg = CacheConfig() + + def test_walk_backup_can_host_evict_on_path_h_leaf(self): + """A crossing node's backup runs at its walk step, so its host eviction + can still take an H-leaf deeper on the inserted path.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + top = next(iter(cache.root_node.children.values())) + self._insert(cache, allocator, req_to_token_pool, list(range(1, 9))) + h_leaf = next(iter(top.children.values())) + self.assertGreater(_write_backup(cache, h_leaf, write_back=True), 0) + cache.writing_check(write_back=True) + cache.evict(EvictParams(num_tokens=4)) + self.assertTrue(h_leaf.evicted) + + # Fill the host pool below len(top) free, keeping the on-path H-leaf + # the oldest host entry and pinning the unbacked path root. + cache.inc_lock_ref(top.id) + host_pool = cache.cache_controller.mem_pool_host + start = 1000 + while host_pool.available_size() >= len(top.key): + count = min(host_pool.available_size() - len(top.key) + 1, 250) + tokens = list(range(start, start + count)) + start += 1000 + self._insert(cache, allocator, req_to_token_pool, tokens) + filler = None + for child in cache.root_node.children.values(): + if child is not top and not child.evicted: + filler = child + self.assertIsNotNone(filler) + self.assertGreater(_write_backup(cache, filler, write_back=True), 0) + cache.writing_check(write_back=True) + cache.evict(EvictParams(num_tokens=count)) + self.assertTrue(filler.evicted) + cache.dec_lock_ref(top.id) + + # The crossing backup evicts exactly the on-path H-leaf, then the + # remaining suffix is recreated as a fresh leaf. + cache.write_through_threshold = top.hit_count + 1 + self._insert(cache, allocator, req_to_token_pool, list(range(1, 13))) + cache.writing_check(write_back=True) + + self.assertTrue(top.backuped) + self.assertNotIn(h_leaf, top.children.values()) + self.assertIsNone(h_leaf.component_data[ComponentType.FULL].host_value) + (child_key_len,) = {len(c.key) for c in top.children.values()} + self.assertEqual(child_key_len, 8) + cache.sanity_check() + + def test_insert_aborts_continuation_when_action_apply_fails(self): + """An exception while executing a barrier's actions aborts the suspended + insert instead of leaking its continuation.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = 2 + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + + with mock.patch.object( + cache, "_execute_and_commit_kv_backup", side_effect=RuntimeError("boom") + ): + with self.assertRaises(RuntimeError): + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + self.assertIsNone(cache.tree_core._ongoing_insert_walk_state) + + # The tree stays usable and the crossing re-fires on the next walk. + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + cache.writing_check(write_back=True) + ancestor = next(iter(cache.root_node.children.values())) + self.assertTrue(ancestor.backuped) + cache.sanity_check() + + def test_begin_insert_rejects_concurrent_walk(self): + """Insert walks are single-flight: beginning a second insert while one + is suspended at a barrier is re-entrancy and must fail fast.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = 2 + self._insert(cache, allocator, req_to_token_pool, [1, 2]) + + # Suspend an insert at its crossing barrier by pumping it directly. + params = InsertParams( + key=RadixKey(array("q", [1, 2, 3, 4])), value=self._alloc(allocator, 4) + ) + step = cache.tree_core.begin_insert(params) + self.assertIsNone(step.result) + with self.assertRaises(AssertionError): + cache.tree_core.begin_insert(params) + cache.tree_core.end_insert() + + def test_insert_abort_drains_pending_deferred_frees(self): + """A mid-insert failure after a deferred dup-free accumulated must still + return those slots to the allocator via the end_insert drain.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + + # The overlap walk defers a 4-slot dup-free; the commit hook then raises. + available = allocator.available_size() + full_comp = cache.components[ComponentType.FULL] + with mock.patch.object( + full_comp, "commit_insert_component_data", side_effect=RuntimeError("boom") + ): + with self.assertRaises(RuntimeError): + self._insert(cache, allocator, req_to_token_pool, list(range(1, 9))) + + # 8 alloc'd for the insert, 4 dup slots drained back on abort. + self.assertEqual(allocator.available_size(), available - 4) + self.assertIsNone(cache.tree_core._ongoing_insert_walk_state) + + def test_deferrable_actions_ride_final_step_without_suspension(self): + """A walk whose only actions are deferrable frees completes in a single + step, the batched frees riding the final step's actions.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + + # The overlap re-insert defers a dup-free; no barrier action fires. + params = InsertParams( + key=RadixKey(array("q", list(range(1, 9)))), + value=self._alloc(allocator, 8), + ) + step = cache.tree_core.begin_insert(params) + self.assertIsNotNone(step.result) + self.assertTrue(any(isinstance(a, FreeDeviceKV) for a in step.actions)) + cache._apply_cache_actions(step.actions) + cache.tree_core.end_insert() + cache.sanity_check() + + def test_backup_executor_skips_already_backed_nodes(self): + """Overlapping BackupKV chains must not back a node twice: a second + backup would allocate a second host copy and leak the first.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + node = next(iter(cache.root_node.children.values())) + + self.assertGreater(_write_backup(cache, node, write_back=True), 0) + cache.writing_check(write_back=True) + self.assertTrue(node.backuped) + host_avail = cache.cache_controller.mem_pool_host.available_size() + + # Re-applying an overlapping chain is a no-op skip, not a re-backup. + self.assertEqual(_write_backup(cache, node, write_back=True), 0) + cache.writing_check(write_back=True) + self.assertEqual( + cache.cache_controller.mem_pool_host.available_size(), host_avail + ) + + def test_shallower_crossing_backs_up_above_backuped_middle(self): + """A shallower crossing node above a backuped middle must back up in + the same insert as the deeper crossing (its own walk barrier).""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + top = next(iter(cache.root_node.children.values())) + + # A storage-prefetch completion host-inserts a backuped node below the + # still-unbacked top, legitimately breaking backup continuity. + host_indices = cache.cache_controller.mem_pool_host.alloc(8) + host_result = cache.tree_core.insert_host( + cache.root_node.id, + RadixKey(array("q", list(range(1, 9)))), + host_indices, + [f"h{i}" for i in range(8)], + ) + cache.cache_controller.mem_pool_host.free( + host_indices[: host_result.prefix_len] + ) + middle = next(iter(top.children.values())) + self.assertTrue(middle.backuped) + self.assertFalse(top.backuped) + + # The device insert unevicts the middle and adds the deep leaf. + self._insert(cache, allocator, req_to_token_pool, list(range(1, 13))) + deep = next(iter(middle.children.values())) + + cache.write_through_threshold = min(top.hit_count, deep.hit_count) + 1 + self._insert(cache, allocator, req_to_token_pool, list(range(1, 17))) + cache.writing_check(write_back=True) + self.assertTrue(top.backuped) + self.assertTrue(deep.backuped) + + def test_evict_drains_collected_frees_when_walk_raises(self): + """A device-eviction walk that raises mid-way must still free the + already-collected slots via the finally drain.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + self._insert(cache, allocator, req_to_token_pool, [10, 11, 12, 13]) + available = allocator.available_size() + + real_next = cache.tree_core.evict_device_next_node + calls = [] + + def raise_on_second(*args, **kwargs): + calls.append(args) + if len(calls) == 2: + raise RuntimeError("boom") + return real_next(*args, **kwargs) + + with mock.patch.object( + cache.tree_core, "evict_device_next_node", side_effect=raise_on_second + ): + with self.assertRaises(RuntimeError): + cache.evict(EvictParams(num_tokens=8)) + self.assertEqual(allocator.available_size(), available + 4) + + def test_match_split_relocation_survives_finalizer_failure(self): + """A match-walk split's pending write-through relocation applies before + the finalizers, so a finalizer failure cannot strand the stale record.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = 1 + # The leaf backs up on insert; its ack stays pending (no writing_check). + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + self.assertTrue(cache.ongoing_write_through) + + full_comp = cache.components[ComponentType.FULL] + with mock.patch.object( + full_comp, + "finalize_match_result_in_cache", + side_effect=RuntimeError("boom"), + ): + with self.assertRaises(RuntimeError): + cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", [1, 2])))) + + # The relocation reached the pending record: the ack clears the + # pending marker on both split halves, not just the stale node. + cache.writing_check(write_back=True) + parent = next(iter(cache.root_node.children.values())) + (child,) = parent.children.values() + self.assertIsNone(parent.write_through_pending_id) + self.assertIsNone(child.write_through_pending_id) + cache.sanity_check() + + +@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") +class TestResumableInsertWalkSWA(_InsertWalkSuite): + cfg = CacheConfig( + components=(ComponentType.FULL, ComponentType.SWA), sliding_window_size=8 + ) + + def test_swa_recovery_keeps_recovered_node_below_window_nodes(self): + """A tombstone recovered during the walk lands below the in-window path + in the SWA LRU, so eviction takes the recovered span first.""" + + sw = self.cfg.sliding_window_size + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + seq = list(range(1, 2 * sw + 1)) + key = RadixKey(array("q", seq)) + cache.insert( + InsertParams( + key=key, value=self._alloc(allocator, len(seq)), swa_evicted_seqlen=sw + ) + ) + prefix_node = next(iter(cache.root_node.children.values())) + window_node = next(iter(prefix_node.children.values())) + self.assertIsNone(prefix_node.component_data[ComponentType.SWA].value) + + # Re-inserting fully in-window recovers the prefix span's SWA data. + cache.insert( + InsertParams( + key=key, value=self._alloc(allocator, len(seq)), swa_evicted_seqlen=0 + ) + ) + self.assertIsNotNone(prefix_node.component_data[ComponentType.SWA].value) + + # SWA eviction takes the recovered span and keeps the window leaf. + cache.evict(EvictParams(num_tokens=0, swa_num_tokens=sw)) + self.assertIsNone(prefix_node.component_data[ComponentType.SWA].value) + self.assertIsNotNone(window_node.component_data[ComponentType.SWA].value) + self.assertIsNotNone(window_node.component_data[ComponentType.FULL].value) + cache.sanity_check() + + def test_dec_swa_lock_only_early_release_keeps_full_lock(self): + """The scheduler's early SWA release (decode past the window) drops + only the SWA lock; the Full path lock stays held until dec_lock_ref.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + seq = self._make_seq(1, self.cfg.sliding_window_size + 4) + self._insert(cache, allocator, req_to_token_pool, seq) + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + node = cache.resolve_node_handle(m.last_device_node) + swa_cd = node.component_data[ComponentType.SWA] + full_cd = node.component_data[ComponentType.FULL] + + lock_result = cache.inc_lock_ref(node.id) + self.assertGreaterEqual(swa_cd.lock_ref, 1) + cache.dec_swa_lock_only(node.id, lock_result.swa_uuid_for_lock) + self.assertEqual(swa_cd.lock_ref, 0) + self.assertGreaterEqual(full_cd.lock_ref, 1) + + cache.dec_lock_ref(node.id, DecLockRefParams(swa_uuid_for_lock=None)) + cache.sanity_check() + + +@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") +class TestResumableInsertWalkWriteBack(_InsertWalkSuite): + cfg = CacheConfig() + + def test_drop_fallback_frees_host_for_later_backups_same_round(self): + """Host slots reclaimed by the drop fallback itself (an interior + host-only descendant) must be reusable by later write-back backups in + the same eviction round (pre-split freed them inline).""" + if self._skip_unsupported_hicache_test(): + return + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_hicache(cache, write_policy="write_back") + host_pool = cache.cache_controller.mem_pool_host + + # Chain 1: unbacked parent -> host-only child -> host-only grandchild. + p1 = self._make_seq(1, 2) + self._insert(cache, allocator, req_to_token_pool, p1) + c1 = p1 + self._make_seq(1000, 2) + self._insert(cache, allocator, req_to_token_pool, c1) + g1 = c1 + self._make_seq(2000, 2) + self._insert(cache, allocator, req_to_token_pool, g1) + cache.evict(EvictParams(num_tokens=4)) + m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", p1)))) + parent = cache.resolve_node_handle(m.last_device_node) + (child,) = parent.children.values() + (grandchild,) = child.children.values() + self.assertTrue(child.evicted and child.backuped) + self.assertTrue(grandchild.evicted and grandchild.backuped) + + # Chain 2: a younger unbacked leaf needing 4 host slots; only 2 can + # come from evict_host (the grandchild leaf) — the other 2 exist only + # if the drop fallback's child slots are drained within the round. + p2 = self._make_seq(5000, 4) + self._insert(cache, allocator, req_to_token_pool, p2) + leaf2 = None + for node in cache.root_node.children.values(): + if list(node.key.token_ids[: len(p2)]) == list(p2): + leaf2 = node + self.assertIsNotNone(leaf2) + self.assertIsNotNone(host_pool.alloc(host_pool.available_size())) + + real_write = cache.cache_controller.write + calls = [] + + def fail_first(*args, **kwargs): + calls.append(args) + if len(calls) == 1: + return None + return real_write(*args, **kwargs) + + with mock.patch.object(cache.cache_controller, "write", side_effect=fail_first): + result = cache.evict(EvictParams(num_tokens=len(p1) + len(p2))) + self.assertGreaterEqual(result.num_tokens_evicted, len(p1) + len(p2)) + self.assertTrue(leaf2.evicted and leaf2.backuped) + cache.sanity_check() + + +@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") +class TestReturnedValuesDrain(_InsertWalkSuite): + """Drain contract of the returned-values eviction API: every tree-core step + result is drained exactly once, in per-component insertion order.""" + + cfg = CacheConfig() + + def test_each_eviction_step_result_is_drained(self): + """Every step wrapper hands its result's frees to _free_values and + passes the payload through; a wrapper that forgets strands pool slots.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + node = next(iter(cache.root_node.children.values())) + tracker = {ct: 0 for ct in cache.tree_components} + sentinel = torch.tensor([7], dtype=torch.int64) + + def make(result_cls, **fields): + result = result_cls(**fields) + result.device_frees[ComponentType.FULL].append(sentinel) + result.host_frees[ComponentType.FULL].append(sentinel) + return result + + cases = [ + ( + "evict_device_next_node", + lambda: make(EvictDeviceNextNodeResult, node_id=node.id), + lambda: cache._evict_device_next_node(ComponentType.FULL, tracker), + node.id, + ), + ( + "evict_device_leaf", + lambda: make(EvictDeviceLeafResult), + lambda: cache._evict_device_leaf(node.id, tracker), + None, + ), + ( + "demote", + lambda: make(DemoteResult), + lambda: cache._demote(node.id, tracker), + None, + ), + ( + "drop_subtree_no_host", + lambda: make(DropSubtreeNoHostResult, is_dropped=True), + lambda: cache._drop_subtree_no_host(node.id, tracker), + True, + ), + ( + "drive_host_eviction", + lambda: make(DriveHostEvictionResult), + lambda: cache.evict_host(4), + 0, + ), + ( + "dec_swa_lock_only", + lambda: make(DecSwaLockOnlyResult), + lambda: cache.dec_swa_lock_only(node.id), + None, + ), + ] + for name, make_result, call, expected in cases: + with self.subTest(step=name): + drained = [] + + def record(device_frees, host_frees): + drained.append((dict(device_frees), dict(host_frees))) + device_frees.clear() + host_frees.clear() + + with mock.patch.object( + cache.tree_core, name, return_value=make_result() + ), mock.patch.object(cache, "_free_values", side_effect=record): + returned = call() + self.assertEqual(returned, expected) + ((device_frees, host_frees),) = drained + self.assertIs(device_frees[ComponentType.FULL][0], sentinel) + self.assertIs(host_frees[ComponentType.FULL][0], sentinel) + + def test_free_values_frees_in_component_insertion_order(self): + """Device frees apply before host frees, each in per-component + insertion order — the allocator free-list order the pre-split inline + frees produced.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + order = [ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA] + device_frees = defaultdict(list) + host_frees = defaultdict(list) + for ct in order: + device_frees[ct].append(torch.tensor([1])) + host_frees[ct].append(torch.tensor([2])) + + freed = [] + fake_components = { + ct: mock.MagicMock( + free_host_values=mock.MagicMock( + side_effect=lambda values, ct=ct: freed.append(("host", ct)) + ) + ) + for ct in order + } + with mock.patch.object( + cache, + "_apply_cache_action", + side_effect=lambda action: freed.append(("device", action.component_type)), + ), mock.patch.dict(cache.components, fake_components): + cache._free_values(device_frees, host_frees) + + self.assertEqual( + freed, [("device", ct) for ct in order] + [("host", ct) for ct in order] + ) + self.assertFalse(device_frees) + self.assertFalse(host_frees) + + def test_free_values_mid_drain_failure_cannot_replay_freed_entries(self): + """A device free that raises must leave only un-attempted entries in + the dict (no replay of successes) while host frees still drain.""" + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + order = [ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA] + device_frees = defaultdict(list) + host_frees = defaultdict(list) + for ct in order: + device_frees[ct].append(torch.tensor([1])) + host_frees[ComponentType.FULL].append(torch.tensor([2])) + + def boom_on_swa(action): + if action.component_type is ComponentType.SWA: + raise RuntimeError("boom") + + host_mock = mock.MagicMock() + with mock.patch.object( + cache, "_apply_cache_action", side_effect=boom_on_swa + ), mock.patch.dict(cache.components, {ComponentType.FULL: host_mock}): + with self.assertRaises(RuntimeError): + cache._free_values(device_frees, host_frees) + + self.assertEqual(list(device_frees), [ComponentType.MAMBA]) + self.assertFalse(host_frees) + host_mock.free_host_values.assert_called_once() + + def test_undrained_result_trips_the_del_assert(self): + """Dropping a result without draining fires the __del__ tripwire (the + only forgotten-drain detection); a drained result stays silent.""" + seen = [] + old_hook = sys.unraisablehook + sys.unraisablehook = lambda unraisable: seen.append(unraisable.exc_value) + try: + undrained = DemoteResult() + undrained.device_frees[ComponentType.FULL].append(torch.tensor([1])) + del undrained + + drained = DemoteResult() + drained.device_frees[ComponentType.FULL].append(torch.tensor([1])) + drained.device_frees.clear() + del drained + finally: + sys.unraisablehook = old_hook + self.assertEqual(len(seen), 1) + self.assertIsInstance(seen[0], AssertionError) + + def test_node_info_facades_forward_to_the_tree(self): + """The scheduler's storage-prefetch facades (is_backuped / is_root / + get_last_hash_value / get_prefix_hash_values) resolve NodeIds tree-side + with node-level semantics.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + self._insert(cache, allocator, req_to_token_pool, list(range(1, 9))) + parent = next(iter(cache.root_node.children.values())) + (child,) = parent.children.values() + + self.assertTrue(cache.is_root(cache.root_node_handle())) + self.assertFalse(cache.is_root(child.id)) + self.assertFalse(cache.is_backuped(parent.id)) + self.assertIsNone(cache.get_last_hash_value(parent.id)) + + self.assertGreater(_write_backup(cache, parent, write_back=True), 0) + cache.writing_check(write_back=True) + self.assertTrue(cache.is_backuped(parent.id)) + + parent.hash_value = ["h1", "h2"] + self.assertEqual(cache.get_last_hash_value(parent.id), "h2") + # The prefix chain carries the ancestors' hashes, not the node's own. + self.assertEqual(cache.get_prefix_hash_values(parent.id), []) + self.assertEqual(cache.get_prefix_hash_values(child.id), ["h1", "h2"]) + + def test_evict_host_drains_freed_host_values_to_the_pool(self): + """Host eviction's returned frees must reach the host pool in the same + call; a dropped drain leaves the pool permanently smaller.""" + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + leaf = next(iter(cache.root_node.children.values())) + self.assertGreater(_write_backup(cache, leaf, write_back=True), 0) + cache.writing_check(write_back=True) + cache.evict(EvictParams(num_tokens=4)) + self.assertTrue(leaf.evicted) + + host_pool = cache.cache_controller.mem_pool_host + available_before = host_pool.available_size() + evicted = cache.evict_host(4) + self.assertGreater(evicted, 0) + self.assertGreater(host_pool.available_size(), available_before) + cache.sanity_check() + + +class TestPrefetchCommitOrdering(CustomTestCase): + """The prefetch commit's action ordering (mock-based).""" + + def test_prefetch_commit_applies_host_insert_actions_before_transfers(self): + """The prefetch commit applies the host-insert walk's actions before + commit_hicache_transfers, whose emissions ride a fresh list.""" + cache = mock.MagicMock() + cache.page_size = 1 + cache.enable_storage_metrics = False + walk_action = object() + insert_result = mock.MagicMock() + insert_result.cache_actions = [walk_action] + insert_result.prefix_len = 4 + cache.tree_core.insert_host.return_value = insert_result + cache.ongoing_prefetch = { + "req": ( + 7, + list(range(8)), + list(range(100, 108)), + mock.MagicMock(), + None, + {}, + ) + } + cache.cache_controller.terminate_prefetch.return_value = ( + 8, + [f"h{i}" for i in range(8)], + ) + cache._sync_and_check_hybrid_prefetch_result.return_value = 8 + cache.cache_controller.prefetch_tokens_occupied = 100 + cache.prefetch_loaded_tokens_by_reqid = {} + + order = mock.MagicMock() + applied = [] + + def record_apply(actions): + applied.append(list(actions)) + actions.clear() + + order.apply.side_effect = record_apply + cache._apply_cache_actions = order.apply + cache.tree_core.commit_hicache_transfers = order.commit + + self.assertTrue(UnifiedRadixCache.check_prefetch_progress(cache, "req")) + + self.assertEqual([c[0] for c in order.mock_calls], ["apply", "commit", "apply"]) + self.assertEqual(applied[0], [walk_action]) + self.assertIsNot( + order.commit.call_args.kwargs["cache_actions"], insert_result.cache_actions + ) + self.assertEqual(cache.ongoing_prefetch, {}) + + if __name__ == "__main__": unittest.main()