Add a host-memory prefix cache (HiCache-style) for the offline engine

The open-source engine has no KV reuse across requests, so every turn of a
multi-turn conversation prefills its whole prompt. This adds deepinfer/hicache:

- Pinned host pools organized in pages (one page = one 128-token KV block)
  plus slots for the SWA state that the production KV store persists
  (select_state_slice_to_store / load_state_from_store_to_prefill_cache).
  Pages are indexed by a hash chain over the prompt, LRU evicted, written
  through at the end of every prefill (engine hook pe_do_save_kvcache).
- Transfers use one small CUDA kernel copying pages between GPU and pinned
  host memory through UVA pointers on the engine stream ("kernel" I/O
  method, ~180 GiB/s on GB300); completion is tracked with CUDA events.
- A hit is capped so that at least yoco keep_tokens remain to recompute
  (bounded replay rebuilds the decoder SWA state from that suffix), then
  snapped to the nearest 128-aligned position with a stored state; states
  are stored at the last few boundaries of a prompt so short follow-ups hit.
  The sequence is pre-allocated, filled from host and scheduled with
  num_cached_tokens set (HostPrefixCache.try_resume).
- Engine: read-only Sequence.block_table binding (to copy a finished
  prefill's blocks out) and its stub.

Validated on 4x GB300 with DeepSeek-V4.1-Flash through an OpenAI-compatible
server (separate PR): greedy outputs byte-identical with and without the
cache over 8 multi-turn requests; InferenceX AgentX replay at concurrency 1
served 224 requests per hour vs 99, TTFT p50 0.59 s vs 8.9 s, with 89-96%
of prompt tokens served from the cache. tests/hicache/test_host_cache.py
covers the cache logic against a fake engine.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK

csrc/engine/pybind_sequence.cpp

@@ -116,6 +116,8 @@ void pybind_engine_sequence(pybind11::module &engine) {
     bind_sequence_readonly(sequence, "sampling_params", [](const Sequence &seq) -> const auto & { return seq.generation.params; });
     bind_sequence_property(sequence, "multi_modal_data", [](Sequence &seq) -> auto & { return seq.input.multi_modal_data; });
     bind_sequence_property(sequence, "cached_hit_tokens", [](Sequence &seq) -> auto & { return seq.cache.hit_tokens; });
+    // 只读：host prefix cache 落盘时需要 seq 当前占用的 block id
+    bind_sequence_readonly(sequence, "block_table", [](const Sequence &seq) -> const auto & { return seq.cache.block_table; });
     bind_sequence_property(sequence, "last_output_token", [](Sequence &seq) -> auto & { return seq.mtp.last_output_token; });
     bind_sequence_property(sequence, "last_accepted_token", [](Sequence &seq) -> auto & { return seq.mtp.accepted_tokens; });
     bind_sequence_property(sequence, "last_next_n_token", [](Sequence &seq) -> auto & { return seq.mtp.draft_tokens; });

deepinfer/_C/engine.pyi

@@ -455,6 +455,9 @@ class Sequence:
     def cached_hit_tokens(self, arg1: typing.SupportsInt | typing.SupportsIndex) -> None:
         ...
     @property
+    def block_table(self) -> list[int]:
+        ...
+    @property
     def chunk_checkpoints(self) -> list[int]:
         ...
     @chunk_checkpoints.setter

deepinfer/hicache/__init__.py

@@ -0,0 +1,3 @@
+from deepinfer.hicache.host_cache import HostPrefixCache
+
+__all__ = ["HostPrefixCache"]

deepinfer/hicache/host_cache.py

@@ -0,0 +1,432 @@
+"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache.
+
+Layout
+    Two pinned host pools, both organized in pages:
+      - KV pages:    one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block.
+      - State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a
+                     128-token-aligned position with GPUCacheInterface.select_state_slice_to_store.
+    Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)),
+    so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool.
+
+Transfer
+    "kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through
+    UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never
+    blocks the CPU. Completion is tracked with CUDA events polled from the engine loop.
+
+Policy
+    Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last
+    128-aligned position are stored (pages already present are only touched). On a new request, the longest
+    page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to
+    be recomputed, then snapped down to the nearest position that has a stored state; the sequence is
+    pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with
+    num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is
+    rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store).
+
+Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's
+prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import time
+from collections import OrderedDict, deque
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+
+from deepinfer import envs
+from deepinfer.logger import init_logger
+from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface
+
+logger = init_logger(__name__)
+
+
+class _Entry:
+    __slots__ = ("key", "slot", "refs", "ready")
+
+    def __init__(self, key: bytes, slot: int) -> None:
+        self.key = key
+        self.slot = slot
+        self.refs = 0  # in-flight transfers (store or load) touching this slot
+        self.ready = False  # host copy complete and hit-able
+
+
+@dataclass
+class _Chain:
+    tokens: np.ndarray
+    hashes: list[bytes] = field(default_factory=list)
+    counted: bool = False  # lookup stats recorded (a deferred request is looked up repeatedly)
+
+
+@dataclass
+class Stats:
+    lookups: int = 0
+    hits: int = 0
+    hit_tokens: int = 0
+    lookup_tokens: int = 0
+    deferred: int = 0  # hit found but GPU allocation not possible yet
+    stores: int = 0
+    store_skipped: int = 0
+    pages_stored: int = 0
+    pages_loaded: int = 0
+    states_stored: int = 0
+    states_loaded: int = 0
+    evicted_pages: int = 0
+    evicted_states: int = 0
+    errors: int = 0
+
+    def as_dict(self) -> dict[str, int]:
+        return dict(self.__dict__)
+
+
+class HostPrefixCache:
+    def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None:
+        assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)"
+        assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1"
+        assert GPUCacheInterface.cache_info_ready, "cache engine not initialized"
+
+        self.engine = engine
+        self.scheduler = engine.scheduler
+        cache_engine = engine.driver_worker.cache_engine
+        self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks
+        self.gpu_states: torch.Tensor = cache_engine.gpu_states
+        assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8
+        self.device = self.gpu_blocks.device
+        self.block_size: int = engine.cache_config.block_size
+        assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN)
+        self.page_bytes: int = self.gpu_blocks.size(1)
+        self.state_bytes: int = GPUCacheInterface.state_size_to_disk
+        assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes)
+        self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS
+        # The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative.
+        self.min_hit_tokens = 2 * self.block_size
+        # A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only
+        # at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a
+        # block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA
+        # ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore.
+        self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4"))
+        self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference
+
+        self.num_pages = max(0, size_bytes // self.page_bytes)
+        self.num_state_slots = num_state_slots
+        self.staging_slots = staging_slots
+        assert self.num_pages > 0 and num_state_slots > 0
+        self.gpu_blocks_ptr = self.gpu_blocks.data_ptr()
+        # Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be
+        # cached before a prefill completes, and an unused cache then costs no pinned memory).
+        self._allocated = False
+        self.host_pages = self.host_states = None
+        self.host_pages_ptr = self.host_states_ptr = 0
+        self._staging: list[torch.Tensor] = []
+        self._staging_idx = 0
+        self.kernel = None
+
+        self.pages: OrderedDict[bytes, _Entry] = OrderedDict()
+        self.states: OrderedDict[bytes, _Entry] = OrderedDict()
+        self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1))
+        self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1))
+        self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque()
+        self._chains: dict[str, _Chain] = {}
+        self.stats = Stats()
+        self.enabled = True
+        # DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel).
+        self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1"
+        logger.info(
+            "HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d",
+            self.num_pages,
+            self.page_bytes,
+            self.num_pages * self.page_bytes / 2**30,
+            num_state_slots,
+            self.state_bytes,
+            num_state_slots * self.state_bytes / 2**30,
+            self.keep_tokens,
+        )
+
+    def _allocate(self) -> None:
+        if self._allocated:
+            return
+        from deepinfer.hicache.kernels import copy_pages_module
+
+        t0 = time.perf_counter()
+        self.kernel = copy_pages_module()
+        self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True)
+        self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True)
+        self.host_pages_ptr = self.host_pages.data_ptr()
+        self.host_states_ptr = self.host_states.data_ptr()
+        # Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream.
+        self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)]
+        self._allocated = True
+        logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0)
+
+    # ------------------------------------------------------------------ hashing
+
+    def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None:
+        chain = self._chains.get(request_id)
+        if chain is None and tokens is not None:
+            chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32))
+        return chain
+
+    def _extend(self, chain: _Chain, num_pages: int) -> None:
+        hashes, toks, bs = chain.hashes, chain.tokens, self.block_size
+        prev = hashes[-1] if hashes else b""
+        for i in range(len(hashes), num_pages):
+            prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest()
+            hashes.append(prev)
+
+    def forget(self, request_id: str) -> None:
+        self._chains.pop(request_id, None)
+
+    # ------------------------------------------------------------------ transfers
+
+    def _ids(self, ids: list[int]) -> torch.Tensor:
+        return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True)
+
+    def _next_staging(self) -> torch.Tensor:
+        t = self._staging[self._staging_idx]
+        self._staging_idx = (self._staging_idx + 1) % len(self._staging)
+        return t
+
+    def _track(self, entries: list[_Entry]) -> None:
+        ev = torch.cuda.Event()
+        ev.record()
+        self._inflight.append((ev, entries))
+
+    def _checkpoint(self, what: str) -> None:
+        if self._debug_sync:
+            from deepinfer.utils import allow_cuda_sync
+
+            with allow_cuda_sync():
+                torch.cuda.synchronize()
+            logger.info("hicache debug: ok after %s", what)
+
+    def poll(self) -> None:
+        """Retire completed transfers (events on one stream complete in order, so only the head is checked)."""
+        while self._inflight and self._inflight[0][0].query():
+            _, entries = self._inflight.popleft()
+            for e in entries:
+                e.refs -= 1
+                e.ready = True
+
+    def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int:
+        """Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found."""
+        victims, busy = [], []
+        for key, e in pool.items():
+            if len(victims) >= need:
+                break
+            if e.refs or not e.ready:
+                busy.append(key)
+                continue
+            victims.append(key)
+        for key in busy:
+            pool.move_to_end(key)
+        if len(victims) < need:
+            return 0
+        for key in victims:
+            free.append(pool.pop(key).slot)
+        return len(victims)
+
+    def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool:
+        if need > len(free) + len(pool):
+            return False
+        if len(free) < need:
+            evicted = self._evict(pool, free, need - len(free))
+            if pool is self.pages:
+                self.stats.evicted_pages += evicted
+            else:
+                self.stats.evicted_states += evicted
+        return len(free) >= need
+
+    # ------------------------------------------------------------------ load (hit)
+
+    def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int:
+        """Fill `seq` (not yet in the scheduler) from the host cache.
+
+        Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not
+        allocatable right now, sequence untouched, retry later).
+        """
+        if not self.enabled:
+            return 0
+        try:
+            return self._try_resume(seq, request_id, prompt_ids)
+        except Exception:  # noqa: BLE001 - never take the engine down for a cache problem
+            logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id)
+            self._fail()
+            return 0
+
+    def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int:
+        bs = self.block_size
+        n = len(prompt_ids)
+        chain = self._chain(request_id, prompt_ids)
+        if not chain.counted:
+            chain.counted = True
+            self.stats.lookups += 1
+            self.stats.lookup_tokens += n
+        limit = ((n - self.keep_tokens) // bs) * bs  # the replay rectangle [n - keep, n) must be recomputed
+        if limit < self.min_hit_tokens:
+            return 0
+        max_pages = limit // bs
+        self._extend(chain, max_pages)
+        matched = 0
+        for h in chain.hashes[:max_pages]:
+            e = self.pages.get(h)
+            if e is None or not e.ready:
+                break
+            matched += 1
+        b = min(matched * bs, limit)
+        while b >= self.min_hit_tokens:
+            st = self.states.get(chain.hashes[b // bs - 1])
+            if st is not None and st.ready:
+                break
+            b -= bs
+        else:
+            return 0
+
+        try:
+            allocated = self.scheduler.try_allocate_seq(seq)
+        except RuntimeError:  # can never fit (prompt longer than the whole KV pool): let the scheduler reject it
+            return 0
+        if not allocated:
+            self.stats.deferred += 1
+            return -1
+        num_pages = b // bs
+        page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)]
+        for e in page_entries:
+            e.refs += 1
+            self.pages.move_to_end(e.key)
+        st.refs += 1
+        self.states.move_to_end(st.key)
+
+        block_table = seq.block_table
+        self.kernel.copy_pages(
+            self.host_pages_ptr,
+            self.gpu_blocks_ptr,
+            self._ids([e.slot for e in page_entries]),
+            self._ids(list(block_table[:num_pages])),
+            self.page_bytes,
+        )
+        self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})")
+        staging = self._next_staging()
+        self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes)
+        self._checkpoint("load: state copy H2D")
+        GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states)
+        self._checkpoint("load: load_state_to_cache")
+        self._track(page_entries + [st])
+
+        seq.num_cached_tokens = b
+        seq.cached_hit_tokens = b
+        self.stats.hits += 1
+        self.stats.hit_tokens += b
+        self.stats.pages_loaded += num_pages
+        self.stats.states_loaded += 1
+        return b
+
+    # ------------------------------------------------------------------ store (write-through at prefill end)
+
+    def on_prefill_step(self, scheduled, finished) -> None:
+        """Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed."""
+        if not self.enabled:
+            return
+        try:
+            for seq in scheduled:
+                if seq.chunked_rest_len == 0:  # last chunk of this prompt
+                    self._store(seq)
+        except Exception:  # noqa: BLE001
+            logger.exception("host prefix cache: store failed")
+            self._fail()
+
+    def _store(self, seq) -> None:
+        chain = self._chains.get(seq.request_id)
+        if chain is None:  # request did not go through try_resume: hash the prompt from the sequence itself
+            chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32))
+        self._allocate()
+        bs = self.block_size
+        b = (seq.get_len() // bs) * bs
+        if b < self.min_hit_tokens:
+            return
+        num_pages = b // bs
+        self._extend(chain, num_pages)
+        missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages]
+        for i in range(num_pages):
+            if i not in missing:
+                self.pages.move_to_end(chain.hashes[i])
+        # State boundaries: b, b - bs, ... while the window ending there is still in the ring.
+        seq_len = seq.get_len()
+        new_states = []
+        for k in range(self.state_lookback_blocks):
+            pos = b - k * bs
+            if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack:
+                break
+            key = chain.hashes[pos // bs - 1]
+            if key in self.states:
+                self.states.move_to_end(key)
+            else:
+                new_states.append((pos, key))
+        if not missing and not new_states:
+            return
+        if not self._reserve(self.pages, self.free_pages, len(missing)):
+            self.stats.store_skipped += 1
+            return
+        # States are best effort: keep the newest boundaries that fit.
+        while new_states and not self._reserve(self.states, self.free_states, len(new_states)):
+            new_states.pop()
+        if not missing and not new_states:
+            self.stats.store_skipped += 1
+            return
+
+        block_table = seq.block_table
+        entries: list[_Entry] = []
+        if missing:
+            new = []
+            for i in missing:
+                e = _Entry(chain.hashes[i], self.free_pages.pop())
+                e.refs = 1
+                self.pages[e.key] = e
+                new.append(e)
+            self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})")
+            self.kernel.copy_pages(
+                self.gpu_blocks_ptr,
+                self.host_pages_ptr,
+                self._ids([block_table[i] for i in missing]),
+                self._ids([e.slot for e in new]),
+                self.page_bytes,
+            )
+            self._checkpoint("store: page copy D2H")
+            entries += new
+            self.stats.pages_stored += len(new)
+        for pos, key in new_states:
+            st = _Entry(key, self.free_states.pop())
+            st.refs = 1
+            self.states[key] = st
+            staging = self._next_staging()
+            GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging)
+            self._checkpoint(f"store: select_state_slice at {pos}")
+            self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes)
+            self._checkpoint("store: state copy D2H")
+            entries.append(st)
+            self.stats.states_stored += 1
+        self._track(entries)
+        self.stats.stores += 1
+
+    # ------------------------------------------------------------------ misc
+
+    def _fail(self) -> None:
+        self.stats.errors += 1
+        if self.stats.errors >= 3:
+            logger.error("host prefix cache: disabled after %d errors", self.stats.errors)
+            self.enabled = False
+
+    def stats_dict(self) -> dict[str, int | float]:
+        d = self.stats.as_dict()
+        d.update(
+            enabled=int(self.enabled),
+            pages_used=len(self.pages),
+            pages_total=self.num_pages,
+            states_used=len(self.states),
+            states_total=len(self.free_states) + len(self.states),
+            inflight=len(self._inflight),
+            hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0,
+        )
+        return d

deepinfer/hicache/kernels.py

@@ -0,0 +1,25 @@
+"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+_module = None
+
+
+def copy_pages_module():
+    global _module
+    if _module is None:
+        from torch.utils.cpp_extension import load
+
+        # Do not touch TORCH_CUDA_ARCH_LIST here: the engine's own JIT kernels (deep_jit / TileLang) run
+        # in this process and must keep their arch-specific (sm_1xxa) targets. torch's loader derives the
+        # arch for this trivial kernel from the visible device.
+        here = Path(__file__).resolve().parent
+        _module = load(
+            name="deepinfer_hicache_page_copy",
+            sources=[str(here / "page_copy.cpp"), str(here / "page_copy.cu")],
+            extra_cuda_cflags=["-O3"],
+            verbose=False,
+        )
+    return _module

deepinfer/hicache/page_copy.cpp

@@ -0,0 +1,9 @@
+#include <torch/extension.h>
+
+void copy_pages(int64_t src_base, int64_t dst_base, torch::Tensor src_ids, torch::Tensor dst_ids, int64_t page_bytes);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("copy_pages", &copy_pages, py::arg("src_base"), py::arg("dst_base"), py::arg("src_ids"), py::arg("dst_ids"),
+        py::arg("page_bytes"),
+        "Copy pages src_base[src_ids[i]] -> dst_base[dst_ids[i]] on the current CUDA stream. Either base may be pinned host memory.");
+}

deepinfer/hicache/page_copy.cu

@@ -0,0 +1,45 @@
+// Page copy between two pools addressed by UVA pointers: either pool may live in GPU memory or in
+// pinned host memory (the "kernel" transfer method of SGLang's HiCache). One thread block (row of the
+// grid) per page; pages are copied as 16-byte vectors on the current CUDA stream.
+#include <torch/extension.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cstdint>
+
+namespace {
+
+__global__ void copy_pages_kernel(const uint8_t *__restrict__ src_base, uint8_t *__restrict__ dst_base,
+    const int32_t *__restrict__ src_ids, const int32_t *__restrict__ dst_ids, int64_t page_bytes) {
+    const int64_t n16 = page_bytes >> 4;
+    const int4 *src = reinterpret_cast<const int4 *>(src_base + static_cast<int64_t>(src_ids[blockIdx.x]) * page_bytes);
+    int4 *dst = reinterpret_cast<int4 *>(dst_base + static_cast<int64_t>(dst_ids[blockIdx.x]) * page_bytes);
+    const int64_t stride = static_cast<int64_t>(gridDim.y) * blockDim.x;
+    for (int64_t i = static_cast<int64_t>(blockIdx.y) * blockDim.x + threadIdx.x; i < n16; i += stride) {
+        dst[i] = src[i];
+    }
+}
+
+} // namespace
+
+void copy_pages(int64_t src_base, int64_t dst_base, torch::Tensor src_ids, torch::Tensor dst_ids, int64_t page_bytes) {
+    TORCH_CHECK(page_bytes > 0 && page_bytes % 16 == 0, "page_bytes must be a positive multiple of 16, got ", page_bytes);
+    TORCH_CHECK(src_ids.is_cuda() && dst_ids.is_cuda(), "page ids must be CUDA tensors");
+    TORCH_CHECK(src_ids.scalar_type() == torch::kInt32 && dst_ids.scalar_type() == torch::kInt32, "page ids must be int32");
+    TORCH_CHECK(src_ids.is_contiguous() && dst_ids.is_contiguous(), "page ids must be contiguous");
+    TORCH_CHECK(src_ids.numel() == dst_ids.numel(), "src/dst page id count mismatch");
+    const int64_t n = src_ids.numel();
+    if (n == 0) {
+        return;
+    }
+    constexpr int threads = 256;
+    const int64_t n16 = page_bytes / 16;
+    const unsigned blocks_y = static_cast<unsigned>(std::min<int64_t>((n16 + threads - 1) / threads, 16));
+    const dim3 grid(static_cast<unsigned>(n), blocks_y);
+    auto stream = c10::cuda::getCurrentCUDAStream();
+    copy_pages_kernel<<<grid, threads, 0, stream.stream()>>>(reinterpret_cast<const uint8_t *>(src_base),
+        reinterpret_cast<uint8_t *>(dst_base), src_ids.data_ptr<int32_t>(), dst_ids.data_ptr<int32_t>(), page_bytes);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}

tests/hicache/test_host_cache.py

@@ -0,0 +1,190 @@
+"""HostPrefixCache logic against a fake engine: hashing/matching, the keep-tokens cap, snapping to a stored
+state, store dedupe, LRU eviction and deferral. The real state slice kernels are replaced by plain copies;
+the page transfer kernel is the real one. Needs one GPU."""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pytest
+import torch
+
+os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact"
+os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = "140"
+os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1"
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+PAGE = 114048
+STATE = 4096
+BLOCK = 128
+
+
+class FakeSeq:
+    def __init__(self, request_id: str, tokens: np.ndarray) -> None:
+        self.request_id = request_id
+        self.tokens = tokens
+        self.block_table: list[int] = []
+        self.state_id = -1
+        self.num_cached_tokens = 0
+        self.cached_hit_tokens = 0
+        self.chunked_rest_len = 0
+        self.data = self  # SequenceData stand-in
+
+    def get_len(self) -> int:
+        return len(self.tokens)
+
+    def get_token_ids(self) -> list[int]:
+        return self.tokens.tolist()
+
+
+class FakeScheduler:
+    def __init__(self, num_blocks: int, num_states: int) -> None:
+        self.free_blocks = list(range(num_blocks))
+        self.free_states = list(range(num_states))
+        self.refuse = False
+
+    def try_allocate_seq(self, seq: FakeSeq) -> bool:
+        n = (len(seq.tokens) + BLOCK - 1) // BLOCK
+        if self.refuse or len(self.free_blocks) < n or not self.free_states:
+            return False
+        seq.block_table = [self.free_blocks.pop(0) for _ in range(n)]
+        seq.state_id = self.free_states.pop(0)
+        return True
+
+    def get_num_unfinished_seqs(self) -> int:
+        return 0
+
+
+class _Obj:
+    pass
+
+
+@pytest.fixture
+def engine(monkeypatch):
+    from deepinfer.worker.cache_engine import GPUCacheInterface
+
+    monkeypatch.setattr(GPUCacheInterface, "cache_info_ready", True, raising=False)
+    monkeypatch.setattr(GPUCacheInterface, "state_size_to_disk", STATE, raising=False)
+    monkeypatch.setattr(GPUCacheInterface, "swa_block_size_kvstore", 134 + 512 + 140, raising=False)
+    monkeypatch.setattr(GPUCacheInterface, "swa_block_size_inference", 134, raising=False)
+    monkeypatch.setattr(
+        GPUCacheInterface,
+        "select_state_slice_to_store",
+        classmethod(lambda cls, sid, pos, cache, state_tensor=None, TO_CE=False: state_tensor.copy_(cache[sid, :STATE])),
+    )
+    monkeypatch.setattr(
+        GPUCacheInterface,
+        "load_state_from_store_to_prefill_cache",
+        classmethod(lambda cls, sid, pos, tensor, cache: cache[sid, :STATE].copy_(tensor)),
+    )
+    e = _Obj()
+    e.scheduler = FakeScheduler(num_blocks=512, num_states=8)
+    e.cache_config = _Obj()
+    e.cache_config.block_size = BLOCK
+    e.driver_worker = _Obj()
+    e.driver_worker.cache_engine = _Obj()
+    e.driver_worker.cache_engine.gpu_blocks = torch.zeros((512, PAGE), dtype=torch.uint8, device="cuda")
+    e.driver_worker.cache_engine.gpu_states = torch.zeros((8, STATE), dtype=torch.uint8, device="cuda")
+    return e
+
+
+def _prefill(engine, cache, seq: FakeSeq) -> None:
+    """Simulate the engine: allocate (if the cache did not), write random KV/state, run the store hook."""
+    if not seq.block_table:
+        assert engine.scheduler.try_allocate_seq(seq)
+    blocks = engine.driver_worker.cache_engine.gpu_blocks
+    states = engine.driver_worker.cache_engine.gpu_states
+    start = seq.num_cached_tokens // BLOCK
+    for b in seq.block_table[start:]:
+        blocks[b] = torch.randint(0, 256, (PAGE,), dtype=torch.uint8, device="cuda")
+    states[seq.state_id] = torch.randint(0, 256, (STATE,), dtype=torch.uint8, device="cuda")
+    cache.on_prefill_step([seq], [])
+    torch.cuda.synchronize()
+    cache.poll()
+
+
+def _tokens(n: int, seed: int) -> np.ndarray:
+    return np.random.default_rng(seed).integers(0, 100000, n, dtype=np.int32)
+
+
+def test_hit_snaps_to_stored_state_and_restores_bytes(engine):
+    from deepinfer.hicache import HostPrefixCache
+
+    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    blocks = engine.driver_worker.cache_engine.gpu_blocks
+    states = engine.driver_worker.cache_engine.gpu_states
+
+    a = FakeSeq("a", _tokens(1000, 1))
+    assert cache.try_resume(a, "a", a.tokens) == 0  # cold
+    _prefill(engine, cache, a)
+    # 896 tokens = 7 pages; states at the last 4 boundaries (896, 768, 640, 512)
+    assert cache.stats.stores == 1 and cache.stats.pages_stored == 7 and cache.stats.states_stored == 4
+    a_pages = [blocks[b].clone() for b in a.block_table[:7]]
+    a_state = states[a.state_id, :STATE].clone()
+
+    # Same prefix plus a new turn: hit at 896 (= 7 full pages, and >= keep_tokens remain to recompute).
+    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
+    assert cache.try_resume(b, "b", b.tokens) == 896
+    torch.cuda.synchronize()
+    cache.poll()
+    assert b.num_cached_tokens == 896 and b.cached_hit_tokens == 896
+    assert all(torch.equal(blocks[b.block_table[i]], a_pages[i]) for i in range(7))
+    assert torch.equal(states[b.state_id, :STATE], a_state)
+    assert cache.stats.hits == 1 and cache.stats.hit_tokens == 896
+
+    # Same 950-token prefix but a different tail at 1000 tokens: pages match up to 896, the keep cap allows
+    # only 768, and a state exists there -> hit at 768 (the lookback states make short extensions hit).
+    c = FakeSeq("c", np.concatenate([a.tokens[:950], _tokens(50, 3)]))
+    assert cache.try_resume(c, "c", c.tokens) == 768
+    torch.cuda.synchronize()
+    cache.poll()
+    _prefill(engine, cache, c)  # pages 0..6 and all four states already exist: nothing new is stored
+    assert cache.stats.stores == 1 and cache.stats.pages_stored == 7 and cache.stats.states_stored == 4
+
+    # b's prefill stores its remaining pages (1500 -> 11 pages, 4 new) and states at 1408..1024.
+    _prefill(engine, cache, b)
+    assert cache.stats.pages_stored == 11 and cache.stats.states_stored == 8
+    d = FakeSeq("d", np.concatenate([b.tokens, _tokens(300, 4)]))
+    assert cache.try_resume(d, "d", d.tokens) == 1408
+
+
+def test_eviction_and_capacity(engine):
+    from deepinfer.hicache import HostPrefixCache
+
+    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=2)
+    a = FakeSeq("a", _tokens(7 * BLOCK + 10, 1))
+    _prefill(engine, cache, a)
+    assert len(cache.pages) == 7
+
+    big = FakeSeq("big", _tokens(60 * BLOCK, 2))  # does not fit at all
+    _prefill(engine, cache, big)
+    assert cache.stats.store_skipped == 1 and len(cache.pages) == 7
+
+    e = FakeSeq("e", _tokens(40 * BLOCK, 3))  # fits without eviction (43 free)
+    _prefill(engine, cache, e)
+    assert cache.stats.evicted_pages == 0 and len(cache.pages) == 47
+
+    f = FakeSeq("f", _tokens(20 * BLOCK, 4))  # needs 20, 3 free -> evict 17 LRU pages (a's 7 first)
+    _prefill(engine, cache, f)
+    assert cache.stats.evicted_pages == 17 and len(cache.pages) == 50
+    assert (
+        cache.try_resume(FakeSeq("a2", np.concatenate([a.tokens, _tokens(300, 5)])), "a2", np.concatenate([a.tokens, _tokens(300, 5)])) == 0
+    )
+    # Only two state slots: each store keeps the two newest boundaries that fit, evicting LRU (a's, then e's).
+    assert len(cache.states) == 2 and cache.stats.evicted_states == 4
+
+
+def test_deferred_when_gpu_cannot_allocate(engine):
+    from deepinfer.hicache import HostPrefixCache
+
+    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=2)
+    a = FakeSeq("a", _tokens(1000, 1))
+    _prefill(engine, cache, a)
+    engine.scheduler.refuse = True
+    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
+    assert cache.try_resume(b, "b", b.tokens) == -1
+    assert b.block_table == [] and b.num_cached_tokens == 0 and cache.stats.deferred == 1
+    engine.scheduler.refuse = False
+    assert cache.try_resume(b, "b", b.tokens) == 896