Support swa HiCache for unified radix cache (#23391)

Co-authored-by: hzh0425 <hzh0425@apache.org>

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:

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:

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

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:

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:

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]

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(

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

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,

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: