[HiCache][HybridModel]: Support mamba state offloading & HybridCacheController (#20457)

Co-authored-by: pansicheng <sicheng.pan.chn@gmail.com>
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>

benchmark/hicache/bench_multiturn.py

@@ -529,6 +529,8 @@ class WorkloadGenerator:
                             f"({expected} clients), releasing {len(next_round_reqs)} "
                             f"requests for round {current_barrier_round + 1}"
                         )
+                        self._send_heartbeat(input_len=100, output_len=100)
+                        time.sleep(10)
                         for req in next_round_reqs:
                             self.ready_queue.append(req)
                         next_round_reqs = []
@@ -541,6 +543,15 @@ class WorkloadGenerator:
                 print(f"Error processing response for client {client_id}: {e}")
                 continue
 
+    def _send_heartbeat(self, input_len=100, output_len=20):
+        """Send a small heartbeat request to the server."""
+        heartbeat_input = [1] * input_len
+        payload = gen_payload(heartbeat_input, output_len, self.lora_path)
+        try:
+            requests.post(self.url, json=payload, timeout=30)
+        except Exception as e:
+            print(f"Heartbeat request failed: {e}")
+
     def run(self):
         request_thread = threading.Thread(target=self.request_sender, daemon=True)
         response_thread = threading.Thread(target=self.response_handler, daemon=True)

python/sglang/srt/disaggregation/decode.py

@@ -195,6 +195,9 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
         effective_mamba_size = (
             mamba_size if mamba_size is not None else size
         ) + pre_alloc_size
+        # TODO: Support PP
+        self.start_layer = 0
+        self.layer_transfer_counter = None
         self._init_mamba_pool(
             size=effective_mamba_size,
             mamba_spec_state_size=size + pre_alloc_size,

python/sglang/srt/managers/schedule_batch.py

@@ -641,7 +641,6 @@ class Req(ReqDllmMixin):
         self.extend_logprob_start_len = 0
         self.last_node: Any = None
         self.last_host_node: Any = None
-        self.last_host_backup_node: Any = None
         self.host_hit_length = 0
         # Tokens loaded from storage backend (L3) during prefetch for this request
         self.storage_hit_length = 0
@@ -915,14 +914,12 @@ class Req(ReqDllmMixin):
                 self.prefix_indices,
                 self.last_node,
                 self.last_host_node,
-                self.last_host_backup_node,
                 self.host_hit_length,
                 self.mamba_branching_seqlen,
             ) = (
                 match_result.device_indices,
                 match_result.last_device_node,
                 match_result.last_host_node,
-                match_result.last_host_backup_node,
                 match_result.host_hit_length,
                 match_result.mamba_branching_seqlen,
             )

python/sglang/srt/managers/schedule_policy.py

@@ -777,6 +777,7 @@ class PrefillAdder:
                     InitLoadBackParams(
                         last_host_node=req.last_host_node,
                         host_hit_length=req.host_hit_length,
+                        req=req,
                     )
                 )
                 req.prefix_indices = torch.cat([req.prefix_indices, new_indices])

python/sglang/srt/managers/scheduler.py

@@ -1818,11 +1818,7 @@ class Scheduler(
     def _prefetch_kvcache(self, req: Req):
         if self.enable_hicache_storage:
             req.init_next_round_input(self.tree_cache, cow_mamba=False)
-            last_host_node = (
+            last_host_node = req.last_host_node
-                req.last_host_backup_node
-                if req.last_host_backup_node is not None
-                else req.last_host_node
-            )
             if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
                 last_hash = last_host_node.get_last_hash_value()
                 matched_len = len(req.prefix_indices) + req.host_hit_length

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -129,8 +129,10 @@ class MatchResult(NamedTuple):
         last_host_node  :   The last TreeNode on the host that was matched.
                             Note that if HiCache is not enabled,
                             this **must** be the same as `last_device_node`.
-        last_host_backup_node: The deepest backuped node for prefetch from storage.
+        host_hit_length :   Length of the host cache hit. For pure-KV caches this is the
-        host_hit_length :   Length of the KV cache hit on the host, if applicable.
+                            number of evicted KV tokens on CPU. For hybrid Mamba models this
+                            is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
+                            host hit still triggers load-back without adding a separate field.
                             0 if HiCache is not enabled.
         mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
                                 page-aligned position that could've been cache hit if there
@@ -140,7 +142,6 @@ class MatchResult(NamedTuple):
     device_indices: torch.Tensor
     last_device_node: Any
     last_host_node: Any
-    last_host_backup_node: Any = None
     host_hit_length: int = 0
     mamba_branching_seqlen: Optional[int] = None
     cache_protected_len: Optional[int] = None

python/sglang/srt/mem_cache/hicache_storage.py

@@ -3,7 +3,8 @@ import logging
 import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from enum import Enum
+from typing import Any, List, Optional, Set
 
 import torch
 
@@ -65,6 +66,61 @@ class HiCacheStorageExtraInfo:
     extra_info: Optional[dict] = None
 
 
+class PoolName(str, Enum):
+    """Well-known pool names used as PoolTransfer/PoolEntry identifiers."""
+
+    KV = "kv"
+    MAMBA = "mamba"
+
+
+class PoolHitPolicy(str, Enum):
+    """Hit policy for batch_exists_v2 per-pool prefix matching.
+
+    ALL_PAGES      : every page in [0, kv_hit) must exist (default).
+    TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states).
+    """
+
+    ALL_PAGES = "all_pages"
+    TRAILING_PAGES = "trailing_pages"
+
+
+@dataclass
+class PoolTransfer:
+    """Unified per-pool transfer descriptor for batch v2 interface.
+
+    device<->host path : host_indices + device_indices
+    host<->storage path: host_indices + keys
+    """
+
+    name: PoolName
+    host_indices: Optional[torch.Tensor] = None
+    device_indices: Optional[torch.Tensor] = None
+    keys: Optional[List[str]] = None
+    hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
+
+
+@dataclass
+class PoolTransferResult:
+    """Tracks how many pages were successfully processed per pool."""
+
+    kv_hit_pages: int
+    extra_pool_hit_pages: dict[str, int]
+
+    @classmethod
+    def empty(cls) -> "PoolTransferResult":
+        return cls(0, {})
+
+    def update_kv_hit_pages(self, kv_hit_pages: int) -> None:
+        """Accumulate kv_hit_pages across batches (max = last successful batch)."""
+        self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages)
+
+    def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None:
+        """Record actual load/write success counts per extra pool."""
+        self.extra_pool_hit_pages.update(
+            {name: sum(rs) for name, rs in results.items()}
+        )
+
+
 class HiCacheStorage(ABC):
     """
     HiCacheStorage is a class that provides a generic key-value interface for storing and retrieving KV cache.
@@ -72,10 +128,69 @@ class HiCacheStorage(ABC):
     """
 
     # todo, the page size of storage backend does not have to be the same as the same as host memory pool
-
     def register_mem_pool_host(self, mem_pool_host: HostKVCache):
         self.mem_pool_host = mem_pool_host
 
+    def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name):
+        if not hasattr(self, "registered_pools"):
+            self.registered_pools = {}
+        self.registered_pools[host_pool_name] = host_pool
+
+    def batch_exists_v2(
+        self,
+        keys: List[str],
+        pool_transfers: Optional[List[PoolTransfer]] = None,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> PoolTransferResult:
+        """Check which cache pages exist in storage, respecting per-pool hit policies.
+
+        Longest-prefix semantics
+        Extra-pool hit policies (``PoolTransfer.hit_policy``)
+        ------------------------------------------------------
+        Each ``PoolTransfer`` in ``pool_transfers`` describes a secondary
+        cache pool (e.g. Mamba SSM states) that must be co-present with the
+        KV pages.  The final ``final_pages`` is the minimum across all pools,
+        so a missing auxiliary page shrinks the usable prefix.
+
+        - ``"all_pages"`` (default):  every page in [0, kv_hit) must exist
+          for this pool.  Used for pools that are required for every token
+          in the prefix (e.g. DeepSeek DSA pool).
+
+        - ``"trailing_pages"``:  only the *last* ``len(transfer.keys)`` pages
+          of the KV prefix need to exist.  Used for pools whose data covers
+          only the tail of a prefix (e.g. Mamba/SWA Pool).
+
+        Returns
+        -------
+        PoolTransferResult
+            ``kv_hit_pages`` = length of the usable KV prefix.
+            ``extra_pool_hit_pages`` maps each pool name to the number of pages
+            that were found.
+        """
+        raise NotImplementedError()
+
+    def batch_get_v2(
+        self,
+        transfers: List[PoolTransfer],
+        extra_info: Optional["HiCacheStorageExtraInfo"] = None,
+    ) -> dict[str, List[bool]]:
+        """Read data from storage into host memory for each PoolTransfer.
+
+        Returns a dict mapping pool name to a per-entry success list.
+        """
+        raise NotImplementedError()
+
+    def batch_set_v2(
+        self,
+        transfers: List[PoolTransfer],
+        extra_info: Optional["HiCacheStorageExtraInfo"] = None,
+    ) -> dict[str, List[bool]]:
+        """Write data from host memory to storage for each PoolTransfer.
+
+        Returns a dict mapping pool name to a per-entry success list.
+        """
+        raise NotImplementedError()
+
     def batch_get_v1(
         self,
         keys: List[str],
@@ -203,7 +318,6 @@ class HiCacheFile(HiCacheStorage):
             self.config_suffix = f"_{model_name}"
         else:
             self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}"
-
         if not os.path.exists(self.file_path) and tp_rank == 0:
             os.makedirs(self.file_path)
             logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
@@ -211,6 +325,18 @@ class HiCacheFile(HiCacheStorage):
     def _get_suffixed_key(self, key: str) -> str:
         return key + self.config_suffix
 
+    def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str:
+        if component_name is None or component_name in ("__default__", PoolName.KV):
+            return self._get_suffixed_key(key)
+        return self._get_suffixed_key(f"{key}.{component_name}")
+
+    def _get_component_path(
+        self, key: str, component_name: Optional[str] = None
+    ) -> str:
+        return os.path.join(
+            self.file_path, f"{self._get_component_key(key, component_name)}.bin"
+        )
+
     def get(
         self,
         key: str,
@@ -280,6 +406,133 @@ class HiCacheFile(HiCacheStorage):
         tensor_path = os.path.join(self.file_path, f"{key}.bin")
         return os.path.exists(tensor_path)
 
+    def _collect_existing_component_keys(
+        self,
+        keys: List[str],
+        pool_transfers: Optional[List[PoolTransfer]] = None,
+    ) -> Set[str]:
+        target_files = {f"{self._get_component_key(key)}.bin" for key in keys}
+        for transfer in pool_transfers or []:
+            for key in keys:
+                target_files.add(f"{self._get_component_key(key, transfer.name)}.bin")
+
+        existing_files = set()
+        with os.scandir(self.file_path) as entries:
+            for entry in entries:
+                if entry.is_file() and entry.name in target_files:
+                    existing_files.add(entry.name)
+        return existing_files
+
+    def batch_exists_v2(
+        self,
+        keys: List[str],
+        pool_transfers: Optional[List[PoolTransfer]] = None,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> PoolTransferResult:
+        existing_files = self._collect_existing_component_keys(keys, pool_transfers)
+
+        def has_component(page_idx: int, name: str) -> bool:
+            return (
+                f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files
+            )
+
+        # Longest contiguous KV prefix present in storage.
+        kv_pages = next(
+            (
+                i
+                for i in range(len(keys))
+                if f"{self._get_component_key(keys[i])}.bin" not in existing_files
+            ),
+            len(keys),
+        )
+
+        hit_count: dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {}
+        final_pages = kv_pages
+
+        for transfer in pool_transfers or []:
+            if final_pages == 0:
+                break
+            name = transfer.name
+            if transfer.hit_policy == PoolHitPolicy.ALL_PAGES:
+                boundary = next(
+                    (i for i in range(kv_pages) if not has_component(i, name)), kv_pages
+                )
+            else:  # trailing_pages
+                trailing = max(1, len(transfer.keys) if transfer.keys else 1)
+                boundary = 0
+                for prefix_len in range(kv_pages, 0, -1):
+                    if all(
+                        has_component(i, name)
+                        for i in range(max(0, prefix_len - trailing), prefix_len)
+                    ):
+                        boundary = prefix_len
+                        break
+            if boundary:
+                hit_count[name] = boundary
+            final_pages = min(final_pages, boundary)
+
+        return PoolTransferResult(final_pages, hit_count)
+
+    def _log_key(self, pool_name: str, key: str) -> str:
+        return key if pool_name == PoolName.KV else f"{key}.{pool_name}"
+
+    def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool:
+        """Read one page from storage into host_pool at page_offset."""
+        storage_key = self._log_key(pool_name, key)
+        data_page = self.get(storage_key, host_pool.get_dummy_flat_data_page())
+        if data_page is None:
+            return False
+        host_pool.set_from_flat_data_page(page_offset, data_page)
+        return True
+
+    def _write_page(
+        self, pool_name: str, key: str, host_pool, page_offset: int
+    ) -> bool:
+        """Write one page from host_pool at page_offset to storage as raw bytes."""
+        storage_key = self._log_key(pool_name, key)
+        data_page = host_pool.get_data_page(page_offset, flat=True)
+        return self.set(storage_key, data_page)
+
+    def _batch_io_v2(self, transfers: List[PoolTransfer], op_fn):
+        results: dict[str, List[bool]] = {}
+        for transfer in transfers:
+            host_pool = self.registered_pools[transfer.name]
+            keys = transfer.keys or []
+            page_size = getattr(host_pool, "page_size", 1) or 1
+            expected = len(keys) * page_size
+            host_indices = transfer.host_indices
+
+            if host_indices is None or host_indices.numel() != expected:
+                logger.error(
+                    "%s indices length mismatch for %s: expected %s, got %s",
+                    op_fn.__name__,
+                    transfer.name,
+                    expected,
+                    host_indices.numel() if host_indices is not None else 0,
+                )
+                results[transfer.name] = [False] * len(keys)
+                continue
+
+            results[transfer.name] = [
+                op_fn(transfer.name, key, host_pool, host_indices[i * page_size].item())
+                for i, key in enumerate(keys)
+            ]
+        return results
+
+    def batch_get_v2(
+        self,
+        transfers: List[PoolTransfer],
+        extra_info: Optional["HiCacheStorageExtraInfo"] = None,
+    ) -> dict[str, List[bool]]:
+        return self._batch_io_v2(transfers, self._read_page)
+
+    def batch_set_v2(
+        self,
+        transfers: List[PoolTransfer],
+        extra_info: Optional["HiCacheStorageExtraInfo"] = None,
+    ) -> dict[str, List[bool]]:
+        return self._batch_io_v2(transfers, self._write_page)
+
     def clear(self) -> bool:
         try:
             for filename in os.listdir(self.file_path):

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -0,0 +1,500 @@
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import torch
+
+from sglang.srt.managers.cache_controller import CacheOperation as BaseCacheOperation
+from sglang.srt.managers.cache_controller import (
+    HiCacheAck,
+)
+from sglang.srt.managers.cache_controller import (
+    HiCacheController as BaseHiCacheController,
+)
+from sglang.srt.managers.cache_controller import (
+    LayerDoneCounter,
+)
+from sglang.srt.managers.cache_controller import (
+    StorageOperation as BaseStorageOperation,
+)
+from sglang.srt.mem_cache.hicache_storage import (
+    HiCacheStorageExtraInfo,
+    PoolHitPolicy,
+    PoolTransfer,
+    PoolTransferResult,
+)
+from sglang.srt.mem_cache.memory_pool_host import PoolEntry
+from sglang.srt.utils import get_device_module
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+
+logger = logging.getLogger(__name__)
+device_module = get_device_module()
+
+
+class CacheOperation(BaseCacheOperation):
+    def __init__(
+        self,
+        host_indices: torch.Tensor,
+        device_indices: torch.Tensor,
+        node_id: int,
+        priority: Optional[int] = None,
+        pool_transfers: Optional[list[PoolTransfer]] = None,
+    ):
+        super().__init__(host_indices, device_indices, node_id, priority)
+        self.pool_transfers = pool_transfers
+
+    @staticmethod
+    def merge_pool_transfers(
+        ops: List["CacheOperation"],
+    ) -> Optional[list[PoolTransfer]]:
+        grouped: dict[str, list[PoolTransfer]] = {}
+        for op in ops:
+            for t in op.pool_transfers or []:
+                grouped.setdefault(t.name, []).append(t)
+        if not grouped:
+            return None
+
+        def cat_or_none(tensors):
+            parts = [x for x in tensors if x is not None]
+            return torch.cat(parts) if parts else None
+
+        return [
+            PoolTransfer(
+                name=name,
+                host_indices=cat_or_none(t.host_indices for t in ts),
+                device_indices=cat_or_none(t.device_indices for t in ts),
+                keys=[k for t in ts if t.keys for k in t.keys] or None,
+            )
+            for name, ts in grouped.items()
+        ]
+
+    @staticmethod
+    def merge_ops(ops: List["CacheOperation"]) -> "CacheOperation":
+        if len(ops) == 1:
+            return ops[0]
+        host_indices = torch.cat([op.host_indices for op in ops])
+        device_indices = torch.cat([op.device_indices for op in ops])
+        node_ids = []
+        priority = min(op.priority for op in ops)
+        for op in ops:
+            node_ids.extend(op.node_ids)
+        merged = CacheOperation(
+            host_indices,
+            device_indices,
+            -1,
+            priority,
+            pool_transfers=CacheOperation.merge_pool_transfers(ops),
+        )
+        merged.node_ids = node_ids
+        return merged
+
+
+class StorageOperation(BaseStorageOperation):
+    def __init__(
+        self,
+        host_indices: torch.Tensor,
+        token_ids: List[int],
+        last_hash: Optional[str] = None,
+        hash_value: Optional[List[str]] = None,
+        prefix_keys: Optional[List[str]] = None,
+        pool_transfers: Optional[list[PoolTransfer]] = None,
+    ):
+        super().__init__(host_indices, token_ids, last_hash, hash_value, prefix_keys)
+        self.pool_transfers = pool_transfers
+        self.pool_storage_result = PoolTransferResult.empty()
+
+
+class PrefetchOperation(StorageOperation):
+    def __init__(
+        self,
+        request_id: str,
+        host_indices: torch.Tensor,
+        token_ids: List[int],
+        last_hash: Optional[str] = None,
+        prefix_keys: Optional[List[str]] = None,
+        pool_transfers: Optional[list[PoolTransfer]] = None,
+    ):
+        self.request_id = request_id
+        self._lock = threading.Lock()
+        self._terminated_flag = False
+        self.start_time = time.monotonic()
+        super().__init__(
+            host_indices,
+            token_ids,
+            last_hash,
+            prefix_keys=prefix_keys,
+            pool_transfers=pool_transfers,
+        )
+
+    def increment(self, num_tokens: int):
+        with self._lock:
+            if self._terminated_flag:
+                return False
+            self.completed_tokens += num_tokens
+            return True
+
+    def mark_terminate(self):
+        with self._lock:
+            self._terminated_flag = True
+
+    def is_terminated(self) -> bool:
+        return self._terminated_flag
+
+
+class HybridCacheController(BaseHiCacheController):
+    def __init__(
+        self,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+        mem_pool_host: Any,
+        page_size: int,
+        tp_group: torch.distributed.ProcessGroup,
+        load_cache_event: threading.Event,
+        write_policy: str = "write_through_selective",
+        io_backend: str = "",
+        storage_backend: Optional[str] = None,
+        prefetch_threshold: int = 256,
+        model_name: Optional[str] = None,
+        storage_backend_extra_config: Optional[dict] = None,
+        pp_rank: int = 0,
+        pp_size: int = 1,
+        transfer_layer_num: Optional[int] = None,
+    ):
+        startup_storage_backend = storage_backend
+        super().__init__(
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+            mem_pool_host=mem_pool_host,
+            page_size=page_size,
+            tp_group=tp_group,
+            load_cache_event=load_cache_event,
+            write_policy=write_policy,
+            io_backend=io_backend,
+            storage_backend=None,
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
+            pp_rank=pp_rank,
+            pp_size=pp_size,
+        )
+        # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)),
+        # not just the full attention layers reported by full_kv_pool.
+        if transfer_layer_num is not None and transfer_layer_num != self.layer_num:
+            self.layer_num = transfer_layer_num
+            self.layer_done_counter = LayerDoneCounter(self.layer_num)
+
+        if startup_storage_backend is not None:
+            self.attach_storage_backend(
+                storage_backend=startup_storage_backend,
+                prefetch_threshold=prefetch_threshold,
+                model_name=model_name,
+                storage_backend_extra_config=storage_backend_extra_config,
+                host_pools=getattr(mem_pool_host, "entries", None),
+            )
+
+    def attach_storage_backend(
+        self,
+        storage_backend: str,
+        prefetch_threshold: int = 256,
+        model_name: Optional[str] = None,
+        storage_backend_extra_config: Optional[dict] = None,
+        host_pools: Optional[list[PoolEntry]] = None,
+    ):
+        super().attach_storage_backend(
+            storage_backend=storage_backend,
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
+        )
+
+        for entry in host_pools or []:
+            self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name)
+
+    def reset(self):
+        super().reset()
+        if self.enable_storage:
+            self.host_mem_release_queue.queue.clear()
+            self.prefetch_tokens_occupied = 0
+
+    def write(
+        self,
+        device_indices: torch.Tensor,
+        priority: Optional[int] = None,
+        node_id: int = -1,
+        extra_pools: Optional[list[PoolTransfer]] = None,
+    ) -> Optional[torch.Tensor]:
+        host_indices = self.mem_pool_host.alloc(len(device_indices))
+        if host_indices is None:
+            return None
+        pool_transfers = self._resolve_pool_transfers_allocation(
+            extra_pools, alloc_host=True
+        )
+        if pool_transfers is None and extra_pools:
+            self.mem_pool_host.free(host_indices)
+            return None
+
+        self.write_queue.append(
+            CacheOperation(
+                host_indices,
+                device_indices,
+                node_id,
+                priority,
+                pool_transfers=pool_transfers or None,
+            )
+        )
+        self.start_writing()
+        return host_indices
+
+    def start_writing(self) -> None:
+        if not self.write_queue:
+            return
+        op = CacheOperation.merge_ops(self.write_queue)
+        host_indices, device_indices = self.move_indices(op)
+        self.write_queue.clear()
+        start_event = device_module.Event()
+        finish_event = device_module.Event()
+        start_event.record()
+        with device_module.stream(self.write_stream):
+            start_event.wait(self.write_stream)
+            self.mem_pool_host.backup_from_device_all_layer(
+                self.mem_pool_device,
+                host_indices,
+                device_indices,
+                self.io_backend,
+                pool_transfers=op.pool_transfers,
+            )
+            finish_event.record()
+            if host_indices.is_cuda:
+                host_indices.record_stream(self.write_stream)
+            if device_indices.is_cuda:
+                device_indices.record_stream(self.write_stream)
+        self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids))
+
+    def load(
+        self,
+        host_indices: torch.Tensor,
+        priority: Optional[int] = None,
+        node_id: int = -1,
+        extra_pools: Optional[list[PoolTransfer]] = None,
+    ) -> Optional[torch.Tensor]:
+        need_load_kv = host_indices.numel() > 0
+        if need_load_kv:
+            device_indices = self.mem_pool_device_allocator.alloc(len(host_indices))
+            if device_indices is None:
+                return None
+        else:
+            device_indices = torch.empty((0,), dtype=torch.int64, device=self.device)
+
+        pool_transfers = self._resolve_pool_transfers_allocation(
+            extra_pools, alloc_host=False
+        )
+        if pool_transfers is None and extra_pools:
+            if need_load_kv:
+                self.mem_pool_device_allocator.free(device_indices)
+            return None
+
+        self.load_queue.append(
+            CacheOperation(
+                host_indices,
+                device_indices,
+                node_id,
+                priority,
+                pool_transfers=pool_transfers or None,
+            )
+        )
+        return device_indices
+
+    def start_loading(self) -> int:
+        if not self.load_queue:
+            return -1
+        producer_id = self.layer_done_counter.update_producer()
+        op = CacheOperation.merge_ops(self.load_queue)
+        host_indices, device_indices = self.move_indices(op)
+        self.load_queue.clear()
+        producer_event = self.layer_done_counter.events[producer_id]
+        producer_event.start_event.record()
+        with device_module.stream(self.load_stream):
+            producer_event.start_event.wait(self.load_stream)
+            for i in range(self.layer_num):
+                self.mem_pool_host.load_to_device_per_layer(
+                    self.mem_pool_device,
+                    host_indices,
+                    device_indices,
+                    i,
+                    self.io_backend,
+                    pool_transfers=op.pool_transfers,
+                )
+                producer_event.complete(i)
+            if host_indices.is_cuda:
+                host_indices.record_stream(self.load_stream)
+            if device_indices.is_cuda:
+                device_indices.record_stream(self.load_stream)
+        self.ack_load_queue.append(
+            HiCacheAck(
+                producer_event.start_event,
+                producer_event.finish_event,
+                op.node_ids,
+            )
+        )
+        return producer_id
+
+    def prefetch(
+        self,
+        request_id: str,
+        host_indices: torch.Tensor,
+        new_input_tokens: List[int],
+        last_hash: Optional[str] = None,
+        prefix_keys: Optional[List[str]] = None,
+        extra_pools: Optional[list[PoolTransfer]] = None,
+    ) -> PrefetchOperation:
+        operation = PrefetchOperation(
+            request_id,
+            host_indices,
+            new_input_tokens,
+            last_hash,
+            prefix_keys=prefix_keys,
+            pool_transfers=extra_pools,
+        )
+        self.prefetch_queue.put(operation)
+        return operation
+
+    def write_storage(
+        self,
+        host_indices: torch.Tensor,
+        token_ids: List[int],
+        hash_value: Optional[List[str]] = None,
+        prefix_keys: Optional[List[str]] = None,
+        extra_pools: Optional[list[PoolTransfer]] = None,
+    ) -> int:
+        operation = StorageOperation(
+            host_indices,
+            token_ids,
+            hash_value=hash_value,
+            prefix_keys=prefix_keys,
+            pool_transfers=extra_pools,
+        )
+        self.backup_queue.put(operation)
+        return operation.id
+
+    def _storage_hit_query(self, operation) -> tuple[list[str], int]:
+        last_hash = operation.last_hash
+        hash_value = []
+        for start in range(0, len(operation.token_ids), self.page_size):
+            last_hash = self.get_hash_str(
+                operation.token_ids[start : start + self.page_size], last_hash
+            )
+            hash_value.append(last_hash)
+
+        extra_info = HiCacheStorageExtraInfo(
+            prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None
+        )
+        if operation.pool_transfers:
+            hit_result = self.storage_backend.batch_exists_v2(
+                hash_value, operation.pool_transfers, extra_info
+            )
+        else:
+            kv_hit_count = self.storage_backend.batch_exists(hash_value, extra_info)
+            hit_result = PoolTransferResult(
+                kv_hit_pages=kv_hit_count, extra_pool_hit_pages={}
+            )
+
+        kv_hit_pages = hit_result.kv_hit_pages
+        operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages)
+
+        if kv_hit_pages > 0 and operation.pool_transfers:
+            self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages)
+
+        return (
+            hash_value[:kv_hit_pages],
+            kv_hit_pages * self.page_size,
+        )
+
+    def _page_transfer(self, operation):
+        # Transfer extra pools
+        if operation.pool_transfers and not operation.is_terminated():
+            results = self.storage_backend.batch_get_v2(operation.pool_transfers)
+            operation.pool_storage_result.update_extra_pool_hit_pages(results)
+
+        # Transfer kv pools
+        super()._page_transfer(operation)
+
+    def _page_backup(self, operation):
+        # Backup extra pools
+        if operation.pool_transfers:
+            results = self.storage_backend.batch_set_v2(operation.pool_transfers)
+            operation.pool_storage_result.update_extra_pool_hit_pages(results)
+
+        # Backup kv pools
+        super()._page_backup(operation)
+
+    def _sync_trailing_keys(
+        self,
+        pool_transfers: list[PoolTransfer],
+        all_hashes: list[str],
+        kv_hit_pages: int,
+    ) -> None:
+        """Re-align trailing-page sidecar keys after KV hit truncation.
+
+        When the storage hit is shorter than the original target prefix, each
+        pool transfer's keys must be updated to the last N hashes of the actual
+        hit range instead of the last N hashes of the original target range.
+        For mamba (N=1) this is just the last hit page hash; for SWA (N>1) it
+        is a sliding window of the last N hit pages.
+        """
+        for transfer in pool_transfers:
+            if transfer.hit_policy != PoolHitPolicy.TRAILING_PAGES:
+                continue
+            trailing_n = len(transfer.keys) if transfer.keys else 1
+            transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages]
+
+    def _resolve_pool_transfers_allocation(
+        self,
+        extra_pools: Optional[list[PoolTransfer]],
+        alloc_host: bool,
+    ) -> Optional[list[PoolTransfer]]:
+        """Auto-alloc host or device indices for PoolTransfers where they are None."""
+        if not extra_pools:
+            return None
+        newly_allocated: list[tuple[PoolTransfer, Any, torch.Tensor]] = []
+        for pool in extra_pools:
+            entry = self.mem_pool_host.entry_map.get(pool.name)
+            if entry is None:
+                continue
+            if alloc_host:
+                if pool.host_indices is not None or pool.device_indices is None:
+                    continue
+                entry_pool, evict_fn, size = (
+                    entry.host_pool,
+                    entry.host_evict_fn,
+                    len(pool.device_indices),
+                )
+            else:
+                if pool.device_indices is not None or pool.host_indices is None:
+                    continue
+                entry_pool, evict_fn, size = (
+                    entry.device_pool,
+                    entry.device_evict_fn,
+                    len(pool.host_indices),
+                )
+            indices = entry_pool.alloc(size)
+            if indices is None and evict_fn:
+                evict_fn(size)
+                indices = entry_pool.alloc(size)
+            if indices is None:
+                # Roll back all previous allocations using each pool's own entry_pool.
+                for prev_pool, prev_entry_pool, prev_indices in newly_allocated:
+                    prev_entry_pool.free(prev_indices)
+                    if alloc_host:
+                        prev_pool.host_indices = None
+                    else:
+                        prev_pool.device_indices = None
+                return None
+            if alloc_host:
+                pool.host_indices = indices
+            else:
+                pool.device_indices = indices
+            newly_allocated.append((pool, entry_pool, indices))
+        return extra_pools

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -74,6 +74,7 @@ class TreeNode:
         self.key: RadixKey = None
         self.value: Optional[torch.Tensor] = None
         self.mamba_value: Optional[torch.Tensor] = None
+        self.mamba_host_value: Optional[torch.Tensor] = None
         # invariant: for any node, if mamba_lock_ref is locked, full_lock_ref must be locked;
         # if full_lock_ref is locked, mamba_lock_ref doesn't need to be locked. So,
         # full_lock_ref is always >= mamba_lock_ref.
@@ -98,6 +99,8 @@ class TreeNode:
         self.next = None
         self.mamba_prev = None
         self.mamba_next = None
+        self.host_mamba_prev = None
+        self.host_mamba_next = None
 
         self.id = TreeNode.counter if id is None else id
         TreeNode.counter += 1
@@ -106,10 +109,18 @@ class TreeNode:
     def evicted(self):
         return self.value is None
 
+    @property
+    def mamba_evicted(self):
+        return self.mamba_value is None
+
     @property
     def backuped(self):
         return self.host_value is not None
 
+    @property
+    def mamba_backuped(self):
+        return self.mamba_host_value is not None
+
     def protect_host(self):
         """Protect the host value from eviction."""
         self.host_ref_counter += 1

python/sglang/srt/mem_cache/memory_pool.py

@@ -475,6 +475,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
         self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1
         self.enable_mamba_extra_buffer = enable_mamba_extra_buffer
         self.enable_memory_saver = enable_memory_saver
+        # TODO: Support PP
+        self.start_layer = 0
+        self.layer_transfer_counter = None
         self._init_mamba_pool(
             size=mamba_size,
             mamba_spec_state_size=mamba_spec_state_size,
@@ -516,6 +519,11 @@ class HybridReqToTokenPool(ReqToTokenPool):
                 )
             )
 
+    def register_layer_transfer_counter(
+        self, layer_transfer_counter: "LayerDoneCounter"
+    ):
+        self.layer_transfer_counter = layer_transfer_counter
+
     # For chunk prefill req, we do not need to allocate mamba cache,
     # We could use allocated mamba cache instead.
     def alloc(self, reqs: List["Req"]) -> Optional[List[int]]:
@@ -570,6 +578,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
 
     def mamba2_layer_cache(self, layer_id: int):
         assert layer_id in self.mamba_map
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
         return self.mamba_pool.mamba2_layer_cache(self.mamba_map[layer_id])
 
     def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState:
@@ -1238,8 +1248,8 @@ class HybridLinearKVPool(KVCache):
         self.device = device
         self.full_layer_nums = len(full_attention_layer_ids)
         self.page_size = page_size
-        # TODO support pp?
+        self.start_layer = 0  # TODO: Support PP
-        self.start_layer = 0
+        self.layer_transfer_counter = None
         self.head_num = head_num
         self.head_dim = head_dim
         self.mamba_pool = mamba_pool
@@ -1323,15 +1333,30 @@ class HybridLinearKVPool(KVCache):
             )
         return self.full_attention_layer_id_mapping[layer_id]
 
+    def register_layer_transfer_counter(
+        self, layer_transfer_counter: "LayerDoneCounter"
+    ):
+        self.layer_transfer_counter = layer_transfer_counter
+        # The layer-wise wait logic is executed at the Hybrid LinearPool level;
+        # no additional wait is needed in the full_kv_pool
+        self.full_kv_pool.register_layer_transfer_counter(None)
+
+    def _wait_for_layer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+
     def get_key_buffer(self, layer_id: int):
+        self._wait_for_layer(layer_id)
         layer_id = self._transfer_full_attention_id(layer_id)
         return self.full_kv_pool.get_key_buffer(layer_id)
 
     def get_value_buffer(self, layer_id: int):
+        self._wait_for_layer(layer_id)
         layer_id = self._transfer_full_attention_id(layer_id)
         return self.full_kv_pool.get_value_buffer(layer_id)
 
     def get_kv_buffer(self, layer_id: int):
+        self._wait_for_layer(layer_id)
         layer_id = self._transfer_full_attention_id(layer_id)
         return self.full_kv_pool.get_kv_buffer(layer_id)
 

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -1,10 +1,17 @@
+from __future__ import annotations
+
 import abc
 import logging
 import threading
 from collections import defaultdict
+from dataclasses import dataclass
 from functools import wraps
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.hicache_storage import PoolName
+
+import numpy as np
 import psutil
 import torch
 
@@ -19,6 +26,7 @@ from sglang.jit_kernel.hicache import (
 )
 from sglang.srt.mem_cache.memory_pool import (
     KVCache,
+    MambaPool,
     MHATokenToKVPool,
     MLATokenToKVPool,
     NSATokenToKVPool,
@@ -1074,6 +1082,526 @@ class MLATokenToKVPoolHost(HostKVCache):
         return ptr_list, element_size_list
 
 
+class MambaPoolHost(HostKVCache):
+
+    def __init__(
+        self,
+        device_pool: MambaPool,
+        host_to_device_ratio: float,
+        host_size: int,
+        pin_memory: bool = True,
+        device: str = "cpu",
+        allocator_type: str = "default",
+        layout: str = "layer_first",
+    ):
+        self.device_pool = device_pool
+        self.page_size = 1
+        assert layout in [
+            "page_first",
+            "page_first_direct",
+            "layer_first",
+        ], "Unsupported layout: {layout}"
+
+        self.layout = layout
+        self.pin_memory = pin_memory
+        self.device = device
+        self.allocator = get_allocator_from_storage(allocator_type)
+        self.num_mamba_layers = device_pool.num_mamba_layers
+
+        self.conv_state_shapes = [
+            conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv
+        ]
+        self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:]
+        self.conv_dtype = device_pool.mamba_cache.conv[0].dtype
+        self.temporal_dtype = device_pool.mamba_cache.temporal.dtype
+        self.dtype = self.conv_dtype
+        self.size_per_token = self.get_size_per_token()
+
+        if host_size > 0:
+            self.size = int(host_size * 1e9 // self.size_per_token)
+        else:
+            self.size = int(device_pool.size * host_to_device_ratio)
+
+        self.page_num = self.size // self.page_size + 1
+        self.size = self.page_num * self.page_size
+
+        assert (
+            self.size > device_pool.size
+        ), "The host memory should be larger than the device memory with the current protocol"
+
+        host_mem = psutil.virtual_memory()
+        requested_bytes = self.size * self.size_per_token
+        ten_gb = 10 * (1024**3)
+        available_bytes = host_mem.available - ten_gb
+        if requested_bytes > available_bytes:
+            raise ValueError(
+                f"Not enough host memory available. Requesting "
+                f"{requested_bytes / 1e9:.2f} GB but only have "
+                f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
+                f"size of the hierarchical cache."
+            )
+        logger.info(
+            "Allocating %.2f GB host memory for hierarchical Mamba cache (layout=%s).",
+            requested_bytes / 1e9,
+            self.layout,
+        )
+
+        self.init_kv_buffer()
+        self.lock = threading.RLock()
+        self.clear()
+
+    def init_kv_buffer(self):
+        alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
+
+        if self.layout in ["page_first", "page_first_direct"]:
+            # page-first: (page_num, num_layers, 1, *shape) — per-page data is contiguous
+            temporal_dims = (
+                self.size,
+                self.num_mamba_layers,
+                1,
+            ) + self.temporal_state_shape
+            self.temporal_buffer = alloc_func(
+                temporal_dims,
+                dtype=self.temporal_dtype,
+                device=self.device,
+                pin_memory=self.pin_memory,
+                allocator=self.allocator,
+            )
+            self.conv_buffer = []
+            for conv_shape in self.conv_state_shapes:
+                conv_dims = (self.size, self.num_mamba_layers, 1) + conv_shape
+                self.conv_buffer.append(
+                    alloc_func(
+                        conv_dims,
+                        dtype=self.conv_dtype,
+                        device=self.device,
+                        pin_memory=self.pin_memory,
+                        allocator=self.allocator,
+                    )
+                )
+        else:
+            # layer-first: (num_layers, size, *shape)
+            temporal_dims = (
+                self.num_mamba_layers,
+                self.size,
+            ) + self.temporal_state_shape
+            self.temporal_buffer = alloc_func(
+                temporal_dims,
+                dtype=self.temporal_dtype,
+                device=self.device,
+                pin_memory=self.pin_memory,
+                allocator=self.allocator,
+            )
+            self.conv_buffer = []
+            for conv_shape in self.conv_state_shapes:
+                conv_dims = (self.num_mamba_layers, self.size) + conv_shape
+                self.conv_buffer.append(
+                    alloc_func(
+                        conv_dims,
+                        dtype=self.conv_dtype,
+                        device=self.device,
+                        pin_memory=self.pin_memory,
+                        allocator=self.allocator,
+                    )
+                )
+
+    def _iter_page_tensors(self, index: int):
+        if self.layout in ["page_first", "page_first_direct"]:
+            yield self.temporal_buffer[index]
+            for conv_buf in self.conv_buffer:
+                yield conv_buf[index]
+        else:
+            yield self.temporal_buffer[:, index : index + self.page_size]
+            for conv_buf in self.conv_buffer:
+                yield conv_buf[:, index : index + self.page_size]
+
+    @staticmethod
+    def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.contiguous().view(torch.uint8).reshape(-1)
+
+    @synchronized
+    def clear(self):
+        self.mem_state = torch.zeros(
+            (self.size,), dtype=torch.uint8, device=self.device
+        )
+        self.free_slots = torch.arange(self.size, dtype=torch.int64)
+
+    def available_size(self):
+        return len(self.free_slots)
+
+    @synchronized
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
+        assert (
+            need_size % self.page_size == 0
+        ), "The requested size should be a multiple of the page size."
+        if need_size > self.available_size():
+            return None
+        select_index = self.free_slots[:need_size]
+        self.free_slots = self.free_slots[need_size:]
+        return select_index
+
+    @synchronized
+    def free(self, indices: torch.Tensor) -> int:
+        self.free_slots = torch.cat([self.free_slots, indices])
+        return len(indices)
+
+    def get_size_per_token(self):
+        conv_total_size = 0
+        for conv_shape in self.conv_state_shapes:
+            conv_total_size += int(np.prod(conv_shape)) * self.conv_dtype.itemsize
+        temporal_size = (
+            int(np.prod(self.temporal_state_shape)) * self.temporal_dtype.itemsize
+        )
+        return (conv_total_size + temporal_size) * self.num_mamba_layers
+
+    def get_ksize_per_token(self):
+        return self.get_size_per_token()
+
+    @staticmethod
+    def _item_size_per_index(tensor: torch.Tensor) -> int:
+        if tensor.shape[0] == 0:
+            return 0
+        return int(tensor[0].numel() * tensor.element_size())
+
+    @staticmethod
+    def _copy_tensor(
+        src: torch.Tensor,
+        dst: torch.Tensor,
+        src_indices: torch.Tensor,
+        dst_indices: torch.Tensor,
+        io_backend: str,
+    ) -> None:
+        if src_indices.numel() == 0:
+            return
+        if io_backend == "kernel":
+            # TODO: Rename the interface for clarity.
+            # Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state.
+            # This has nothing to do with MLA; it's only reused because this interface happens to transfer a single Pool.
+            transfer_kv_per_layer_mla(
+                src=src,
+                dst=dst,
+                src_indices=src_indices,
+                dst_indices=dst_indices,
+                item_size=MambaPoolHost._item_size_per_index(src),
+            )
+        elif io_backend == "direct":
+            transfer_kv_direct(
+                src_layers=[src],
+                dst_layers=[dst],
+                src_indices=src_indices,
+                dst_indices=dst_indices,
+                page_size=1,
+            )
+        else:
+            raise ValueError(f"Unsupported io_backend: {io_backend}")
+
+    @staticmethod
+    def _copy_tensor_pf_lf(
+        src: torch.Tensor,
+        dst: torch.Tensor,
+        src_indices: torch.Tensor,
+        dst_indices: torch.Tensor,
+        layer_id: int,
+        num_layers: int,
+        io_backend: str,
+    ) -> None:
+        if src_indices.numel() == 0:
+            return
+        if io_backend == "kernel":
+            item_size = MambaPoolHost._item_size_per_index(dst)
+            transfer_kv_per_layer_mla_pf_lf(
+                src=src,
+                dst=dst,
+                src_indices=src_indices,
+                dst_indices=dst_indices,
+                layer_id=layer_id,
+                item_size=item_size,
+                src_layout_dim=item_size * num_layers,
+            )
+        elif io_backend == "direct":
+            transfer_kv_per_layer_direct_pf_lf(
+                src_ptrs=[src],
+                dst_ptrs=[dst],
+                src_indices=src_indices,
+                dst_indices=dst_indices,
+                layer_id=layer_id,
+                page_size=1,
+            )
+        else:
+            raise ValueError(f"Unsupported io_backend: {io_backend}")
+
+    @staticmethod
+    def _copy_tensor_all_layers_lf_pf(
+        src_layers: torch.Tensor,
+        dst: torch.Tensor,
+        src_indices: torch.Tensor,
+        dst_indices: torch.Tensor,
+        num_layers: int,
+        device: str,
+        io_backend: str,
+    ) -> None:
+        if src_indices.numel() == 0:
+            return
+        if io_backend == "kernel":
+            item_size = MambaPoolHost._item_size_per_index(src_layers[0])
+            src_ptrs = torch.tensor(
+                [src_layers[i].data_ptr() for i in range(num_layers)],
+                dtype=torch.uint64,
+                device=device,
+            )
+            transfer_kv_all_layer_mla_lf_pf(
+                src_layers=src_ptrs,
+                dst=dst,
+                src_indices=src_indices,
+                dst_indices=dst_indices,
+                item_size=item_size,
+                dst_layout_dim=item_size * num_layers,
+                num_layers=num_layers,
+            )
+        elif io_backend == "direct":
+            src_ptrs = [src_layers[i] for i in range(num_layers)]
+            transfer_kv_all_layer_direct_lf_pf(
+                src_ptrs=src_ptrs,
+                dst_ptrs=[dst],
+                src_indices=src_indices,
+                dst_indices=dst_indices,
+                page_size=1,
+            )
+        else:
+            raise ValueError(f"Unsupported io_backend: {io_backend}")
+
+    def load_to_device_per_layer(
+        self,
+        device_pool,
+        host_indices,
+        device_indices,
+        layer_id,
+        io_backend="kernel",
+    ):
+        if self.layout in ["page_first", "page_first_direct"]:
+            self._copy_tensor_pf_lf(
+                src=self.temporal_buffer,
+                dst=device_pool.mamba_cache.temporal[layer_id],
+                src_indices=host_indices,
+                dst_indices=device_indices,
+                layer_id=layer_id,
+                num_layers=self.num_mamba_layers,
+                io_backend=io_backend,
+            )
+            for conv_idx in range(len(self.conv_state_shapes)):
+                self._copy_tensor_pf_lf(
+                    src=self.conv_buffer[conv_idx],
+                    dst=device_pool.mamba_cache.conv[conv_idx][layer_id],
+                    src_indices=host_indices,
+                    dst_indices=device_indices,
+                    layer_id=layer_id,
+                    num_layers=self.num_mamba_layers,
+                    io_backend=io_backend,
+                )
+        else:
+            self._copy_tensor(
+                self.temporal_buffer[layer_id],
+                device_pool.mamba_cache.temporal[layer_id],
+                host_indices,
+                device_indices,
+                io_backend,
+            )
+            for conv_idx in range(len(self.conv_state_shapes)):
+                self._copy_tensor(
+                    self.conv_buffer[conv_idx][layer_id],
+                    device_pool.mamba_cache.conv[conv_idx][layer_id],
+                    host_indices,
+                    device_indices,
+                    io_backend,
+                )
+
+    def backup_from_device_all_layer(
+        self, device_pool, host_indices, device_indices, io_backend="kernel"
+    ):
+        if self.layout in ["page_first", "page_first_direct"]:
+            self._copy_tensor_all_layers_lf_pf(
+                src_layers=device_pool.mamba_cache.temporal,
+                dst=self.temporal_buffer,
+                src_indices=device_indices,
+                dst_indices=host_indices,
+                num_layers=self.num_mamba_layers,
+                device=self.device_pool.device,
+                io_backend=io_backend,
+            )
+            for conv_idx in range(len(self.conv_state_shapes)):
+                self._copy_tensor_all_layers_lf_pf(
+                    src_layers=device_pool.mamba_cache.conv[conv_idx],
+                    dst=self.conv_buffer[conv_idx],
+                    src_indices=device_indices,
+                    dst_indices=host_indices,
+                    num_layers=self.num_mamba_layers,
+                    device=self.device_pool.device,
+                    io_backend=io_backend,
+                )
+        else:
+            for layer_id in range(self.num_mamba_layers):
+                self._copy_tensor(
+                    device_pool.mamba_cache.temporal[layer_id],
+                    self.temporal_buffer[layer_id],
+                    device_indices,
+                    host_indices,
+                    io_backend,
+                )
+                for conv_idx in range(len(self.conv_state_shapes)):
+                    self._copy_tensor(
+                        device_pool.mamba_cache.conv[conv_idx][layer_id],
+                        self.conv_buffer[conv_idx][layer_id],
+                        device_indices,
+                        host_indices,
+                        io_backend,
+                    )
+
+    def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
+        data_page = torch.cat(
+            [
+                self._flatten_tensor_bytes(tensor)
+                for tensor in self._iter_page_tensors(index)
+            ]
+        )
+        return data_page.flatten() if flat else data_page
+
+    def get_dummy_flat_data_page(self) -> torch.Tensor:
+        return torch.zeros(
+            self.page_size * self.size_per_token,
+            dtype=torch.uint8,
+            device=self.device,
+            pin_memory=self.pin_memory,
+        )
+
+    def set_from_flat_data_page(
+        self,
+        index: int,
+        data_page: torch.Tensor,
+    ) -> None:
+        flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1)
+        start = 0
+        for tensor in self._iter_page_tensors(index):
+            num_bytes = tensor.numel() * tensor.element_size()
+            tensor_bytes = flat_bytes[start : start + num_bytes]
+            start += num_bytes
+            restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape)
+            tensor.copy_(restored)
+
+
+@dataclass
+class PoolEntry:
+    name: PoolName
+    host_pool: Any
+    device_pool: Any
+    layer_mapper: Callable[[int], Optional[int]]
+    is_primary_index_anchor: bool = False
+    # Optional eviction callbacks for auto-alloc in HybridCacheController.
+    # host_evict_fn(n): evict n slots from the host pool (used by write()).
+    # device_evict_fn(n): evict n slots from the device pool (used by load()).
+    host_evict_fn: Optional[Callable] = None
+    device_evict_fn: Optional[Callable] = None
+
+
+class HostPoolGroup:
+    def __init__(self, entries: list[PoolEntry]):
+        if not entries:
+            raise ValueError("HostPoolGroup requires at least one pool entry.")
+        self.entries = entries
+        self.entry_map = {entry.name: entry for entry in entries}
+        self.anchor_entry = next(
+            (entry for entry in entries if entry.is_primary_index_anchor),
+            entries[0],
+        )
+
+        self.layout = self.anchor_entry.host_pool.layout
+        self.page_size = self.anchor_entry.host_pool.page_size
+        self.device = self.anchor_entry.host_pool.device
+        self.size = self.anchor_entry.host_pool.size
+
+    def clear(self) -> None:
+        for entry in self.entries:
+            entry.host_pool.clear()
+
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
+        return self.anchor_entry.host_pool.alloc(need_size)
+
+    def free(self, indices: torch.Tensor) -> int:
+        return self.anchor_entry.host_pool.free(indices)
+
+    def get_data_page(self, index, flat: bool = True):
+        return self.anchor_entry.host_pool.get_data_page(index, flat)
+
+    def get_dummy_flat_data_page(self):
+        return self.anchor_entry.host_pool.get_dummy_flat_data_page()
+
+    def set_from_flat_data_page(self, index: int, data_page) -> None:
+        return self.anchor_entry.host_pool.set_from_flat_data_page(index, data_page)
+
+    def load_to_device_per_layer(
+        self,
+        device_pool,
+        host_indices,
+        device_indices,
+        layer_id,
+        io_backend,
+        pool_transfers: Optional[list] = None,
+    ) -> None:
+        # 1. Anchor (KV) transfer
+        anchor = self.anchor_entry
+        local_layer_id = anchor.layer_mapper(layer_id)
+        if local_layer_id is not None and host_indices.numel() > 0:
+            anchor.host_pool.load_to_device_per_layer(
+                anchor.device_pool,
+                host_indices,
+                device_indices,
+                local_layer_id,
+                io_backend,
+            )
+
+        # 2. Extra pool transfers
+        for transfer in pool_transfers or []:
+            entry = self.entry_map.get(transfer.name)
+            if entry is None or transfer.host_indices is None:
+                continue
+            local_layer_id = entry.layer_mapper(layer_id)
+            if local_layer_id is None:
+                continue
+            entry.host_pool.load_to_device_per_layer(
+                entry.device_pool,
+                transfer.host_indices,
+                transfer.device_indices,
+                local_layer_id,
+                io_backend,
+            )
+
+    def backup_from_device_all_layer(
+        self,
+        device_pool,
+        host_indices,
+        device_indices,
+        io_backend,
+        pool_transfers: Optional[list] = None,
+    ) -> None:
+        # 1. Anchor (KV) backup
+        self.anchor_entry.host_pool.backup_from_device_all_layer(
+            self.anchor_entry.device_pool,
+            host_indices,
+            device_indices,
+            io_backend,
+        )
+        # 2. Extra pool backup
+        for transfer in pool_transfers or []:
+            entry = self.entry_map.get(transfer.name)
+            if entry is None or transfer.host_indices is None:
+                continue
+            entry.host_pool.backup_from_device_all_layer(
+                entry.device_pool,
+                transfer.host_indices,
+                transfer.device_indices,
+                io_backend,
+            )
+
+
 class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
     device_pool: NSATokenToKVPool
 

test/registered/4-gpu-models/test_qwen35_models.py

@@ -1,5 +1,6 @@
 import shutil
 import tempfile
+import time
 import unittest
 from types import SimpleNamespace
 
@@ -254,6 +255,10 @@ class TestQwen35WithHiCache(CustomTestCase):
             other_args=[
                 "--tp-size",
                 "4",
+                "--max-mamba-cache-size",
+                "500",
+                "--max-total-tokens",
+                "120000",
                 "--chunked-prefill-size",
                 "2048",
                 "--mamba-scheduler-strategy",
@@ -289,14 +294,14 @@ class TestQwen35WithHiCache(CustomTestCase):
         kill_process_tree(cls.process.pid)
         shutil.rmtree(cls.storage_dir, ignore_errors=True)
 
-    def test_gsm8k(self):
+    def _run_gsm8k(self):
         args = SimpleNamespace(
             model=self.model,
             eval_name="gsm8k",
             num_shots=5,
-            num_examples=200,
+            num_examples=100,
-            max_tokens=16000,
+            max_tokens=4000,
-            num_threads=128,
+            num_threads=50,
             repeat=1,
             temperature=0.6,
             top_p=0.95,
@@ -305,9 +310,30 @@ class TestQwen35WithHiCache(CustomTestCase):
             host="http://127.0.0.1",
             port=int(self.base_url.split(":")[-1]),
         )
-        metrics = run_eval(args)
+        return run_eval(args)
-        print(f"{metrics=}")
+
-        self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"])
+    def test_gsm8k(self):
+        first_metrics = self._run_gsm8k()
+        print(f"first_metrics={first_metrics}")
+        self.assertGreaterEqual(
+            first_metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]
+        )
+
+        print(f"flush cache")
+        time.sleep(2)
+        requests.post(f"{self.base_url}/flush_cache", timeout=10)
+
+        second_metrics = self._run_gsm8k()
+        print(f"second_metrics={second_metrics}")
+        self.assertGreaterEqual(
+            second_metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]
+        )
+        self.assertLessEqual(
+            abs(second_metrics["score"] - first_metrics["score"]),
+            0.05,
+            f"HiCache prefetch accuracy drift too large: "
+            f"first={first_metrics['score']}, second={second_metrics['score']}",
+        )
 
 
 if __name__ == "__main__":