Merge 625606db4eb8fafd8dc5243dd02c328bf38dafc5 into 553741557a57df44a46fb3ab72b6aafeaf945337

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,642 @@
+"""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).
+
+L3 (optional): with DEEPINFER_HICACHE_MOONCAKE_STORE set, every stored page group and state is also written behind
+to a Mooncake Store (deepinfer/hicache/mooncake_store.py), and a host miss is extended from the Store, so prefixes
+are shared across DP ranks, server instances and restarts.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import time
+from collections import OrderedDict, deque
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+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
+
+if TYPE_CHECKING:
+    from deepinfer._C.engine import Sequence
+    from deepinfer.engine.llm_engine import LLMEngine
+
+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)
+    l3_tried: bool = False  # L3 consulted once for this request
+
+
+@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
+    l3_lookups: int = 0
+    l3_hits: int = 0
+    l3_pages_loaded: int = 0
+    l3_states_loaded: int = 0
+    l3_get_failed: int = 0
+    l3_get_seconds: float = 0.0
+    l3_errors: int = 0
+
+    def as_dict(self) -> dict[str, int]:
+        return dict(self.__dict__)
+
+
+class HostPrefixCache:
+    def __init__(self, engine: LLMEngine, *, 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
+        self.l3 = self._make_l3()
+        self.l3_group = self.l3.group_pages if self.l3 is not None else 0
+        self._l3_uploaded: OrderedDict[str, None] = OrderedDict()  # recently put/seen L3 keys (skip re-puts)
+        # 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 _make_l3(self):  # noqa: ANN202
+        from deepinfer.hicache import mooncake_store
+
+        config = mooncake_store.config_from_env()
+        if config is None:
+            return None
+        model_cfg = os.path.join(str(self.engine.model_config.model), "config.json")
+        with open(model_cfg, "rb") as f:
+            cfg_hash = hashlib.sha256(f.read()).hexdigest()
+        fp = mooncake_store.layout_fingerprint(
+            dict(
+                config=cfg_hash,
+                page_bytes=self.page_bytes,
+                state_bytes=self.state_bytes,
+                block_size=self.block_size,
+                extra_state=int(envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE),
+                extra_size=int(envs.DEEPINFER_PREFILL_STORE_EXTRA_SIZE),
+                backtracking=int(envs.DEEPINFER_MAX_BACKTRACKING_TOKENS),
+                swa_kvstore=int(GPUCacheInterface.swa_block_size_kvstore),
+            )
+        )
+        group = int(os.environ.get("DEEPINFER_HICACHE_L3_GROUP_PAGES", "4"))
+        return mooncake_store.MooncakeL3(config, fp, group)
+
+    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)]
+        if self.l3 is not None:
+            self.l3.register(self.host_pages_ptr, self.host_pages.numel())
+            self.l3.register(self.host_states_ptr, self.host_states.numel())
+        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], uploads: list | None = None) -> None:
+        ev = torch.cuda.Event()
+        ev.record()
+        self._inflight.append((ev, entries, uploads or []))
+
+    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, uploads = self._inflight.popleft()
+            for e in entries:
+                e.refs -= 1
+                e.ready = True
+            for up in uploads:
+                self._l3_submit(*up)
+        if self.l3 is not None:
+            for held in self.l3.drain_done():
+                for e in held:
+                    e.refs -= 1
+
+    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: Sequence, 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: Sequence, 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)
+        if self.l3 is not None and not chain.l3_tried and self.l3.readable():
+            chain.l3_tried = True
+            try:
+                self._l3_fill(chain, max_pages, limit)
+            except Exception:  # noqa: BLE001 - fall back to the host-only lookup
+                logger.exception("host prefix cache: L3 fetch failed for %s", request_id)
+                self.stats.l3_errors += 1
+                self.l3.breaker.trip("L3 fetch raised")
+        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: list[Sequence], finished: list[Sequence]) -> 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: Sequence) -> 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))
+        uploads = self._l3_plan(chain, num_pages, [key for _, key in new_states])
+        if not missing and not new_states:
+            for up in uploads:  # everything is already on the host: upload right away
+                self._l3_submit(*up)
+            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, uploads)
+        self.stats.stores += 1
+
+    # ------------------------------------------------------------------ L3 (Mooncake Store)
+
+    def _l3_seen(self, key: str) -> bool:
+        if key in self._l3_uploaded:
+            self._l3_uploaded.move_to_end(key)
+            return True
+        return False
+
+    def _l3_mark(self, key: str) -> None:
+        self._l3_uploaded[key] = None
+        while len(self._l3_uploaded) > 1_000_000:
+            self._l3_uploaded.popitem(last=False)
+
+    def _l3_plan(self, chain: _Chain, num_pages: int, new_state_keys: list[bytes]) -> list:
+        """Page groups of this prompt and its new states the Store has not seen yet: [(page hashes, groups, state keys)]."""
+        if self.l3 is None:
+            return []
+        g = self.l3_group
+        groups = [i for i in range(num_pages // g) if not self._l3_seen(self.l3.group_key(chain.hashes[(i + 1) * g - 1]))]
+        state_keys = [k for k in new_state_keys if not self._l3_seen(self.l3.state_key(k))]
+        if not groups and not state_keys:
+            return []
+        return [(chain.hashes[:num_pages], groups, state_keys)]
+
+    def _l3_submit(self, hashes: list[bytes], groups: list[int], state_keys: list[bytes]) -> None:
+        """Pin the host slots of the planned groups/states and queue a write-behind put. Entries that are gone (evicted)
+        or not ready are skipped."""
+        g = self.l3_group
+        keys, ptrs, sizes, held = [], [], [], []
+        for i in groups:
+            es = [self.pages.get(h) for h in hashes[i * g : (i + 1) * g]]
+            if any(e is None or not e.ready for e in es):
+                continue
+            keys.append(self.l3.group_key(hashes[(i + 1) * g - 1]))
+            ptrs.append([self.host_pages_ptr + e.slot * self.page_bytes for e in es])
+            sizes.append([self.page_bytes] * g)
+            held += es
+        for k in state_keys:
+            st = self.states.get(k)
+            if st is None or not st.ready:
+                continue
+            keys.append(self.l3.state_key(k))
+            ptrs.append([self.host_states_ptr + st.slot * self.state_bytes])
+            sizes.append([self.state_bytes])
+            held.append(st)
+        if not keys:
+            return
+        for e in held:
+            e.refs += 1
+        for k in keys:
+            self._l3_mark(k)
+        self.l3.submit_put(keys, ptrs, sizes, held)
+
+    def _l3_fill(self, chain: _Chain, max_pages: int, limit: int) -> None:
+        """Extend the host cache with page groups and a state from the Store, so the host lookup that follows hits
+        deeper. No-op when the Store adds nothing beyond what the host already has."""
+        bs, g, l3 = self.block_size, self.l3_group, self.l3
+        self._allocate()  # pools are lazy; a fresh instance can hit L3 before it ever stored anything
+        host_pages = 0
+        for h in chain.hashes[:max_pages]:
+            e = self.pages.get(h)
+            if e is None or not e.ready:
+                break
+            host_pages += 1
+        first_group = host_pages // g
+        n_groups = max_pages // g - first_group
+        if n_groups <= 0:
+            return
+        self.stats.l3_lookups += 1
+        group_keys = [l3.group_key(chain.hashes[(first_group + i + 1) * g - 1]) for i in range(n_groups)]
+        present = l3.exists(group_keys)
+        k = 0
+        while k < n_groups and present[k]:
+            k += 1
+        if k == 0:
+            return
+        l3_pages = (first_group + k) * g  # pages available in host + L3
+        # deepest boundary <= min(l3_pages * bs, limit) beyond the host-only reach that has a state (host or L3)
+        top = min(l3_pages * bs, limit) // bs
+        candidates = [p for p in range(top, host_pages, -1) if p * bs >= self.min_hit_tokens]
+        if not candidates:
+            return
+        host_state = {p: (self.states.get(chain.hashes[p - 1]) is not None) for p in candidates}
+        need_l3 = [p for p in candidates if not host_state[p]]
+        in_l3 = dict(zip(need_l3, l3.exists([l3.state_key(chain.hashes[p - 1]) for p in need_l3]), strict=True))
+        pick = next((p for p in candidates if host_state[p] or in_l3.get(p)), None)
+        if pick is None:
+            return
+        state_from_l3 = not host_state[pick]
+        fetch_groups = list(range(first_group, (pick + g - 1) // g))  # groups covering pages [first_group*g, pick)
+
+        # Pin the host prefix so reserving slots below cannot evict it, then reserve.
+        prefix = [self.pages[h] for h in chain.hashes[:host_pages]]
+        for e in prefix:
+            e.refs += 1
+        try:
+            need_pages = len(fetch_groups) * g
+            if not self._reserve(self.pages, self.free_pages, need_pages):
+                return
+            if state_from_l3 and not self._reserve(self.states, self.free_states, 1):
+                return
+            slots = [self.free_pages.pop() for _ in range(need_pages)]
+            keys = [l3.group_key(chain.hashes[(i + 1) * g - 1]) for i in fetch_groups]
+            ptrs = [[self.host_pages_ptr + s * self.page_bytes for s in slots[j * g : (j + 1) * g]] for j in range(len(fetch_groups))]
+            sizes = [[self.page_bytes] * g for _ in fetch_groups]
+            st_slot = None
+            if state_from_l3:
+                st_slot = self.free_states.pop()
+                keys.append(l3.state_key(chain.hashes[pick - 1]))
+                ptrs.append([self.host_states_ptr + st_slot * self.state_bytes])
+                sizes.append([self.state_bytes])
+            t0 = time.perf_counter()
+            ok = l3.get(keys, ptrs, sizes)
+            self.stats.l3_get_seconds += time.perf_counter() - t0
+            if not all(ok):  # evicted from the Store in between, or a transfer error: drop the whole fetch
+                self.stats.l3_get_failed += 1
+                self.free_pages.extend(slots)
+                if st_slot is not None:
+                    self.free_states.append(st_slot)
+                return
+            loaded = 0
+            for j, gi in enumerate(fetch_groups):
+                for t in range(g):
+                    page = gi * g + t
+                    slot = slots[j * g + t]
+                    h = chain.hashes[page]
+                    if page >= pick or h in self.pages:  # beyond the hit, or the host already holds this page
+                        self.free_pages.append(slot)
+                        continue
+                    e = _Entry(h, slot)
+                    e.ready = True
+                    self.pages[h] = e
+                    loaded += 1
+            if st_slot is not None:
+                st = _Entry(chain.hashes[pick - 1], st_slot)
+                st.ready = True
+                self.states[st.key] = st
+                self.stats.l3_states_loaded += 1
+            self.stats.l3_hits += 1
+            self.stats.l3_pages_loaded += loaded
+        finally:
+            for e in prefix:
+                e.refs -= 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,
+        )
+        if self.l3 is not None:
+            d.update(self.l3.stats_dict())
+        return d

deepinfer/hicache/kernels.py

@@ -0,0 +1,26 @@
+"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import ModuleType
+
+_module = None
+
+
+def copy_pages_module() -> ModuleType:
+    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/mooncake_store.py

@@ -0,0 +1,224 @@
+"""Mooncake Store as an L3 tier under HostPrefixCache (GPU -> pinned host pool -> Mooncake Store).
+
+Enabled by DEEPINFER_HICACHE_MOONCAKE_STORE: a JSON object (or "@/path/to/file.json") with Mooncake's
+`MooncakeDistributedStore.setup(config)` keys, e.g.
+
+    {"local_hostname": "127.0.0.1", "metadata_server": "http://127.0.0.1:18080/metadata",
+     "master_server_addr": "127.0.0.1:50551", "protocol": "tcp", "global_segment_size": 0}
+
+Objects
+    page groups  dsv4kv/<fp>/g<G>/<chain hash of the group's last page>  G consecutive KV pages (G * page_bytes)
+    states       state/<fp>/<chain hash at the 128-aligned boundary>      the SWA stored region (state_bytes)
+    <fp> versions the layout (checkpoint config, page/state sizes, env-dependent state layout), so a layout change
+    is a clean miss. The state namespace follows the contract discussed in kvcache-ai/Mooncake#4803. Pages hold all
+    layers of a block, so they are independent of the DP rank: any rank or instance can reuse them.
+    Pages are grouped (DEEPINFER_HICACHE_L3_GROUP_PAGES, default 4 = 512 tokens) because per-object overhead dominates
+    for single 114 KB pages (0.17 GB/s vs 1.7 GB/s at 4 pages over TCP on one host).
+
+Data moves between the Store and the (registered) pinned host pools with batch_put_from_multi_buffers /
+batch_get_into_multi_buffers, so no extra copies. Puts run on a background thread (write-behind); the engine thread
+learns about completions through drain_done(). Gets are synchronous (called from try_resume).
+
+Failure handling. A Store transfer that cannot complete blocks for Mooncake's fixed 60 s transfer timeout, so one bad
+get stalls the engine thread that long. A circuit breaker therefore opens on the first transfer failure (any put or
+get error other than "not found" / "already exists") that persists after one retry, or on a get slower than
+DEEPINFER_HICACHE_L3_SLOW_S (default 5 s). The retry absorbs sporadic transfer errors (seen at ~0.015% of EFA puts).
+While it is open, reads are skipped (host-only lookup) and queued puts are dropped. It closes after
+DEEPINFER_HICACHE_L3_COOLDOWN_S (default 30 s, doubled on each consecutive trip up to 10 min).
+
+TCP transport limits. Mooncake <= 0.3.13 queues at most 1024 transfers per peer, and one long prompt stores more
+buffers than that in a single batch put; the rejected transfers hang until the 60 s timeout. Unless they are already
+set, MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER and MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER default to 65535, which is the
+default since kvcache-ai/Mooncake#3710.
+"""
+
+from __future__ import annotations
+
+import collections
+import hashlib
+import json
+import os
+import queue
+import threading
+import time
+from typing import Any
+
+from deepinfer.logger import init_logger
+
+logger = init_logger(__name__)
+
+ENV = "DEEPINFER_HICACHE_MOONCAKE_STORE"
+FORMAT_VERSION = 1
+OBJECT_NOT_FOUND, OBJECT_ALREADY_EXISTS = -704, -705
+TCP_ENV_DEFAULTS = {"MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER": "65535", "MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER": "65535"}
+
+
+def config_from_env() -> dict | None:
+    raw = os.environ.get(ENV, "").strip()
+    if not raw:
+        return None
+    if raw.startswith("@"):
+        with open(raw[1:]) as f:
+            raw = f.read()
+    return json.loads(raw)
+
+
+def layout_fingerprint(parts: dict[str, Any]) -> str:
+    """Version byte + short hash of everything that determines the bytes of a page / state."""
+    blob = json.dumps(parts, sort_keys=True).encode()
+    return f"v{FORMAT_VERSION}-{hashlib.sha256(blob).hexdigest()[:16]}"
+
+
+class _Breaker:
+    """Opens on a transfer failure; closes after a cooldown that doubles on consecutive trips."""
+
+    def __init__(self, cooldown_s: float, max_cooldown_s: float = 600.0) -> None:
+        self.base, self.max = cooldown_s, max_cooldown_s
+        self.cooldown = cooldown_s
+        self.open_until = 0.0
+        self.trips = 0
+        self._lock = threading.Lock()
+
+    def allow(self) -> bool:
+        return time.monotonic() >= self.open_until
+
+    def trip(self, why: str) -> None:
+        with self._lock:
+            now = time.monotonic()
+            if now < self.open_until:  # already open (e.g. the put thread and a get failed together)
+                return
+            if now - self.open_until > self.cooldown:  # healthy for a while since the last trip: reset the backoff
+                self.cooldown = self.base
+            self.open_until = now + self.cooldown
+            self.trips += 1
+            logger.warning("HostPrefixCache L3: %s; Store disabled for %.0f s", why, self.cooldown)
+            self.cooldown = min(self.cooldown * 2, self.max)
+
+
+class MooncakeL3:
+    def __init__(self, config: dict, fingerprint: str, group_pages: int) -> None:
+        for k, v in TCP_ENV_DEFAULTS.items():  # read by Mooncake's TCP transport when setup() creates it
+            os.environ.setdefault(k, v)
+        from mooncake.store import MooncakeDistributedStore  # optional dependency (pip install mooncake-transfer-engine)
+
+        cfg = dict(config)
+        cfg.setdefault("global_segment_size", 0)  # memory is contributed by a standalone mooncake_client by default
+        self.store = MooncakeDistributedStore()
+        rc = self.store.setup(cfg)
+        if rc != 0:
+            raise RuntimeError(f"MooncakeDistributedStore.setup failed: {rc}")
+        self.fingerprint = fingerprint
+        self.group_pages = group_pages
+        self._page_prefix = f"dsv4kv/{fingerprint}/g{group_pages}/"
+        self._state_prefix = f"state/{fingerprint}/"
+        self._puts: queue.Queue = queue.Queue()
+        self._done: collections.deque = collections.deque()  # tokens of finished puts (appended by the worker thread)
+        self.put_objects = self.put_failed = self.put_bytes = self.put_dropped = self.put_retried = self.get_retried = 0
+        self.put_seconds = 0.0
+        self.slow_get_s = float(os.environ.get("DEEPINFER_HICACHE_L3_SLOW_S", "5"))
+        self.breaker = _Breaker(float(os.environ.get("DEEPINFER_HICACHE_L3_COOLDOWN_S", "30")))
+        threading.Thread(target=self._put_worker, name="hicache-l3-put", daemon=True).start()
+        logger.info("HostPrefixCache L3: Mooncake Store ready (fingerprint %s, %d pages per object)", fingerprint, group_pages)
+
+    def register(self, ptr: int, nbytes: int) -> None:
+        rc = self.store.register_buffer(ptr, nbytes)
+        if rc != 0:
+            raise RuntimeError(f"Mooncake register_buffer({nbytes} bytes) failed: {rc}")
+
+    def group_key(self, last_page_hash: bytes) -> str:
+        return self._page_prefix + last_page_hash.hex()
+
+    def state_key(self, boundary_hash: bytes) -> str:
+        return self._state_prefix + boundary_hash.hex()
+
+    # ---------------------------------------------------------------- write-behind (any thread may submit)
+
+    def submit_put(self, keys: list[str], ptrs: list[list[int]], sizes: list[list[int]], token: object) -> None:
+        """Queue a batch put; `token` comes back from drain_done() when the put has returned (success or not)."""
+        self._puts.put((keys, ptrs, sizes, token))
+
+    def drain_done(self) -> list[object]:
+        out = []
+        while self._done:
+            out.append(self._done.popleft())
+        return out
+
+    def _put_worker(self) -> None:
+        while True:
+            keys, ptrs, sizes, token = self._puts.get()
+            if not self.breaker.allow():  # Store unhealthy: drop the put, release the host entries right away
+                self.put_dropped += len(keys)
+                self._done.append(token)
+                continue
+            t0 = time.perf_counter()
+            try:
+                rcs = list(self.store.batch_put_from_multi_buffers(keys, ptrs, sizes))
+                # retry the failed objects once: sporadic transfer errors are not a sick Store
+                redo = [i for i, rc in enumerate(rcs) if rc not in (0, OBJECT_ALREADY_EXISTS)]
+                if redo:
+                    self.put_retried += len(redo)
+                    again = self.store.batch_put_from_multi_buffers(
+                        [keys[i] for i in redo], [ptrs[i] for i in redo], [sizes[i] for i in redo]
+                    )
+                    for i, rc in zip(redo, again, strict=True):
+                        rcs[i] = rc
+                # an object that already exists (another rank or instance stored it first) is not a failure
+                bad = [rc for rc in rcs if rc not in (0, OBJECT_ALREADY_EXISTS)]
+                self.put_objects += len(keys) - len(bad)
+                self.put_failed += len(bad)
+                self.put_bytes += sum(sum(s) for s, rc in zip(sizes, rcs, strict=True) if rc == 0)
+                if bad:
+                    self.breaker.trip(f"batch put failed for {len(bad)}/{len(keys)} objects (rc {bad[0]})")
+            except Exception:  # noqa: BLE001 - a Store problem must not take the engine down
+                logger.exception("HostPrefixCache L3: batch put failed")
+                self.put_failed += len(keys)
+                self.breaker.trip("batch put raised")
+            self.put_seconds += time.perf_counter() - t0
+            self._done.append(token)
+
+    # ---------------------------------------------------------------- read (engine thread)
+
+    def readable(self) -> bool:
+        """False while the breaker is open: skip the Store and use the host cache only."""
+        return self.breaker.allow()
+
+    def exists(self, keys: list[str]) -> list[bool]:
+        if not keys:
+            return []
+        return [rc == 1 for rc in self.store.batch_is_exist(keys)]
+
+    def get(self, keys: list[str], ptrs: list[list[int]], sizes: list[list[int]]) -> list[bool]:
+        """Read each object into its list of buffers; True where the whole object arrived."""
+        if not keys:
+            return []
+        t0 = time.perf_counter()
+        rcs = list(self.store.batch_get_into_multi_buffers(keys, ptrs, sizes))
+        redo = [i for i, (rc, s) in enumerate(zip(rcs, sizes, strict=True)) if rc != sum(s) and rc != OBJECT_NOT_FOUND]
+        if redo:  # one retry, as for puts
+            self.get_retried += len(redo)
+            again = self.store.batch_get_into_multi_buffers([keys[i] for i in redo], [ptrs[i] for i in redo], [sizes[i] for i in redo])
+            for i, rc in zip(redo, again, strict=True):
+                rcs[i] = rc
+        dt = time.perf_counter() - t0
+        ok = [rc == sum(s) for rc, s in zip(rcs, sizes, strict=True)]
+        # "not found" just means evicted since exists(); anything else is a transfer problem
+        bad = [rc for rc, good in zip(rcs, ok, strict=True) if not good and rc != OBJECT_NOT_FOUND]
+        if bad:
+            self.breaker.trip(f"batch get failed for {len(bad)}/{len(keys)} objects (rc {bad[0]}, {dt:.1f} s)")
+        elif dt > self.slow_get_s:
+            self.breaker.trip(f"batch get of {len(keys)} objects took {dt:.1f} s")
+        return ok
+
+    def stats_dict(self) -> dict[str, int | float]:
+        return dict(
+            l3_put_objects=self.put_objects,
+            l3_put_failed=self.put_failed,
+            l3_put_gib=round(self.put_bytes / 2**30, 3),
+            l3_put_seconds=round(self.put_seconds, 3),
+            l3_put_queue=self._puts.qsize(),
+            l3_put_dropped=self.put_dropped,
+            l3_put_retried=self.put_retried,
+            l3_get_retried=self.get_retried,
+            l3_breaker_trips=self.breaker.trips,
+            l3_breaker_open=int(not self.breaker.allow()),
+        )

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

tests/hicache/test_host_cache_l3.py

@@ -0,0 +1,193 @@
+"""HostPrefixCache L3 (Mooncake Store) logic against a dict-backed fake store: write-behind of page groups and
+states, a second cache (another rank / instance) extending its empty host pool from L3, no re-puts, a fetch that
+loses an object between exists() and get(), and the circuit breaker after a transfer failure. Same fake engine as
+test_host_cache.py. Needs one GPU."""
+
+from __future__ import annotations
+
+import ctypes
+import time
+
+import numpy as np
+import pytest
+import torch
+
+from deepinfer.hicache.mooncake_store import _Breaker
+from tests.hicache.test_host_cache import BLOCK, PAGE, STATE, FakeSeq, _prefill, _tokens, engine  # noqa: F401
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+GROUP = 4
+
+
+class FakeL3:
+    """In-memory stand-in for MooncakeL3 with the same interface; puts complete synchronously."""
+
+    def __init__(self) -> None:
+        self.group_pages = GROUP
+        self.objects: dict[str, bytes] = {}
+        self.puts = 0
+        self.drop_on_get: set[str] = set()  # keys that vanish between exists() and get()
+        self.fail_gets = False  # every get is a transfer failure (trips the breaker, like MooncakeL3.get)
+        self.gets = 0
+        self.breaker = _Breaker(30.0)
+        self._done: list[object] = []
+
+    def readable(self) -> bool:
+        return self.breaker.allow()
+
+    def register(self, ptr: int, nbytes: int) -> None:
+        pass
+
+    def group_key(self, h: bytes) -> str:
+        return "dsv4kv/test/g4/" + h.hex()
+
+    def state_key(self, h: bytes) -> str:
+        return "state/test/" + h.hex()
+
+    def submit_put(self, keys, ptrs, sizes, token) -> None:  # noqa: ANN001
+        for k, ps, ss in zip(keys, ptrs, sizes, strict=True):
+            if k not in self.objects:
+                self.objects[k] = b"".join(ctypes.string_at(p, s) for p, s in zip(ps, ss, strict=True))
+                self.puts += 1
+        self._done.append(token)
+
+    def drain_done(self) -> list[object]:
+        out, self._done = self._done, []
+        return out
+
+    def exists(self, keys: list[str]) -> list[bool]:
+        return [k in self.objects for k in keys]
+
+    def get(self, keys, ptrs, sizes) -> list[bool]:  # noqa: ANN001
+        self.gets += 1
+        if self.fail_gets:
+            self.breaker.trip("test transfer failure")
+            return [False] * len(keys)
+        ok = []
+        for k, ps, ss in zip(keys, ptrs, sizes, strict=True):
+            data = None if k in self.drop_on_get else self.objects.get(k)
+            if data is None or len(data) != sum(ss):
+                ok.append(False)
+                continue
+            off = 0
+            for p, s in zip(ps, ss, strict=True):
+                ctypes.memmove(p, data[off : off + s], s)
+                off += s
+            ok.append(True)
+        return ok
+
+    def stats_dict(self) -> dict:
+        return {"l3_put_objects": self.puts}
+
+
+@pytest.fixture
+def l3(monkeypatch):
+    from deepinfer.hicache import HostPrefixCache
+
+    store = FakeL3()
+    monkeypatch.setattr(HostPrefixCache, "_make_l3", lambda self: store)
+    return store
+
+
+def _settle(cache) -> None:  # noqa: ANN001
+    torch.cuda.synchronize()
+    cache.poll()  # retires D2H copies and submits the planned L3 puts
+    cache.poll()  # drains the put completions (refs back to 0)
+
+
+def test_second_cache_hits_from_l3_with_exact_bytes(engine, l3):  # noqa: F811
+    from deepinfer.hicache import HostPrefixCache
+
+    blocks = engine.driver_worker.cache_engine.gpu_blocks
+    states = engine.driver_worker.cache_engine.gpu_states
+    a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    a = FakeSeq("a", _tokens(1000, 1))
+    _prefill(engine, a_cache, a)
+    _settle(a_cache)
+    # 7 pages -> one full group (pages 0..3) + states at 896/768/640/512
+    assert len([k for k in l3.objects if k.startswith("dsv4kv/")]) == 1
+    assert len([k for k in l3.objects if k.startswith("state/")]) == 4
+    assert all(e.refs == 0 for e in list(a_cache.pages.values()) + list(a_cache.states.values()))
+    a_pages = [blocks[b].clone() for b in a.block_table[:4]]
+    a_state = states[a.state_id, :STATE].clone()
+
+    # Another rank / instance: empty host pool, same prefix + a new turn. Host alone would miss; L3 gives the group
+    # and the state at 512 (the only boundary inside the stored group).
+    b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
+    assert b_cache.try_resume(b, "b", b.tokens) == 4 * BLOCK
+    torch.cuda.synchronize()
+    b_cache.poll()
+    assert all(torch.equal(blocks[b.block_table[i]], a_pages[i]) for i in range(4))
+    assert torch.equal(states[b.state_id, :STATE], a_state)
+    s = b_cache.stats
+    assert s.l3_lookups == 1 and s.l3_hits == 1 and s.l3_pages_loaded == 4 and s.l3_states_loaded == 1
+    assert all(e.refs == 0 for e in list(b_cache.pages.values()) + list(b_cache.states.values()))
+    assert len(b_cache.free_pages) + len(b_cache.pages) == b_cache.num_pages
+
+
+def test_no_reput_of_known_objects(engine, l3):  # noqa: F811
+    from deepinfer.hicache import HostPrefixCache
+
+    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    a = FakeSeq("a", _tokens(1000, 1))
+    _prefill(engine, cache, a)
+    _settle(cache)
+    puts = l3.puts
+    again = FakeSeq("a-again", a.tokens.copy())
+    _prefill(engine, cache, again)  # identical prompt: nothing new on the host, nothing new for L3
+    _settle(cache)
+    assert l3.puts == puts
+
+
+def test_object_lost_between_exists_and_get_falls_back(engine, l3):  # noqa: F811
+    from deepinfer.hicache import HostPrefixCache
+
+    a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    a = FakeSeq("a", _tokens(1000, 1))
+    _prefill(engine, a_cache, a)
+    _settle(a_cache)
+    l3.drop_on_get = {k for k in l3.objects if k.startswith("dsv4kv/")}
+
+    b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
+    assert b_cache.try_resume(b, "b", b.tokens) == 0  # miss, sequence untouched
+    assert b.block_table == [] and b_cache.stats.l3_get_failed == 1
+    assert len(b_cache.pages) == 0 and len(b_cache.free_pages) == b_cache.num_pages
+
+
+def test_transfer_failure_opens_the_breaker(engine, l3):  # noqa: F811
+    from deepinfer.hicache import HostPrefixCache
+
+    a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    a = FakeSeq("a", _tokens(1000, 1))
+    _prefill(engine, a_cache, a)
+    _settle(a_cache)
+
+    l3.fail_gets = True
+    b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
+    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
+    assert b_cache.try_resume(b, "b", b.tokens) == 0 and l3.gets == 1
+    assert not l3.readable() and l3.breaker.trips == 1
+    # while open, later requests do not touch the Store at all (no 60 s stalls), even once it is healthy again
+    l3.fail_gets = False
+    c = FakeSeq("c", np.concatenate([a.tokens, _tokens(500, 3)]))
+    assert b_cache.try_resume(c, "c", c.tokens) == 0 and l3.gets == 1
+    # after the cooldown the Store is used again
+    l3.breaker.open_until = 0.0
+    d = FakeSeq("d", np.concatenate([a.tokens, _tokens(500, 4)]))
+    assert b_cache.try_resume(d, "d", d.tokens) == 4 * BLOCK and l3.gets == 2
+
+
+def test_breaker_backoff() -> None:
+    br = _Breaker(10.0, max_cooldown_s=25.0)
+    assert br.allow()
+    br.trip("x")
+    assert not br.allow() and br.trips == 1
+    br.trip("y")  # already open: no second trip
+    assert br.trips == 1
+    for want in (20.0, 25.0):  # consecutive trips right after reopening double the cooldown, capped
+        br.open_until = time.monotonic() - 1e-3
+        br.trip("z")
+        assert abs(br.open_until - time.monotonic() - want) < 1.0