From eb5f0fbeef5b90aa522abf0103f6e0eaf00c29d9 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Wed, 6 May 2026 22:19:25 +0800 Subject: [PATCH] Support swa HiCache for unified radix cache (#23391) Co-authored-by: hzh0425 --- .../sglang/srt/mem_cache/base_prefix_cache.py | 4 + .../sglang/srt/mem_cache/hicache_storage.py | 1 + .../hybrid_cache/hybrid_cache_controller.py | 54 ++-- .../hybrid_cache/hybrid_pool_assembler.py | 155 ++++++++- .../sglang/srt/mem_cache/memory_pool_host.py | 7 + .../sglang/srt/mem_cache/swa_memory_pool.py | 32 ++ .../mamba_component.py | 2 +- .../unified_cache_components/swa_component.py | 286 +++++++++++++++-- .../tree_component.py | 13 + .../srt/mem_cache/unified_radix_cache.py | 54 +++- .../test_unified_radix_cache_unittest.py | 299 +++++++++++++++++- 11 files changed, 814 insertions(+), 93 deletions(-) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 999f46f87..4316697ef 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -95,6 +95,10 @@ class IncLockRefResult: delta: Optional[int] = None swa_uuid_for_lock: Optional[int] = None + def to_dec_params(self) -> "DecLockRefParams": + """Convert to the corresponding DecLockRefParams for dec_lock_ref.""" + return DecLockRefParams(swa_uuid_for_lock=self.swa_uuid_for_lock) + @dataclasses.dataclass class DecLockRefParams: diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 21bb4a9ec..203b0ec60 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -41,6 +41,7 @@ class PoolName(str, Enum): KV = "kv" MAMBA = "mamba" + SWA = "swa" INDEXER = "indexer" def __str__(self) -> str: diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 86c5db24d..456bd72c8 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging import threading import time -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Callable, List, Optional import torch @@ -296,12 +296,18 @@ class HybridCacheController(BaseHiCacheController): extra_pools: Optional[list[PoolTransfer]] = None, ) -> Optional[torch.Tensor]: need_load_kv = host_indices.numel() > 0 - if need_load_kv: - device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) + + full_allocator = getattr( + self.mem_pool_device_allocator, + "full_attn_allocator", + self.mem_pool_device_allocator, + ) + if not need_load_kv: + device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) + else: + device_indices = full_allocator.alloc(len(host_indices)) if device_indices is None: return None - else: - device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) pool_transfers = self._resolve_pool_transfers_allocation( extra_pools, @@ -311,7 +317,7 @@ class HybridCacheController(BaseHiCacheController): ) if pool_transfers is None and extra_pools: if need_load_kv: - self.mem_pool_device_allocator.free(device_indices) + full_allocator.free(device_indices) return None self.load_queue.append( @@ -535,7 +541,8 @@ class HybridCacheController(BaseHiCacheController): """Auto-alloc host or device indices for PoolTransfers where they are None.""" if not extra_pools: return None - newly_allocated: list[tuple[PoolTransfer, Any, torch.Tensor]] = [] + # (pool, free_fn, indices) for atomic rollback on failure. + newly_allocated: list[tuple[PoolTransfer, Callable, torch.Tensor]] = [] for pool in extra_pools: entry = self.mem_pool_host.entry_map.get(pool.name) if entry is None: @@ -547,27 +554,28 @@ class HybridCacheController(BaseHiCacheController): if alloc_host: if pool.host_indices is not None or pool.device_indices is None: continue - entry_pool, evict_fn, size = ( - entry.host_pool, - entry.host_evict_fn, - len(pool.device_indices), - ) + alloc_fn = entry.host_pool.alloc + free_fn = entry.host_pool.free + evict_fn = entry.host_evict_fn + size = len(pool.device_indices) else: if pool.device_indices is not None or pool.host_indices is None: continue - entry_pool, evict_fn, size = ( - entry.device_pool, - entry.device_evict_fn, - len(pool.host_indices), - ) - indices = entry_pool.alloc(size) + # device_alloc_fn / device_free_fn override entry.device_pool's + # methods for pools whose device_pool is a raw KV pool (layout) + # rather than an allocator (e.g. SWA). + alloc_fn = entry.device_alloc_fn or entry.device_pool.alloc + free_fn = entry.device_free_fn or entry.device_pool.free + evict_fn = entry.device_evict_fn + size = len(pool.host_indices) + indices = alloc_fn(size) if indices is None and evict_fn: evict_fn(size) - indices = entry_pool.alloc(size) + indices = alloc_fn(size) if indices is None: - # Roll back all previous allocations using each pool's own entry_pool. - for prev_pool, prev_entry_pool, prev_indices in newly_allocated: - prev_entry_pool.free(prev_indices) + # Atomic rollback: free everything we successfully allocated. + for prev_pool, prev_free_fn, prev_indices in newly_allocated: + prev_free_fn(prev_indices) if alloc_host: prev_pool.host_indices = None else: @@ -577,5 +585,5 @@ class HybridCacheController(BaseHiCacheController): pool.host_indices = indices else: pool.device_indices = indices - newly_allocated.append((pool, entry_pool, indices)) + newly_allocated.append((pool, free_fn, indices)) return extra_pools diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index d88ba3da2..c4d857b5d 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -74,6 +74,8 @@ def build_pool_entry( share_indices_with_anchor: bool = False, host_evict_fn: Optional[Callable[[int], Any]] = None, device_evict_fn: Optional[Callable[[int], Any]] = None, + device_alloc_fn: Optional[Callable[[int], Any]] = None, + device_free_fn: Optional[Callable[[Any], Any]] = None, ) -> PoolEntry: return PoolEntry( name=name, @@ -84,6 +86,8 @@ def build_pool_entry( share_indices_with_anchor=share_indices_with_anchor, host_evict_fn=host_evict_fn, device_evict_fn=device_evict_fn, + device_alloc_fn=device_alloc_fn, + device_free_fn=device_free_fn, ) @@ -153,6 +157,90 @@ def build_kv_only_stack( return host_pool_group, cache_controller +def build_hybrid_swa_stack( + *, + params: CacheInitParams, + server_args: ServerArgs, + full_kv_pool: Any, + swa_kv_pool: Any, + full_layer_mapping: dict[int, int], + swa_layer_mapping: dict[int, int], + page_size: int, + tp_group, + load_cache_event, + storage_backend: Optional[str], + use_mla: bool, + host_swa_evict_fn: Optional[Callable[[int], Any]] = None, + device_swa_evict_fn: Optional[Callable[[int], Any]] = None, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + pp_rank: int = 0, + pp_size: int = 1, + attn_cp_rank: int = 0, + attn_cp_size: int = 1, + enable_storage_metrics: bool = False, +) -> tuple[HostPoolGroup, HybridCacheController]: + transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) + kv_host_pool = build_kv_host_pool( + kv_pool=full_kv_pool, + page_size=page_size, + server_args=server_args, + use_mla=use_mla, + ) + swa_host_pool = build_kv_host_pool( + kv_pool=swa_kv_pool, + page_size=page_size, + server_args=server_args, + use_mla=use_mla, + ) + + # For SWA hybrid, the device alloc/free goes through the inner swa_attn_allocator + swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator + entries = [ + build_pool_entry( + name=PoolName.KV, + host_pool=kv_host_pool, + device_pool=full_kv_pool, + layer_mapping=full_layer_mapping, + transfer_layer_num=transfer_layer_num, + is_anchor=True, + ), + build_pool_entry( + name=PoolName.SWA, + host_pool=swa_host_pool, + device_pool=swa_kv_pool, + layer_mapping=swa_layer_mapping, + transfer_layer_num=transfer_layer_num, + host_evict_fn=host_swa_evict_fn, + device_evict_fn=device_swa_evict_fn, + device_alloc_fn=swa_attn_allocator.alloc, + device_free_fn=swa_attn_allocator.free, + ), + ] + host_pool_group = HostPoolGroup(entries) + cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + page_size, + tp_group, + load_cache_event=load_cache_event, + write_policy=server_args.hicache_write_policy, + io_backend=server_args.hicache_io_backend, + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + pp_rank=pp_rank, + pp_size=pp_size, + attn_cp_rank=attn_cp_rank, + attn_cp_size=attn_cp_size, + transfer_layer_num=transfer_layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + return host_pool_group, cache_controller + + def build_hybrid_mamba_stack( *, params: CacheInitParams, @@ -330,17 +418,29 @@ def attach_hybrid_pool_to_unified_cache( MLATokenToKVPool, NSATokenToKVPool, ) + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_cache_components import ComponentType try: kvcache = params.token_to_kv_pool_allocator.get_kvcache() - if isinstance(kvcache, HybridLinearKVPool): + swa_stack = isinstance(kvcache, SWAKVPool) + mamba_stack = isinstance(kvcache, HybridLinearKVPool) + nsa_stack = isinstance(kvcache, NSATokenToKVPool) + + if mamba_stack: full_kv_pool = kvcache.full_kv_pool use_mla = kvcache.use_mla assert set(cache.components.keys()) == { ComponentType.FULL, ComponentType.MAMBA, }, "HybridLinearKVPool currently only supports FULL + MAMBA in UnifiedRadixCache." + elif swa_stack: + full_kv_pool = kvcache.full_kv_pool + use_mla = False + assert set(cache.components.keys()) == { + ComponentType.FULL, + ComponentType.SWA, + }, "SWAKVPool currently only supports FULL + SWA in UnifiedRadixCache." else: full_kv_pool = kvcache use_mla = isinstance(kvcache, MLATokenToKVPool) @@ -348,8 +448,6 @@ def attach_hybrid_pool_to_unified_cache( ComponentType.FULL }, "Non-hybrid KV pool currently only supports FULL-only UnifiedRadixCache." - mamba_stack = isinstance(kvcache, HybridLinearKVPool) - nsa_stack = isinstance(kvcache, NSATokenToKVPool) if mamba_stack: full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping) mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) @@ -386,6 +484,47 @@ def attach_hybrid_pool_to_unified_cache( cache_controller.layer_done_counter ) transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping) + elif swa_stack: + full_layer_mapping = { + global_id: local_id + for global_id, (local_id, is_swa) in kvcache.layers_mapping.items() + if not is_swa + } + swa_layer_mapping = { + global_id: local_id + for global_id, (local_id, is_swa) in kvcache.layers_mapping.items() + if is_swa + } + host_pool_group, cache_controller = build_hybrid_swa_stack( + params=params, + server_args=server_args, + full_kv_pool=full_kv_pool, + swa_kv_pool=kvcache.swa_kv_pool, + full_layer_mapping=full_layer_mapping, + swa_layer_mapping=swa_layer_mapping, + page_size=cache.page_size, + tp_group=params.tp_cache_group, + load_cache_event=load_cache_event, + storage_backend=None, + use_mla=False, + host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA), + device_swa_evict_fn=lambda n: cache.evict( + EvictParams(swa_num_tokens=n) + ), + pp_rank=params.pp_rank, + pp_size=params.pp_size, + ) + cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) + cache.host_pool_group = host_pool_group + cache.cache_controller = cache_controller + cache.components[ComponentType.FULL]._full_kv_pool_host = ( + cache.full_kv_pool_host + ) + cache.swa_kv_pool_host = host_pool_group.get_pool(PoolName.SWA) + cache.components[ComponentType.SWA]._swa_kv_pool_host = ( + cache.swa_kv_pool_host + ) + transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) elif nsa_stack: full_layer_mapping = { layer_id: layer_id for layer_id in range(full_kv_pool.layer_num) @@ -456,9 +595,17 @@ def attach_hybrid_pool_to_unified_cache( cache.cache_controller.layer_done_counter ) + if mamba_stack: + pools_desc = "KV + MAMBA" + elif swa_stack: + pools_desc = "KV + SWA" + elif nsa_stack: + pools_desc = "KV + INDEXER" + else: + pools_desc = "KV" logger.info( "Attached hybrid pool stack to UnifiedRadixCache: pools=%s, transfer_layer_num=%s", - "KV + MAMBA" if mamba_stack else "KV + INDEXER" if nsa_stack else "KV", + pools_desc, transfer_layer_num, ) except Exception: diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index f45c1d616..37fc5b4cc 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1671,6 +1671,13 @@ class PoolEntry: # device_evict_fn(n): evict n slots from the device pool (used by load()). host_evict_fn: Optional[Callable] = None device_evict_fn: Optional[Callable] = None + # Optional alloc/free overrides for the device side, used by + # _resolve_pool_transfers_allocation. Set when entry.device_pool is the + # raw KV pool (layout) rather than an allocator (e.g. SWA, where alloc + # lives on a separate sub-allocator inside SWATokenToKVPoolAllocator). + # When None, fall back to entry.device_pool.alloc/free. + device_alloc_fn: Optional[Callable] = None + device_free_fn: Optional[Callable] = None class HostPoolGroup: diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 5c1012554..67d30295c 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -51,9 +51,11 @@ class SWAKVPool(KVCache): self.device = device self.swa_layer_nums = len(swa_attention_layer_ids) self.full_layer_nums = len(full_attention_layer_ids) + self.layer_num = self.full_layer_nums + self.swa_layer_nums self.start_layer = 0 self.page_size = page_size self.swa_loc = None + self.layer_transfer_counter = None kwargs["page_size"] = page_size kwargs["enable_memory_saver"] = False @@ -100,6 +102,16 @@ class SWAKVPool(KVCache): def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping + def register_layer_transfer_counter(self, layer_transfer_counter): + # Wait happens at this wrapper. Inner pools must not wait again. + self.layer_transfer_counter = layer_transfer_counter + self.full_kv_pool.register_layer_transfer_counter(None) + self.swa_kv_pool.register_layer_transfer_counter(None) + + def _wait_for_layer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + def get_kv_size_bytes(self): k_size, v_size = self.full_kv_pool.get_kv_size_bytes() k_size_swa, v_size_swa = self.swa_kv_pool.get_kv_size_bytes() @@ -123,6 +135,7 @@ class SWAKVPool(KVCache): return swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens def get_key_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] if is_swa_layer: return self.swa_kv_pool.get_key_buffer(layer_id_pool) @@ -130,6 +143,7 @@ class SWAKVPool(KVCache): return self.full_kv_pool.get_key_buffer(layer_id_pool) def get_value_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] if is_swa_layer: return self.swa_kv_pool.get_value_buffer(layer_id_pool) @@ -137,6 +151,7 @@ class SWAKVPool(KVCache): return self.full_kv_pool.get_value_buffer(layer_id_pool) def get_kv_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] if is_swa_layer: return self.swa_kv_pool.get_kv_buffer(layer_id_pool) @@ -461,6 +476,23 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): ) assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size + def set_full_to_swa_mapping( + self, full_indices: torch.Tensor, swa_indices: torch.Tensor + ) -> None: + """Write full_to_swa_index_mapping[full_indices[i]] = swa_indices[i]. + + Used by HiCache load-back path to rebuild the mapping after FULL and SWA device alloc. + """ + if full_indices.numel() == 0: + return + assert full_indices.numel() == swa_indices.numel() + if _is_npu: + self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = ( + swa_indices.to(torch.int64) + ) + else: + self.full_to_swa_index_mapping[full_indices] = swa_indices + def free_swa(self, free_index: torch.Tensor): swa_indices = self.full_to_swa_index_mapping[free_index] swa_indices = swa_indices[swa_indices > 0] 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 9f75352e5..59dfc4768 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 @@ -347,7 +347,7 @@ class MambaComponent(TreeComponent): if cd.value is not None: return None - # restore single node if host_value exists and + # restore single node if host_value exists if cd.host_value is not None and cd.value is None: transfers.append( PoolTransfer( 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 4eccfecfc..ec78978a1 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 @@ -10,9 +10,13 @@ from sglang.srt.mem_cache.base_prefix_cache import ( IncLockRefResult, InsertParams, InsertResult, + MatchPrefixParams, + MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer from sglang.srt.mem_cache.unified_cache_components.tree_component import ( BASE_COMPONENT_TYPE, + CacheTransferPhase, ComponentType, EvictLayer, TreeComponent, @@ -46,6 +50,8 @@ class SWAComponent(TreeComponent): ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(cache.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 + self._swa_kv_pool_host = None component_type = ComponentType.SWA @@ -54,13 +60,25 @@ class SWAComponent(TreeComponent): full_indices ) + 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 create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: sliding_window_size = self.sliding_window_size ct = self.component_type state = {"len": float("inf")} def validator(node: UnifiedTreeNode) -> bool: - if node.component_data[ct].value is None: + cd = node.component_data[ct] + # HiCache: a host-only tombstone is a valid match boundary too + # — load_back will restore SWA from host before use. + if cd.value is None and cd.host_value is None: state["len"] = 0 return False state["len"] += len(node.key) @@ -68,6 +86,30 @@ class SWAComponent(TreeComponent): return validator + def finalize_match_result( + self, + result: MatchResult, + params: MatchPrefixParams, + value_chunks: list[torch.Tensor], + best_value_len: int, + ) -> MatchResult: + ct = self.component_type + n_swa = 0 + node = result.last_device_node + root = self.cache.root_node + while node is not root and n_swa < self.sliding_window_size: + cd = node.component_data[ct] + if cd.value is None and cd.host_value is not None: + return result._replace(host_hit_length=max(result.host_hit_length, 1)) + if cd.value is not None: + n_swa += len(cd.value) + elif cd.host_value is not None: + n_swa += len(cd.host_value) + else: + break + node = node.parent + return result + def update_component_on_insert_overlap( self, node: UnifiedTreeNode, @@ -100,9 +142,7 @@ class SWAComponent(TreeComponent): 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) + self._restore_device_value(node, swa_value) return 0 elif swa_evicted_seqlen < total_prefix_len + prefix_len: # Branch 2: value_slice[start_idx:] is within SWA window — partial recover @@ -117,9 +157,7 @@ class SWAComponent(TreeComponent): 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) + self._restore_device_value(node, swa_value) return start_idx else: # Branch 3: entire value_slice is outside SWA window — not consumed @@ -130,6 +168,39 @@ class SWAComponent(TreeComponent): ) -> bool: return params.swa_evicted_seqlen >= total_prefix_len + key_len + def recover_after_unevict( + self, + node: UnifiedTreeNode, + prefix_len: int, + total_prefix_len: int, + params: InsertParams, + ) -> 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 + # the in-window portion. There is no old SWA slot to free here. + ct = self.component_type + if node.component_data[ct].value is not None: + return + assert ( + node.component_data[ct].lock_ref == 0 + ), 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 + ), 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) + 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) + else: + return + self._restore_device_value(node, swa_value) + def commit_insert_component_data( self, node: UnifiedTreeNode, @@ -181,6 +252,23 @@ class SWAComponent(TreeComponent): else: new_parent.component_data[self.component_type].value = None + child_swa_host_value = child.component_data[self.component_type].host_value + if child_swa_host_value is not None: + split_len = len(new_parent.key) + new_parent.component_data[self.component_type].host_value = ( + child_swa_host_value[:split_len].clone() + ) + 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] + if new_parent.component_data[self.component_type].value is None: + host_lru.insert_mru(new_parent) + if child.component_data[ + self.component_type + ].value is None and not host_lru.in_list(child): + host_lru.insert_mru(child) + # parent inherits the swa_uuid from child for swa lock ref new_parent.component_data[self.component_type].metadata["uuid"] = ( child.component_data[self.component_type].metadata.get("uuid") @@ -192,23 +280,43 @@ class SWAComponent(TreeComponent): node: UnifiedTreeNode, target: EvictLayer = EvictLayer.DEVICE, ) -> tuple[int, int]: - if target is EvictLayer.HOST: - return 0, 0 # TODO:SWA has no host layer currently + ct = self.component_type + cd = node.component_data[ct] + freed = 0 + host_freed = 0 - swa_value = node.component_data[self.component_type].value - if swa_value is None: - return 0, 0 - # Direct swa_attn_allocator.free(swa_value) would double-free - # free_swa(full_value) has the mapping guard to avoid double-free - # TODO: decoupling full and swa free, need further discussion on mapping necessity - self.cache.token_to_kv_pool_allocator.free_swa( - node.component_data[BASE_COMPONENT_TYPE].value - ) - freed = len(swa_value) - self.cache.component_evictable_size_[self.component_type] -= freed - if target is EvictLayer.DEVICE: - node.component_data[self.component_type].value = None - return freed, 0 + # Device layer + if EvictLayer.DEVICE in target and cd.value is not None: + # 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( + node.component_data[BASE_COMPONENT_TYPE].value + ) + freed = len(cd.value) + self.cache.component_evictable_size_[ct] -= freed + cd.value = None + + # Host layer + host_lru = self.cache.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) + cd.host_value = None + if host_lru.in_list(node): + host_lru.remove_node(node) + + # After device tombstone: if host_value remains, move into host LRU + if ( + target is EvictLayer.DEVICE + and cd.value is None + and cd.host_value is not None + ): + if not host_lru.in_list(node): + host_lru.insert_mru(node) + + return freed, host_freed def eviction_priority(self, is_leaf: bool) -> int: return 0 if is_leaf else 1 @@ -247,12 +355,15 @@ class SWAComponent(TreeComponent): swa_lock_size = 0 swa_uuid_for_lock = None + # Tombstoned nodes (cd.value is None) have no SWA chunk to protect + # skip them and keep walking up. This path is hit when HiCache + # backs up a FULL present internal node whose SWA was already evicted. cur = node while cur != root and swa_lock_size < sliding_window_size: - assert ( - cur.component_data[ct].value is not None - ), f"acquire_component_lock({ct}) on tombstoned node {cur.id}" comp = cur.component_data[ct] + if comp.value is None: + cur = cur.parent + continue if comp.lock_ref == 0: key_len = len(cur.key) self.cache.component_evictable_size_[ct] -= key_len @@ -276,15 +387,15 @@ class SWAComponent(TreeComponent): swa_uuid_for_lock = params.swa_uuid_for_lock if params else None dec_swa = True + # lock_ref == 0 means acquire_component_lock skipped this node + # (tombstone at acquire time) or load_back revived a tombstone between + # acquire and release. Either way, there is nothing for us to undo here. cur = node while cur != root and dec_swa: - assert ( - cur.component_data[ct].value is not None - ), f"release_component_lock({ct}) on tombstoned node {cur.id}" comp = cur.component_data[ct] - assert ( - comp.lock_ref > 0 - ), f"release_component_lock({ct}) on node with lock_ref=0, node {cur.id}" + if comp.lock_ref == 0: + cur = cur.parent + continue if comp.lock_ref == 1: key_len = len(cur.key) self.cache.component_evictable_size_[ct] += key_len @@ -304,3 +415,114 @@ class SWAComponent(TreeComponent): if is_finished: insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen return None + + # ---- HiCache Hooks ---- + + def build_hicache_transfers( + self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + ) -> Optional[list[PoolTransfer]]: + ct = self.component_type + + if phase == CacheTransferPhase.BACKUP_HOST: + cd = node.component_data[ct] + if cd.value is None: + return None + # cd.value already holds SWA-pool indices (translated at insert time). + # Host pool indexing wants int64. + return [ + PoolTransfer( + name=PoolName.SWA, + device_indices=cd.value.to(torch.int64), + ) + ] + + if phase == CacheTransferPhase.LOAD_BACK: + n_swa = 0 + backed_up: list[torch.Tensor] = [] + nodes: list = [] + while node is not self.cache.root_node and n_swa < self.sliding_window_size: + cd = node.component_data[ct] + assert cd.host_value is not None or cd.value is not None + if cd.value is not None: + # device exists, skip it + n_swa += len(cd.value) + else: + # host only, collect it + backed_up.append(cd.host_value) + nodes.append(node) + n_swa += len(cd.host_value) + node = node.parent + + if not backed_up: + return None + + backed_up.reverse() + nodes.reverse() + + return [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.cat(backed_up), + device_indices=None, + nodes_to_load=nodes, + ) + ] + + return None + + def commit_hicache_transfer( + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + transfers: list[PoolTransfer] = (), + ) -> None: + ct = self.component_type + + if phase == CacheTransferPhase.BACKUP_HOST: + if transfers and transfers[0].host_indices is not None: + cd = node.component_data[ct] + if cd.host_value is None: + cd.host_value = transfers[0].host_indices.clone() + return + + if phase == CacheTransferPhase.LOAD_BACK: + 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 + + offset = 0 + for n in xfer.nodes_to_load or []: + 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) + 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) + offset += n_tokens + assert offset == len(xfer.host_indices) + return + + def drive_host_eviction( + self, num_tokens: int, tracker: dict[ComponentType, int] + ) -> 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] + x = host_lru.get_lru_no_lock() + while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x): + x_next = host_lru.get_prev_no_lock(x) + cd = x.component_data[ct] + if x in self.cache.evictable_host_leaves: + self.cache._evict_host_leaf(x, tracker) + else: + assert cd.host_value is not None + self.cache._evict_component_and_detach_lru( + x, self, target=EvictLayer.HOST, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST) + x = x_next 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 75e3123ed..317ee3012 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 @@ -163,6 +163,19 @@ class TreeComponent(ABC): be a tombstone for this component.""" return False + def recover_after_unevict( + self, + node: UnifiedTreeNode, + prefix_len: int, + total_prefix_len: int, + params: InsertParams, + ) -> 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 + their own data from the freshly assigned base value when their entry + is still tombstoned. Default no-op.""" + return None + def commit_insert_component_data( self, node: UnifiedTreeNode, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 71e0158a1..3dbe2c212 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -275,8 +275,10 @@ class UnifiedRadixCache(BasePrefixCache): ct: UnifiedLRUList(ct, self.tree_components, use_host_ptr=True) for ct in self.tree_components } - self.ongoing_write_through: dict[int, UnifiedTreeNode] = {} - self.ongoing_load_back: dict[int, UnifiedTreeNode] = {} + self.ongoing_write_through: dict[ + int, tuple[UnifiedTreeNode, Optional[DecLockRefParams]] + ] = {} + self.ongoing_load_back: dict[int, tuple[UnifiedTreeNode, DecLockRefParams]] = {} self.enable_storage = False self.ongoing_prefetch: dict = {} self.ongoing_backup: dict = {} @@ -806,6 +808,18 @@ class UnifiedRadixCache(BasePrefixCache): 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 @@ -1185,9 +1199,10 @@ class UnifiedRadixCache(BasePrefixCache): transfers=xfers, ) - self.ongoing_write_through[node.id] = node + lock_params = None if not write_back: - self.inc_lock_ref(node) + lock_params = self.inc_lock_ref(node).to_dec_params() + self.ongoing_write_through[node.id] = (node, lock_params) return len(host_indices) def load_back( @@ -1210,6 +1225,7 @@ class UnifiedRadixCache(BasePrefixCache): nodes_to_load = kv_xfer.nodes_to_load ancestor_node = nodes_to_load[0].parent if nodes_to_load else last_hit_node result = self.inc_lock_ref(ancestor_node) + ancestor_lock_params = result.to_dec_params() kv_tokens = len(kv_xfer.host_indices) # Build aux transfers, keyed per component. @@ -1233,7 +1249,7 @@ class UnifiedRadixCache(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(ancestor_node) + self.dec_lock_ref(ancestor_node, ancestor_lock_params) return None avail = self.token_to_kv_pool_allocator.available_size() @@ -1241,7 +1257,7 @@ class UnifiedRadixCache(BasePrefixCache): needed = kv_tokens - avail result = self.evict(EvictParams(num_tokens=needed)) if result.num_tokens_evicted < needed: - self.dec_lock_ref(ancestor_node) + self.dec_lock_ref(ancestor_node, ancestor_lock_params) return None # Load H→D @@ -1253,7 +1269,7 @@ class UnifiedRadixCache(BasePrefixCache): extra_pools=aux_xfers or None, ) - self.dec_lock_ref(ancestor_node) + self.dec_lock_ref(ancestor_node, ancestor_lock_params) if device_indices is None: return None @@ -1272,8 +1288,10 @@ class UnifiedRadixCache(BasePrefixCache): ) self._update_evictable_leaf_sets(ancestor_node) - self.inc_lock_ref(last_hit_node) - self.ongoing_load_back[last_hit_node.id] = last_hit_node + self.ongoing_load_back[last_hit_node.id] = ( + last_hit_node, + self.inc_lock_ref(last_hit_node).to_dec_params(), + ) return device_indices def _inc_hit_count(self, node: UnifiedTreeNode, chunked: bool = False) -> None: @@ -1302,7 +1320,11 @@ class UnifiedRadixCache(BasePrefixCache): for _, finish_event, ack_list in cc.ack_write_queue: finish_event.synchronize() for ack_id in ack_list: - self.ongoing_write_through.pop(ack_id, None) + entry = self.ongoing_write_through.pop(ack_id, None) + if entry is not None: + node, params = entry + if params is not None: + self.dec_lock_ref(node, params) cc.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -1329,8 +1351,8 @@ class UnifiedRadixCache(BasePrefixCache): _, finish_event, ack_list = cc.ack_write_queue.pop(0) finish_event.synchronize() for ack_id in ack_list: - node = self.ongoing_write_through.pop(ack_id) - self.dec_lock_ref(node) + node, params = self.ongoing_write_through.pop(ack_id) + self.dec_lock_ref(node, params) finish_count -= 1 def loading_check(self) -> None: @@ -1344,8 +1366,8 @@ class UnifiedRadixCache(BasePrefixCache): break finish_count += 1 for ack_id in ack_list: - node = self.ongoing_load_back.pop(ack_id) - self.dec_lock_ref(node) + node, lock_params = self.ongoing_load_back.pop(ack_id) + self.dec_lock_ref(node, lock_params) del cc.ack_load_queue[:finish_count] # ---- HiCache: Scheduler Entry Points ---- @@ -1752,14 +1774,14 @@ class UnifiedRadixCache(BasePrefixCache): ) # ── PART 5: Ongoing Operations ── - for nid, n in self.ongoing_write_through.items(): + 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(): + 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: 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 66bf55994..10cf822f0 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 @@ -17,9 +17,11 @@ from sglang.srt.mem_cache.base_prefix_cache import ( EvictResult, InsertParams, MatchPrefixParams, + MatchResult, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import available_and_evictable_str +from sglang.srt.mem_cache.hicache_storage import PoolName from sglang.srt.mem_cache.memory_pool import ( HybridLinearKVPool, HybridReqToTokenPool, @@ -28,7 +30,10 @@ 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, SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.unified_cache_components.tree_component import ComponentType +from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + CacheTransferPhase, + ComponentType, +) from sglang.srt.mem_cache.unified_radix_cache import ( UnifiedRadixCache, UnifiedTreeNode, @@ -1201,6 +1206,24 @@ class UnifiedRadixCacheSuite: self.skipTest("HiCache tests do not run on SWA stacks") return False + def _simulate_backup(self, tree, node): + """Simulate D->H backup by setting host_value on each component.""" + for ct in (ComponentType.FULL, ComponentType.MAMBA, ComponentType.SWA): + if ct not in self.cfg.components: + continue + cd = node.component_data[ct] + if cd.value is not None and cd.host_value is None: + cd.host_value = cd.value.clone() + + def _simulate_backup_tree(self, tree): + """Backup all non-root nodes (simulates write-through).""" + stack = [tree.root_node] + while stack: + node = stack.pop() + if node is not tree.root_node: + self._simulate_backup(tree, node) + stack.extend(node.children.values()) + def _init_hicache(self, tree): import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler @@ -1516,33 +1539,275 @@ class UnifiedRadixCacheSuite: tree.sanity_check() def test_hicache_evict_to_host_updates_aux_lru(self): - """Aux components move from device LRU to host LRU on device-to-host eviction.""" - if self._skip_unsupported_hicache_test(): - return - if not self.cfg.has_mamba: - self.skipTest("requires Mamba component") - tree, allocator, req_to_token_pool = self._build_hicache_fixture() + """Aux components (MAMBA / SWA) move from device LRU to host LRU on D->H eviction.""" + aux_types = [ + ct + for ct in (ComponentType.MAMBA, ComponentType.SWA) + if ct in self.cfg.components + ] + if not aux_types: + self.skipTest("requires at least one aux component") + + tree, allocator, req_to_token_pool = build_fixture(self.cfg) seq = self._make_seq(1, 2) self._insert(tree, allocator, req_to_token_pool, seq) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) node = m.last_device_node - # Check mamba is in device LRU - mamba_lru = tree.lru_lists[ComponentType.MAMBA] - host_mamba_lru = tree.host_lru_lists[ComponentType.MAMBA] - self.assertTrue(mamba_lru.in_list(node)) - self.assertFalse(host_mamba_lru.in_list(node)) + for aux in aux_types: + self.assertTrue(tree.lru_lists[aux].in_list(node)) + self.assertFalse(tree.host_lru_lists[aux].in_list(node)) - self._backup_node(tree, node) + self._simulate_backup(tree, node) tree.evict(EvictParams(num_tokens=len(seq))) - # Mamba should move to host LRU - self.assertFalse(mamba_lru.in_list(node)) - if node.component_data[ComponentType.MAMBA].host_value is not None: - self.assertTrue(host_mamba_lru.in_list(node)) + for aux in aux_types: + self.assertFalse(tree.lru_lists[aux].in_list(node)) + if node.component_data[aux].host_value is not None: + self.assertTrue(tree.host_lru_lists[aux].in_list(node)) tree.sanity_check() + def _build_chain_pages(self, tree, allocator, req_to_token_pool, num_pages): + """Insert an incremental chain of single-page extensions. + + Returns the chain root-to-leaf. Length may differ from num_pages + when the radix tree merges or splits nodes. + """ + seq: list[int] = [] + for i in range(num_pages): + seq = seq + self._make_seq(1000 * (i + 1), 1) + self._insert(tree, allocator, req_to_token_pool, seq) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + chain: list = [] + cur = m.last_device_node + while cur is not tree.root_node: + chain.append(cur) + cur = cur.parent + chain.reverse() + return chain + + def test_hicache_swa_load_back_min_suffix(self): + """LOAD_BACK collects only the suffix nodes needed to cover sliding_window_size.""" + if not self.cfg.has_swa: + self.skipTest("requires SWA") + if self.cfg.has_mamba: + # Mamba's per-insert req allocation exhausts max_num_reqs on long chains. + self.skipTest("SWA-only path keeps the chain construction simple") + ps = self.cfg.page_size + sw = self.cfg.sliding_window_size + expected_pages = (sw + ps - 1) // ps + chain_pages = expected_pages + 2 + if chain_pages * ps > self.cfg.kv_size // 2: + self.skipTest("kv_size too small for the desired chain") + + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + chain = self._build_chain_pages(tree, allocator, req_to_token_pool, chain_pages) + if len(chain) <= expected_pages: + self.skipTest("chain collapsed below the suffix length being tested") + + self._simulate_backup_tree(tree) + + # Tombstone every chain node on the device side without going through + # the tree-wide eviction loop. This isolates build_hicache_transfers + # from LRU and cascade ordering. + for n in chain: + n.component_data[ComponentType.FULL].value = None + n.component_data[ComponentType.SWA].value = None + + leaf = chain[-1] + swa_comp = tree.components[ComponentType.SWA] + transfers = swa_comp.build_hicache_transfers(leaf, CacheTransferPhase.LOAD_BACK) + self.assertIsNotNone(transfers) + self.assertEqual(len(transfers), 1) + xfer = transfers[0] + self.assertEqual(xfer.name, PoolName.SWA) + self.assertEqual(len(xfer.nodes_to_load), expected_pages) + # 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) + self.assertEqual(xfer.nodes_to_load, chain[-expected_pages:]) + + def test_hicache_swa_host_independent_of_full(self): + """FULL host and SWA host are physically independent. + Freeing one component's host_value must not touch the other. + """ + if not self.cfg.has_swa: + self.skipTest("requires SWA") + + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seq = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, seq) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + node = m.last_device_node + + self._simulate_backup(tree, node) + tree.evict(EvictParams(num_tokens=len(seq))) + + cd_full = node.component_data[ComponentType.FULL] + cd_swa = node.component_data[ComponentType.SWA] + self.assertIsNotNone(cd_full.host_value) + self.assertIsNotNone(cd_swa.host_value) + self.assertIn(node, tree.evictable_host_leaves) + self.assertTrue(tree.host_lru_lists[ComponentType.SWA].in_list(node)) + + # Drop FULL host bookkeeping. SWA side must stay intact. + tree.evictable_host_leaves.discard(node) + cd_full.host_value = None + self.assertIsNotNone(cd_swa.host_value) + self.assertTrue(tree.host_lru_lists[ComponentType.SWA].in_list(node)) + self.assertNotIn(node, tree.evictable_host_leaves) + + # Drop SWA host bookkeeping. FULL side (already cleared) stays cleared. + tree.host_lru_lists[ComponentType.SWA].remove_node(node) + cd_swa.host_value = None + self.assertIsNone(cd_full.host_value) + self.assertIsNone(cd_swa.host_value) + self.assertFalse(tree.host_lru_lists[ComponentType.SWA].in_list(node)) + self.assertNotIn(node, tree.evictable_host_leaves) + + def _swa_finalize_setup(self): + """Build a SWA chain long enough to fill at least the window + plus one extra page, and host-back every node so we can flip + SWA tombstones at will.""" + ps = self.cfg.page_size + sw = self.cfg.sliding_window_size + window_pages = (sw + ps - 1) // ps + chain_pages = window_pages + 2 + if chain_pages * ps > self.cfg.kv_size // 2: + self.skipTest("kv_size too small for the desired chain") + + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + chain = self._build_chain_pages(tree, allocator, req_to_token_pool, chain_pages) + if len(chain) <= window_pages: + self.skipTest("chain collapsed below the window length") + self._simulate_backup_tree(tree) + return tree, allocator, req_to_token_pool, chain, window_pages + + def test_hicache_swa_finalize_match_result(self): + """finalize_match_result bumps host_hit_length to 1 iff some SWA node + within the trailing window is tombstoned (cd.value is None, + cd.host_value is not None). Out-of-window tombstones and chains fully + on device must leave host_hit_length untouched. + + Sentinel only — never the real SWA token count, since SWA load-back + does not grow req.prefix_indices and any non-zero value gets + subtracted from extend_input_len in schedule_policy. + """ + if not self.cfg.has_swa: + self.skipTest("requires SWA") + if self.cfg.has_mamba: + self.skipTest("SWA-only path keeps the chain construction simple") + + tree, _, _, chain, window_pages = self._swa_finalize_setup() + leaf = chain[-1] + swa_comp = tree.components[ComponentType.SWA] + + cases = [ + ("all_on_device", None, 0), + ("tombstone_in_window", chain[-window_pages], 1), + ("tombstone_outside_window", chain[-(window_pages + 1)], 0), + ] + for name, victim, expected in cases: + with self.subTest(name): + # Reset SWA state for each subcase. + for n in chain: + cd = n.component_data[ComponentType.SWA] + if cd.value is None and cd.host_value is not None: + cd.value = cd.host_value.clone() + if victim is not None: + victim.component_data[ComponentType.SWA].value = None + + result = MatchResult( + device_indices=torch.empty( + (0,), dtype=torch.int64, device=tree.device + ), + last_device_node=leaf, + last_host_node=leaf, + host_hit_length=0, + ) + result = swa_comp.finalize_match_result( + result=result, + params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))), + value_chunks=[], + best_value_len=0, + ) + self.assertEqual(result.host_hit_length, expected) + + def test_hicache_swa_commit_load_back_rebuilds_mapping(self): + """LOAD_BACK commit must: + (1) restore SWA cd.value via _restore_device_value (host LRU -> device LRU), + (2) rewrite full_to_swa_index_mapping[full_idx] = new_swa_idx for every + loaded chunk so subsequent SWA reads via translate_loc_from_full_to_swa + return the freshly allocated SWA slot.""" + if not self.cfg.has_swa: + self.skipTest("requires SWA") + if self.cfg.has_mamba: + self.skipTest("SWA-only path keeps the chain construction simple") + + tree, allocator, _, chain, window_pages = self._swa_finalize_setup() + + # Tombstone every SWA node in the trailing window. + loaded_nodes = chain[-window_pages:] + for n in loaded_nodes: + 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. + tree.lru_lists[ComponentType.SWA].remove_node(n) + tree.host_lru_lists[ComponentType.SWA].insert_mru(n) + + # Build the LOAD_BACK transfer the same way load_back() would. + swa_comp = tree.components[ComponentType.SWA] + transfers = swa_comp.build_hicache_transfers( + chain[-1], CacheTransferPhase.LOAD_BACK + ) + self.assertIsNotNone(transfers) + xfer = transfers[0] + self.assertEqual(xfer.nodes_to_load, loaded_nodes) + + # Allocate SWA device slots from the inner allocator (mirrors how + # _resolve_pool_transfers_allocation routes via device_alloc_fn -> + # swa_attn_allocator.alloc on the load-back path). + n_swa = int(xfer.host_indices.numel()) + new_swa = allocator.swa_attn_allocator.alloc(n_swa) + self.assertIsNotNone(new_swa) + xfer.device_indices = new_swa + + # Snapshot pre-commit state for invariants checks. + pre_evictable = tree.component_evictable_size_[ComponentType.SWA] + + swa_comp.commit_hicache_transfer( + chain[-1], CacheTransferPhase.LOAD_BACK, transfers=transfers + ) + + # (1) cd.value restored, host LRU -> device LRU swap done. + offset = 0 + for n in loaded_nodes: + cd = n.component_data[ComponentType.SWA] + self.assertIsNotNone(cd.value) + chunk_len = int(cd.value.numel()) + self.assertEqual( + cd.value.tolist(), + new_swa[offset : offset + chunk_len].tolist(), + ) + offset += chunk_len + self.assertTrue(tree.lru_lists[ComponentType.SWA].in_list(n)) + self.assertFalse(tree.host_lru_lists[ComponentType.SWA].in_list(n)) + self.assertEqual(offset, n_swa) + + # (2) full_to_swa_index_mapping rebuilt for every loaded chunk. + for n in loaded_nodes: + full_idx = n.component_data[ComponentType.FULL].value + swa_idx = n.component_data[ComponentType.SWA].value + translated = allocator.translate_loc_from_full_to_swa(full_idx) + self.assertEqual(translated.tolist(), swa_idx.tolist()) + + # Evictable size moved up by the restored token count. + self.assertEqual( + tree.component_evictable_size_[ComponentType.SWA] - pre_evictable, + n_swa, + ) + def test_hicache_mixed_backup_evict_insert(self): """Complex scenario: backup some, evict, insert new, verify invariants.""" if self._skip_unsupported_hicache_test():