Merge d54fe25993a44a7f09f5bbe6420fec6438c1ea79 into 553741557a57df44a46fb3ab72b6aafeaf945337

csrc/engine/pybind_scheduler.cpp

@@ -91,6 +91,10 @@ void pybind_engine_scheduler(pybind11::module &engine) {
             },
             py::arg("seq"))
         .def("add_sequence", &Scheduler::add_sequence, py::arg("seq"))
+        .def("free_finished_seqs", &Scheduler::free_finished_seqs,
+            "把 running 中已结束的 seq 移出并返回（不释放 block/state，调用方按需 free_seq）；PROMPT 引擎送走 prefill 结果后用")
+        .def("add_prompted_sequence", &Scheduler::add_prompted_sequence, py::arg("seq"),
+            "PD 分离：把 KV/state 已填好的 seq 放入 prompted 队列，由 COMPLETION 引擎直接续写")
         .def("abort_sequence", &Scheduler::abort_sequence, py::arg("request_ids"))
         .def("has_unfinished_seqs", &Scheduler::has_unfinished_seqs)
         .def("get_num_unfinished_seqs", &Scheduler::get_num_unfinished_seqs)

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; });

csrc/engine/scheduler.hpp

@@ -1058,6 +1058,19 @@ struct Scheduler {
         }
     }
 
+    /// PD 分离：把在 prefill 引擎上算完 prefill、KV block 与 state 已由调用方填好的 seq 直接放入 prompted
+    /// 队列（不 clone）。token 为 prompt + 首个生成 token，其中除最后一个 token 外 KV 均已就绪。
+    void add_prompted_sequence(std::shared_ptr<Sequence> seq) {
+        TORCH_CHECK(seq->schedule.status == SequenceStatus::WAITING, "status ", int(seq->schedule.status));
+        TORCH_CHECK(!seq->cache.block_table.empty(), "prompted sequence needs allocated KV blocks");
+        TORCH_CHECK(seq->cache.num_tokens == seq->get_len() - 1, "prompted sequence must have all but its last token computed: ",
+            seq->cache.num_tokens, " vs len ", seq->get_len());
+        seq->request.arrival_time = get_timestamp_ms();
+        seq->reserve_output_tokens();
+        seq->cache.block_table.reserve(_block_table_capacity(*seq));
+        prompted.emplace_back(std::move(seq), false);
+    }
+
     SequencePtr fork_seq(const SequencePtr &seq, py::object blocks, std::string &&target_branch_id, py::object sampling_params,
         py::object so_matcher, int sp_token, int replace_last_token_id, int64_t seq_id, double arrival_time) {
         TORCH_CHECK(seq->cache.num_tokens > 0, "Cannot fork a sequence without a computed prefix");

deepinfer/_C/engine.pyi

@@ -313,6 +313,10 @@ class Scheduler:
         ...
     def add_sequence(self, seq: Sequence) -> None:
         ...
+    def add_prompted_sequence(self, seq: Sequence) -> None:
+        ...
+    def free_finished_seqs(self) -> list[Sequence]:
+        ...
     def fork_seq(self, seq: Sequence, blocks: typing.Any, target_branch_id: str, sampling_params: typing.Any, so_matcher: typing.Any, sp_token: typing.SupportsInt | typing.SupportsIndex, replace_last_token_id: typing.SupportsInt | typing.SupportsIndex, seq_id: typing.SupportsInt | typing.SupportsIndex, arrival_time: typing.SupportsFloat | typing.SupportsIndex) -> Sequence:
         ...
     def free_seq(self, seq: Sequence) -> None:
@@ -455,6 +459,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,437 @@
+"""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
+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)
+
+
+@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: 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
+        # 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: 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)
+        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))
+        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,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/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();
+}

deepinfer/hicache/pd_mooncake.py

@@ -0,0 +1,106 @@
+"""Mooncake Transfer Engine data plane for the prefill/decode split (serve/pd_server.py --kv-transport mooncake).
+
+Each engine rank registers one GPU arena with a Mooncake TransferEngine. The prefill engine exports a sequence
+(KV pages + full "to CE" state, PDTransfer.export_into) into a region of its arena and moves it with one
+contiguous write into a region the decode rank reserved in its own arena; the decode rank admits from there with
+the same page-copy kernels as the TCP path. Control messages (header, reservation, completion) stay on the TCP
+socket of the TCP path. Within a node this is Mooncake's nvlink_intra transport (CUDA IPC + copy engines); across
+hosts the same code runs over RDMA / EFA.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+
+import torch
+
+ALIGN = 4096
+
+
+def _align(n: int, a: int = ALIGN) -> int:
+    return (n + a - 1) // a * a
+
+
+class GpuArena:
+    """One device buffer with a first-fit allocator (thread-safe). Regions are ALIGN-aligned."""
+
+    def __init__(self, nbytes: int, device: torch.device) -> None:
+        self.nbytes = _align(nbytes)
+        # Not torch.empty: the worker runs PyTorch with expandable_segments (worker/base.py), whose VMM memory has no
+        # cudaIpc handle, and Mooncake's nvlink_intra transport registers buffers with cudaIpcGetMemHandle. Use the
+        # same legacy cuMemAlloc path as the KV pools (CacheEngine._alloc_gpu_uint8).
+        try:
+            from deepinfer._C import cuda_utils  # pyright: ignore[reportMissingModuleSource]
+
+            raw = cuda_utils.alloc_tensor(self.nbytes)
+            if raw is None:
+                raise MemoryError(f"cuMemAlloc of a {self.nbytes >> 20} MiB PD arena failed (lower --pd-arena-gb)")
+            self.buf = raw.view(torch.uint8).reshape(-1)
+        except ImportError:  # outside a DeepInfer build (unit tests)
+            self.buf = torch.empty(self.nbytes, dtype=torch.uint8, device=device)
+        assert self.buf.numel() == self.nbytes, (self.buf.numel(), self.nbytes)
+        self.base: int = self.buf.data_ptr()
+        self._free: list[tuple[int, int]] = [(0, self.nbytes)]  # sorted, coalesced (offset, size)
+        self._lock = threading.Lock()
+
+    def alloc(self, n: int) -> int | None:
+        n = _align(n)
+        with self._lock:
+            for i, (off, size) in enumerate(self._free):
+                if size >= n:
+                    if size == n:
+                        del self._free[i]
+                    else:
+                        self._free[i] = (off + n, size - n)
+                    return off
+        return None
+
+    def release(self, off: int, n: int) -> None:
+        n = _align(n)
+        with self._lock:
+            free = self._free
+            i = 0
+            while i < len(free) and free[i][0] < off:
+                i += 1
+            free.insert(i, (off, n))
+            if i + 1 < len(free) and off + n == free[i + 1][0]:  # merge with the next
+                free[i] = (off, n + free[i + 1][1])
+                del free[i + 1]
+            if i > 0 and free[i - 1][0] + free[i - 1][1] == off:  # merge with the previous
+                free[i - 1] = (free[i - 1][0], free[i - 1][1] + free[i][1])
+                del free[i]
+
+    def view(self, off: int, n: int) -> torch.Tensor:
+        return self.buf[off : off + n]
+
+
+def region_layout(num_pages: int, page_bytes: int, state_bytes: int) -> tuple[int, int]:
+    """(state offset, total bytes) of a transfer region: pages first, then the state."""
+    state_off = _align(num_pages * page_bytes, 256)
+    return state_off, state_off + state_bytes
+
+
+class MooncakeLink:
+    """A TransferEngine bound to one arena. hostname is the address peers use to reach this engine's handshake
+    server (P2PHANDSHAKE metadata: no etcd); the resulting segment name is hostname:rpc_port."""
+
+    def __init__(self, hostname: str, protocol: str, arena: GpuArena, device_name: str = "") -> None:
+        if protocol == "nvlink_intra":
+            os.environ.setdefault("MC_INTRANODE_NVLINK", "true")
+        from mooncake.engine import TransferEngine  # optional dependency: pip install mooncake-transfer-engine[-cuda13]
+
+        self.te = TransferEngine()
+        rc = self.te.initialize(hostname, "P2PHANDSHAKE", protocol, device_name)
+        if rc != 0:
+            raise RuntimeError(f"Mooncake TransferEngine.initialize({hostname!r}, {protocol!r}) failed: {rc}")
+        self.segment = f"{hostname}:{self.te.get_rpc_port()}"
+        rc = self.te.register_memory(arena.base, arena.nbytes)
+        if rc != 0:
+            raise RuntimeError(f"Mooncake register_memory({arena.nbytes} bytes) failed: {rc}")
+        self.arena = arena
+
+    def write(self, segment: str, src_off: int, dst_addr: int, n: int) -> None:
+        rc = self.te.transfer_sync_write(segment, self.arena.base + src_off, dst_addr, n)
+        if rc != 0:
+            raise RuntimeError(f"Mooncake transfer_sync_write to {segment} ({n} bytes) failed: {rc}")

deepinfer/hicache/pd_transfer.py

@@ -0,0 +1,141 @@
+"""Prefill/decode split: move a prefilled sequence's KV blocks and SWA state between engines.
+
+The prefill engine (EngineType.PROMPT) exports, after the last prefill chunk, the sequence's full KV blocks
+(all tokens of the prompt) and the complete SWA/compression state in the "to CE" layout
+(GPUCacheInterface.select_state_slice_to_store(..., TO_CE=True)), both into pinned host buffers through the
+page-copy kernel. The decode engine (EngineType.COMPLETION) receives those bytes into pinned buffers, allocates
+blocks and a state slot for a new sequence (prompt + first sampled token), fills them with the kernel, and admits
+the sequence into its `prompted` queue (Scheduler.add_prompted_sequence). Transport between the two is left to
+the caller (plain TCP in serve/pd_server.py).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from deepinfer import envs
+from deepinfer.hicache.kernels import copy_pages_module
+from deepinfer.worker.cache_engine import GPUCacheInterface
+
+if TYPE_CHECKING:
+    from deepinfer._C.engine import Sequence
+    from deepinfer.engine.llm_engine import LLMEngine
+
+
+def load_state_from_ce_transfer(state_id: int, last_position: int, state_tensor: torch.Tensor, state_cache: torch.Tensor) -> None:
+    """Inverse of select_state_slice_to_store(..., TO_CE=True): write a full inference-layout state (stored SWA,
+    dropped/consumer SWA and compression state) into state_cache[state_id] with the ring ending at last_position."""
+    from deepinfer.model_executor.layers.tilelang_ops.state_cache import load_state_to_cache
+
+    cls = GPUCacheInterface
+    cache_views, tensor_views = cls._carve_state_views(state_id, state_cache, state_tensor, TO_CE=True)
+    for tensor_view, cache_view in zip(tensor_views, cache_views, strict=True):
+        if tensor_view.numel():
+            load_state_to_cache(tensor_view, cache_view, last_position)
+    dropped_tensor_offset = cls.swa_total_bytes_inference_pad
+    dropped_cache_offset = cls.swa_total_bytes_kvstore_pad
+    compress_bytes = cls.dropped_bytes_inference - cls.swa_drop_total_bytes_inference_pad
+    if compress_bytes:
+        state_cache[state_id, dropped_cache_offset : dropped_cache_offset + compress_bytes].copy_(
+            state_tensor[dropped_tensor_offset : dropped_tensor_offset + compress_bytes]
+        )
+
+
+class PDTransfer:
+    """Per-engine helper owning the GPU pools' pointers and small pinned staging for transfers."""
+
+    def __init__(self, engine: LLMEngine) -> None:
+        cache_engine = engine.driver_worker.cache_engine
+        self.engine = engine
+        self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks
+        self.gpu_states: torch.Tensor = cache_engine.gpu_states
+        self.device = self.gpu_blocks.device
+        self.block_size: int = engine.cache_config.block_size
+        self.page_bytes: int = self.gpu_blocks.size(1)
+        self.state_bytes: int = GPUCacheInterface.total_bytes_inference  # full "to CE" state
+        assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0
+        self.kernel = copy_pages_module()
+
+    def _ids(self, ids: list[int]) -> torch.Tensor:
+        return torch.tensor(ids, dtype=torch.int32).pin_memory().to(self.device, non_blocking=True)
+
+    # ---------------------------------------------------------------- prefill side
+
+    def export(self, seq: Sequence, prompt_len: int | None = None) -> tuple[torch.Tensor, torch.Tensor, torch.cuda.Event]:
+        """Copy seq's KV blocks for its prompt and the full state to pinned host buffers.
+
+        prompt_len: tokens with KV (default: seq.get_len() - 1, i.e. after the first output token was appended;
+        pass seq.get_len() when called before outputs are processed). KV exists for the prompt only, so a prompt
+        ending exactly on a block boundary must not count an extra block.
+        Returns (pages[num_blocks, page_bytes] pinned, state[state_bytes] pinned, event); the buffers are valid once
+        the event has completed. Enqueued on the current stream, so it is ordered after the prefill's kernels.
+        """
+        if prompt_len is None:
+            prompt_len = seq.get_len() - 1
+        pages = torch.empty((self.num_blocks(prompt_len), self.page_bytes), dtype=torch.uint8, pin_memory=True)
+        state = torch.empty(self.state_bytes, dtype=torch.uint8, pin_memory=True)
+        return pages, state, self.export_into(seq, prompt_len, pages, state)
+
+    def num_blocks(self, prompt_len: int) -> int:
+        return (prompt_len + self.block_size - 1) // self.block_size
+
+    def export_into(self, seq: Sequence, prompt_len: int, pages: torch.Tensor, state: torch.Tensor) -> torch.cuda.Event:
+        """export() into caller-provided buffers: pages [num_blocks(prompt_len), page_bytes] and state [state_bytes],
+        pinned host or device memory (e.g. a region of a Mooncake-registered arena, pd_mooncake.GpuArena)."""
+        num_blocks = pages.size(0)
+        block_table = list(seq.block_table)[:num_blocks]
+        self.kernel.copy_pages(
+            self.gpu_blocks.data_ptr(), pages.data_ptr(), self._ids(block_table), self._ids(list(range(num_blocks))), self.page_bytes
+        )
+        ev = torch.cuda.Event()
+        if state.is_cuda:
+            GPUCacheInterface.select_state_slice_to_store(seq.state_id, prompt_len, self.gpu_states, state_tensor=state, TO_CE=True)
+            ev.record()
+            return ev
+        staging = torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device)
+        GPUCacheInterface.select_state_slice_to_store(seq.state_id, prompt_len, self.gpu_states, state_tensor=staging, TO_CE=True)
+        self.kernel.copy_pages(staging.data_ptr(), state.data_ptr(), self._ids([0]), self._ids([0]), self.state_bytes)
+        ev.record()
+        # keep the staging alive until the copy ran (same-stream ordering makes later reuse safe)
+        ev.staging = staging  # type: ignore[attr-defined]
+        return ev
+
+    # ---------------------------------------------------------------- decode side
+
+    def admit(self, seq: Sequence, pages: torch.Tensor, state: torch.Tensor, prompt_len: int) -> torch.cuda.Event:
+        """Fill a freshly created sequence (prompt + first token, created with add_request(enqueue=False)) from
+        pinned host or device buffers and put it in the scheduler's prompted queue. Caller must have called
+        try_allocate_seq, and must keep `pages`/`state` alive until the returned event has completed.
+
+        prompt_len: tokens with KV in `pages` (= seq.get_len() - 1); the state ring ends there.
+        """
+        num_blocks = pages.size(0)
+        block_table = list(seq.block_table)
+        assert len(block_table) >= num_blocks, (len(block_table), num_blocks)
+        self.kernel.copy_pages(
+            pages.data_ptr(),
+            self.gpu_blocks.data_ptr(),
+            self._ids(list(range(num_blocks))),
+            self._ids(block_table[:num_blocks]),
+            self.page_bytes,
+        )
+        if state.is_cuda:
+            staging = state
+        else:
+            staging = torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device)
+            self.kernel.copy_pages(state.data_ptr(), staging.data_ptr(), self._ids([0]), self._ids([0]), self.state_bytes)
+        load_state_from_ce_transfer(seq.state_id, prompt_len, staging, self.gpu_states)
+        seq.num_cached_tokens = prompt_len
+        seq.cached_hit_tokens = prompt_len
+        # A finished prefill leaves `nextn` draft tokens on the sequence (pad tokens when the prompt step produced
+        # none, as a prompt-only engine does); the decode sampler requires exactly that many.
+        nextn = int(envs.DEEPINFER_USE_NEXTN_DECODE)
+        if nextn > 0:
+            seq.last_next_n_token = self.engine.get_nextn_pad_token_ids().tolist()[:nextn]
+        self.engine.scheduler.add_prompted_sequence(seq)
+        ev = torch.cuda.Event()
+        ev.record()
+        ev.staging = staging  # type: ignore[attr-defined]
+        return ev

serve/README.md

@@ -0,0 +1,134 @@
+# Serving DeepInfer + host prefix cache (local additions, not part of upstream)
+
+## Pieces
+
+- `serve/openai_server.py` — OpenAI-compatible `/v1/chat/completions` (streaming, `tools`, usage) on top of the
+  offline engine: one engine process per GPU (attention DP + MoE EP), prompts rendered with the checkpoint's own
+  `encoding.py`. `serve/server.sh start|stop` runs it as a process group (stop reaps the engine helpers, which
+  otherwise keep the GPUs).
+- `deepinfer/hicache/` — host-memory prefix cache (SGLang HiCache style): pinned host pools organized in pages
+  (one page = one 128-token KV block) plus SWA state slots, hash-chain index, LRU eviction, write-through at the
+  end of every prefill, "kernel" transfers (a CUDA kernel copying pages through UVA pointers on the engine
+  stream, `page_copy.cu`, JIT-built on first use). Needs the read-only `Sequence.block_table` binding added in
+  `csrc/engine/pybind_sequence.cpp` (rebuild `deepinfer._C` with `bash scripts/build.sh build`).
+- `tests/hicache/test_host_cache.py` — cache logic against a fake engine (needs one GPU).
+- `serve/check_prefix_cache.py` — A/B on a live server: each turn sent with `"prefix_cache": false` and with
+  the cache, greedy; compares outputs and TTFT.
+
+## Run
+
+```bash
+export DEEPINFER_ACCELERATOR=nvidia DEEPINFER_HF_MODEL_DIR=/path/to/DeepSeek-V4.1-Flash-1017
+serve/server.sh start --dp 4 --port 8888                       # cache on: 32 GiB pages + 4096 states per rank
+serve/server.sh start --dp 4 --port 8888 --hicache-size-gb 0   # cache off
+python serve/check_prefix_cache.py --port 8888 --turns 3
+curl -s localhost:8888/deepinfer/metrics | grep hicache         # per-rank hit/store counters
+serve/server.sh stop
+```
+
+Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-loaded` (affinity keeps a
+conversation's turns on one rank, keyed by the first user message; required for hits since the cache is
+per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, the default 140).
+
+EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearrangement stalls all ranks for about
+0.3 s, and with `--eplb-after-prefill 1` it fires after nearly every step under multi-turn traffic; enable it with
+`--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. With EPLB off the server also drops
+the 32 redundant expert slots (`DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=0`): without rearrangement they are fixed copies
+of experts 0-31, and removing them speeds decode up by 4-5%.
+
+`--prefill-any-interval` (default 30) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and
+a rank with a waiting prompt forces one at most this many decode steps later. The default favours throughput. With 2,
+almost every arriving request triggers a prompt step, which keeps TTFT low but slows decode at high concurrency. The
+engine default of 100 adds a ~1.4 s TTFT floor. AgentX at c256 on one GB300 tray (20-minute runs):
+
+| interval | tok/s/GPU | TTFT p50 / p90 | ITL p50 / p90 | P90 interactivity |
+|---|---|---|---|---|
+| 2 | 258K | 1.26 / 4.53 s | 14.3 / 31.5 ms | 32 |
+| 30 | 283K | 1.62 / 6.51 s | 10.0 / 14.1 ms | 71 |
+| 100 | 276K | 3.00 / 14.95 s | 8.1 / 13.5 ms | 74 |
+
+B200 shows the same trend: at c256, interval 100 gave 255K vs 211K tok/s/GPU for interval 2.
+
+## Prompt tokenization
+
+Agentic clients resend the whole conversation every turn, and `tokenizer.encode` needs ~1.7 ms per 1K tokens on a
+GB300 Grace core (p50 278 ms per request on replayed AgentX prompts of p50 163K tokens), which dominates TTFT once the
+prefix cache hits. `serve/prompt_tokenizer.py` gives the same ids faster; `--tokenizer-backend` picks one:
+
+- `fastokens` (`pip install fastokens`; x86-64 and aarch64 wheels): Rust reimplementation of the HF pipeline, built
+  from the loaded tokenizer. 3.5 ms p50 on the same prompts.
+- `segment` (no dependency): splits the prompt at added tokens, where HF tokenizers already encode the pieces
+  independently, and caches each piece's ids (LRU, 512 MiB), so a turn only encodes what it adds. 18 ms p50.
+  Enabled only when the tokenizer config makes the split exact (no `single_word` tokens, no `prepend_scheme="first"`).
+- `hf`: `tokenizer.encode`.
+
+`auto` (default) takes the first of these that matches `tokenizer.encode` on startup probes built from the
+tokenizer's added tokens. `--tokenizer-verify-rate` (default 0.005) re-encodes a sample of prompts with
+`tokenizer.encode` in the background; a mismatch switches to `hf`. The backend, encode latency and check counters are
+in `/deepinfer/metrics` (`deepinfer_tokenizer:*`). Tests: `tests/serve/test_prompt_tokenizer.py`.
+
+## How a hit works
+
+A new prompt's page hashes are matched from the start; the match is capped so that at least `keep_tokens`
+tokens remain to be recomputed (the decoder SWA state is rebuilt from that suffix by bounded replay), then
+snapped down to the nearest 128-aligned position that has a stored SWA state. The sequence is allocated up
+front, pages and state are copied in on the engine stream, `num_cached_tokens` is set, and the sequence is
+enqueued; the scheduler prefills only the suffix. Generated tokens are not stored at decode end (they are
+stored when they come back as part of the next turn's prompt).
+
+## AgentX
+
+The server speaks what InferenceX's AgentX replay (aiperf, `--endpoint /v1/chat/completions --streaming
+--use-server-token-count`) needs. At low concurrency, use `--prefill-any-interval 2`: a single lane keeps most
+ranks idle, and the default adds a few hundred ms of TTFT.
+
+## Prefill/decode split (`pd_server.py`)
+
+`serve/pd_server.py` runs the same OpenAI-compatible frontend over two engines on two hosts: a prompt-only
+engine (PE, `--role pe`, host prefix cache on) and a decode-only engine (CE, `--role ce`, user-facing).
+Per request the CE frontend renders the prompt, picks a CE rank and POSTs `/prefill` to the PE; the PE
+prefills, exports the sequence's KV blocks and SWA state (`deepinfer/hicache/pd_transfer.py`) plus the first
+sampled token, pushes the bytes to the CE rank over TCP and the CE admits the sequence straight into decode
+(`Scheduler.add_prompted_sequence`).
+
+```bash
+serve/pd_run.sh start <pe-host> <ce-host>          # PE HTTP 8889, CE HTTP 8888, CE KV receivers 8900+rank
+python serve/check_pd.py --ref http://<single-engine-host>:8888 --pd http://<ce-host>:8888   # token-identical?
+serve/pd_run.sh stop  <pe-host> <ce-host>
+```
+
+`serve/nccl_xtray_test.py` reports the transport and bandwidth NCCL picks between the two hosts.
+Results on 4+4 GB300 vs a single 4-GPU tray: `doc/benchmark_gb300.md`.
+
+### KV transport: Mooncake Transfer Engine (`--kv-transport mooncake`)
+
+`--kv-transport mooncake` (both roles) moves the KV blocks and SWA state with a
+[Mooncake](https://github.com/kvcache-ai/Mooncake) `TransferEngine` instead of pinned host buffers over TCP
+(`deepinfer/hicache/pd_mooncake.py`).
+- Each rank registers one GPU arena (`--pd-arena-gb`, default 4).
+- The PE exports a sequence into a region of its arena (`PDTransfer.export_into`) and writes it with one contiguous
+  `transfer_sync_write` into a region the CE rank reserved in its own arena.
+- The CE admits from there with the same page-copy kernels.
+- Control (header, reservation, completion) stays on the KV socket. Each write is retried, then falls back to the
+  TCP path, so a failing peer cannot strand a request.
+
+`--mooncake-protocol` picks the transport: `nvlink_intra` (default; PE and CE on one host, CUDA IPC + copy engines),
+or `rdma` / `efa` across hosts. Use `--mooncake-device` for the NICs and `--mooncake-host` for the address peers
+use for the handshake.
+
+```bash
+pip install mooncake-transfer-engine-cuda13   # or mooncake-transfer-engine for CUDA 12
+# one 8-GPU host: PE on GPUs 0-3, CE on GPUs 4-7
+python serve/pd_server.py --role pe --port 8889 --kv-transport mooncake --mooncake-host 127.0.0.1 ...
+python serve/pd_server.py --role ce --port 8888 --kv-port 8900 --kv-host 127.0.0.1 --pe-url http://127.0.0.1:8889 \
+    --kv-transport mooncake --mooncake-host 127.0.0.1 ...
+```
+
+Notes for `nvlink_intra`:
+- **Allocation.** The arena is allocated with `cuMemAlloc` (`cuda_utils.alloc_tensor`, like the KV pools), not
+  `torch.empty`. The worker runs PyTorch with `expandable_segments` (`deepinfer/worker/base.py`), and that VMM memory
+  has no cudaIpc handle.
+- **Containers.** If PE and CE run in separate containers, share the IPC **and** PID namespaces
+  (`docker run --ipc=host --pid=host`). A CUDA IPC handle carries the exporter's PID. With per-container PID
+  namespaces, the PE's and CE's rank-r workers get the same PID, and `cudaIpcOpenMemHandle` rejects the peer handle
+  as its own ("invalid device context"). Only the same-index rank pairs fail.

serve/check_pd.py

@@ -0,0 +1,81 @@
+"""Compare a PD (prefill/decode split) deployment against a single-engine server on the same prompts.
+
+    python serve/check_pd.py --ref http://localhost:8888 --pd http://tray10:8888 --turns 3 --prefix-tokens 3000
+
+Greedy, multi-turn with a long shared prefix (so the PE's prefix cache gets exercised); prints per-turn TTFT and
+whether the outputs are identical. Exits non-zero on any mismatch.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+
+import requests
+
+p = argparse.ArgumentParser()
+p.add_argument("--ref", default="http://localhost:8888")
+p.add_argument("--pd", required=True)
+p.add_argument("--model", default="deepseek-ai/DeepSeek-V4.1-Flash-1017")
+p.add_argument("--turns", type=int, default=3)
+p.add_argument("--prefix-tokens", type=int, default=3000)
+p.add_argument("--max-tokens", type=int, default=64)
+args = p.parse_args()
+
+words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa", "lambda", "mu"]
+doc = " ".join(f"{words[i % len(words)]}{i % 97}" for i in range(args.prefix_tokens))
+questions = [
+    "Which word appears right after 'gamma2'? Answer briefly.",
+    "Quote the first five words of the document.",
+    "What is the last word of the document?",
+    "How many distinct word stems are used? Answer briefly.",
+]
+
+
+def ask(url: str, messages: list[dict]) -> tuple[str, float, float, int]:
+    body = {
+        "model": args.model,
+        "messages": messages,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        "temperature": 0,
+        "max_tokens": args.max_tokens,
+        "seed": 1,
+    }
+    t0 = time.perf_counter()
+    ttft, text, n = None, [], 0
+    with requests.post(f"{url}/v1/chat/completions", json=body, stream=True, timeout=1800) as r:
+        r.raise_for_status()
+        for line in r.iter_lines():
+            if not line or not line.startswith(b"data: ") or line == b"data: [DONE]":
+                continue
+            obj = json.loads(line[6:])
+            if obj.get("usage"):
+                n = obj["usage"]["completion_tokens"]
+            for ch in obj.get("choices", []):
+                delta = ch["delta"].get("content")
+                if delta:
+                    if ttft is None:
+                        ttft = time.perf_counter() - t0
+                    text.append(delta)
+    return "".join(text), ttft or 0.0, time.perf_counter() - t0, n
+
+
+messages = [{"role": "system", "content": "You answer questions about the document below.\n\n" + doc}]
+mismatches = 0
+for turn in range(args.turns):
+    messages.append({"role": "user", "content": questions[turn % len(questions)]})
+    ref, ttft_r, tot_r, n_r = ask(args.ref, messages)
+    pd, ttft_p, tot_p, n_p = ask(args.pd, messages)
+    same = ref == pd
+    mismatches += not same
+    print(
+        f"turn {turn + 1}: ref TTFT {ttft_r * 1e3:6.0f} ms total {tot_r * 1e3:6.0f} ms ({n_r} tok) | PD TTFT {ttft_p * 1e3:6.0f} ms total {tot_p * 1e3:6.0f} ms ({n_p} tok) | "
+        f"outputs {'identical' if same else 'DIFFER'}"
+    )
+    if not same:
+        print("   ref:", repr(ref[:160]))
+        print("   pd :", repr(pd[:160]))
+    messages.append({"role": "assistant", "content": ref})
+raise SystemExit(1 if mismatches else 0)

serve/check_prefix_cache.py

@@ -0,0 +1,87 @@
+"""A/B check of the host prefix cache on a running serve/openai_server.py.
+
+Builds a multi-turn conversation with a long shared prefix, sends every turn twice (cache bypassed via
+"prefix_cache": false, then with the cache) with greedy sampling, and compares outputs and TTFT.
+
+    python serve/check_prefix_cache.py --port 8888 --turns 3 --prefix-tokens 6000
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+
+import requests
+
+p = argparse.ArgumentParser()
+p.add_argument("--port", type=int, default=8888)
+p.add_argument("--model", default="deepseek-ai/DeepSeek-V4.1-Flash-1017")
+p.add_argument("--turns", type=int, default=3)
+p.add_argument("--prefix-tokens", type=int, default=6000, help="approximate size of the shared document")
+p.add_argument("--max-tokens", type=int, default=64)
+args = p.parse_args()
+URL = f"http://localhost:{args.port}/v1/chat/completions"
+
+# A deterministic pseudo-document; ~1 token per word for this tokenizer is close enough.
+words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa", "lambda", "mu"]
+doc = " ".join(f"{words[i % len(words)]}{i % 97}" for i in range(args.prefix_tokens))
+questions = [
+    "Which word appears right after 'gamma2'? Answer briefly.",
+    "How many distinct word stems are used in the document? Answer briefly.",
+    "Quote the first five words of the document.",
+    "What is the last word of the document?",
+]
+
+
+def ask(messages: list[dict], use_cache: bool) -> tuple[str, float, float, dict]:
+    body = {
+        "model": args.model,
+        "messages": messages,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        "temperature": 0,
+        "max_tokens": args.max_tokens,
+        "seed": 1,
+        "prefix_cache": use_cache,
+    }
+    t0 = time.perf_counter()
+    ttft = None
+    text, usage = [], {}
+    with requests.post(URL, json=body, stream=True, timeout=1800) as r:
+        r.raise_for_status()
+        for line in r.iter_lines():
+            if not line or not line.startswith(b"data: ") or line == b"data: [DONE]":
+                continue
+            obj = json.loads(line[6:])
+            if obj.get("usage"):
+                usage = obj["usage"]
+            for ch in obj.get("choices", []):
+                delta = ch["delta"].get("content")
+                if delta:
+                    if ttft is None:
+                        ttft = time.perf_counter() - t0
+                    text.append(delta)
+    return "".join(text), ttft or 0.0, time.perf_counter() - t0, usage
+
+
+messages = [{"role": "system", "content": "You answer questions about the document below.\n\n" + doc}]
+mismatches = 0
+for turn in range(args.turns):
+    messages.append({"role": "user", "content": questions[turn % len(questions)]})
+    base, ttft_b, tot_b, usage_b = ask(messages, use_cache=False)
+    cached, ttft_c, tot_c, usage_c = ask(messages, use_cache=True)
+    same = base == cached
+    mismatches += not same
+    print(
+        f"turn {turn + 1}: prompt {usage_c.get('prompt_tokens')} tok | TTFT no-cache {ttft_b * 1e3:7.0f} ms, cache {ttft_c * 1e3:7.0f} ms "
+        f"({ttft_b / max(ttft_c, 1e-6):.1f}x) | outputs {'identical' if same else 'DIFFER'}"
+    )
+    if not same:
+        print("   no-cache:", repr(base[:200]))
+        print("   cache   :", repr(cached[:200]))
+    messages.append({"role": "assistant", "content": base})
+
+stats = requests.get(f"http://localhost:{args.port}/deepinfer/metrics", timeout=30).text
+print("\n".join(line for line in stats.splitlines() if "hicache:hits" in line or "hicache:lookups" in line or "hit_rate" in line))
+raise SystemExit(1 if mismatches else 0)

serve/nccl_xtray_test.py

@@ -0,0 +1,30 @@
+"""Cross-tray NCCL check: run one rank on each of two trays and report the transport NCCL picks + bandwidth.
+
+# on tray A (rank 0) and tray B (rank 1), same MASTER_ADDR:
+MASTER_ADDR=<trayA> MASTER_PORT=29999 RANK=0 WORLD_SIZE=2 CUDA_VISIBLE_DEVICES=0 NCCL_DEBUG=INFO python serve/nccl_xtray_test.py
+"""
+
+import os
+import time
+
+import torch
+import torch.distributed as dist
+
+dist.init_process_group("nccl")
+rank = dist.get_rank()
+torch.cuda.set_device(0)
+x = torch.ones(256 << 20, dtype=torch.uint8, device="cuda")  # 256 MiB
+for _ in range(3):
+    dist.all_reduce(x)
+torch.cuda.synchronize()
+t0 = time.perf_counter()
+n = 10
+for _ in range(n):
+    dist.all_reduce(x)
+torch.cuda.synchronize()
+dt = (time.perf_counter() - t0) / n
+print(
+    f"rank {rank} on {os.uname().nodename}: all_reduce 256 MiB avg {dt * 1e3:.2f} ms -> {x.numel() / dt / 2**30:.1f} GiB/s algbw",
+    flush=True,
+)
+dist.destroy_process_group()

serve/openai_server.py

@@ -0,0 +1,577 @@
+"""Minimal OpenAI-compatible chat server on top of the DeepInfer offline engine.
+
+DeepInfer ships offline inference only. This wraps one `LLM` per data-parallel rank
+(attention DP + MoE EP, same layout as examples/benchmark.py) behind a FastAPI frontend
+so request-replay benchmarks such as InferenceX AgentX (aiperf) can drive it.
+
+    python serve/openai_server.py --model ~/models/DeepSeek-V4.1-Flash-1017 --dp 4 --port 8888
+
+Limitations, by design of the open-source engine:
+- No cross-request prefix cache: every request prefills its full prompt.
+- Responses are streamed as plain `content`; DSML tool calls are not parsed into `tool_calls`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import collections
+import hashlib
+import json
+import multiprocessing as mp
+import os
+import queue
+import sys
+import threading
+import time
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+# Module level so FastAPI can resolve the string annotation `request: Request`.
+from fastapi import FastAPI, Request
+
+if TYPE_CHECKING:
+    from deepinfer._C.engine import Sequence
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(REPO_ROOT / "examples"))
+sys.path.insert(0, str(REPO_ROOT / "serve"))
+from prompt_tokenizer import PromptTokenizer  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", required=True, help="local HF checkpoint directory")
+    p.add_argument("--served-model-name", default=None)
+    p.add_argument("--dp", type=int, default=4, help="data-parallel ranks on this node (= EP size)")
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8888)
+    p.add_argument("--max-model-len", type=int, default=1 << 20)
+    p.add_argument("--max-num-seqs", type=int, default=128, help="per DP rank")
+    p.add_argument("--max-num-batched-tokens", type=int, default=32768, help="prefill chunk per DP rank")
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    p.add_argument("--yoco-keep-tokens", type=int, default=140, help="0 disables decoder SWA bounded replay")
+    p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"])
+    p.add_argument("--log-interval-steps", type=int, default=200)
+    p.add_argument("--hicache-size-gb", type=float, default=32, help="pinned host KV page pool per DP rank; 0 disables the prefix cache")
+    p.add_argument("--hicache-state-slots", type=int, default=4096, help="host SWA state slots per DP rank")
+    p.add_argument(
+        "--routing",
+        default="affinity",
+        choices=["affinity", "least-loaded"],
+        help="affinity: a conversation's turns go to one rank (needed for cache hits)",
+    )
+    p.add_argument(
+        "--prefill-any-interval",
+        type=int,
+        default=30,
+        help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this "
+        "many decode steps later even if other ranks are idle. A prompt step pauses decode on every rank, so a larger "
+        "value gives faster decode and more throughput at the cost of TTFT. 30 favours throughput; use 2 for low TTFT",
+    )
+    p.add_argument(
+        "--eplb-interval",
+        type=int,
+        default=0,
+        help="DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL: rearrange experts every N steps; 0 (default) disables EPLB. "
+        "Each rearrangement stalls every rank (~0.3 s on B200/GB300), which under online traffic costs far more than "
+        "the balance it buys",
+    )
+    p.add_argument(
+        "--eplb-after-prefill",
+        type=int,
+        default=0,
+        choices=[0, 1],
+        help="DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE: also rearrange after every prefill step (only with "
+        "--eplb-interval > 0); with agentic traffic nearly every step batch contains a prefill",
+    )
+    p.add_argument(
+        "--tokenizer-backend",
+        default="auto",
+        choices=["auto", "fastokens", "segment", "hf"],
+        help="prompt tokenization (serve/prompt_tokenizer.py); auto = the first of fastokens (if installed), segment, hf "
+        "that matches tokenizer.encode on startup probes. All produce tokenizer.encode's ids",
+    )
+    p.add_argument(
+        "--tokenizer-verify-rate",
+        type=float,
+        default=0.005,
+        help="fraction of prompts re-encoded with tokenizer.encode in the background; a mismatch switches to hf",
+    )
+    return p.parse_args()
+
+
+# --------------------------------------------------------------------------- engine side
+
+
+def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None:
+    """One DP rank: owns an engine and steps it in lockstep with the other ranks."""
+    import ctypes
+    import signal
+
+    # Die with the frontend instead of lingering as an orphan that holds the GPU.
+    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL)  # PR_SET_PDEATHSIG
+
+    from _config import create_llm, setup_common_env
+
+    if args.yoco_keep_tokens > 0:
+        os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact"
+        os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens)
+        os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1"
+    if args.hicache_size_gb > 0:
+        # Keep more SWA history in the state ring so the prefix cache can store states at several
+        # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks).
+        os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512")
+    os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval))
+    if args.eplb_interval == 0:
+        # The 32 redundant expert slots only pay off with EPLB rearrangement; without it they are fixed copies of
+        # experts 0-31 and only add expert weight to read per rank (decode +4-5% without them on B200 and GB300).
+        os.environ.setdefault("DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS", "0")
+    setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False)
+    if args.eplb_interval > 0:
+        os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval)
+        os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill)
+
+    llm = create_llm(
+        args.model,
+        dp_rank,
+        ep_size=args.dp,
+        env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)},
+        data_parallel_size=args.dp,
+        max_num_seqs=args.max_num_seqs,
+        max_num_batched_tokens=args.max_num_batched_tokens,
+        max_model_len=args.max_model_len,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
+    )
+
+    import numpy as np
+
+    from deepinfer._C.engine import RequestOutput  # pyright: ignore[reportMissingModuleSource]
+    from deepinfer.sampling_params import SamplingParams
+
+    engine = llm.llm_engine
+    cache = None
+    if args.hicache_size_gb > 0:
+        try:
+            from deepinfer.hicache import HostPrefixCache
+        except ImportError:  # tree without the host prefix cache: serve without it
+            print(f"[server] deepinfer.hicache not available, prefix cache disabled (rank {dp_rank})", flush=True)
+        else:
+            cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots)
+            engine.pe_do_save_kvcache = cache.on_prefill_step  # write-through at the end of every prefill
+    out_q.put(("ready", dp_rank, None))
+
+    def mm_data():
+        return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None
+
+    def enqueue(rid: str, seq: Sequence, num_cached_tokens: int) -> None:
+        # Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict
+        # (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length.
+        engine.scheduler.add_sequence(seq)
+        if engine.multi_modal_processor is not None:
+            engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached_tokens)
+
+    pending = collections.deque()  # (rid, seq, prompt_ids, use_cache): created, waiting for a cache decision
+    inflight = False
+    last_stats = time.monotonic()
+    last_logged_lookups = 0
+    while True:
+        # Requests must be enqueued before the DP-wide has_unfinished_requests all-reduce,
+        # so every rank takes the same step / flush / idle branch in the same iteration.
+        while True:
+            try:
+                kind, rid, payload = in_q.get_nowait()
+            except queue.Empty:
+                break
+            if kind == "add":
+                try:
+                    prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32)
+                    seq = engine.add_request(
+                        user_id="0",
+                        request_id=rid,
+                        sampling_params=SamplingParams(**payload["sampling"]),
+                        prompt_token_ids=prompt_ids,
+                        multi_modal_data=mm_data(),
+                        enqueue=cache is None,
+                    )
+                    if cache is not None:
+                        if seq is None:
+                            out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
+                        else:
+                            pending.append((rid, seq, prompt_ids, payload.get("use_cache", True)))
+                except Exception as e:  # noqa: BLE001 - surface to the client, keep the rank alive
+                    out_q.put(("error", rid, repr(e)))
+            elif kind == "abort":
+                if any(p[0] == rid for p in pending):
+                    pending = collections.deque(p for p in pending if p[0] != rid)
+                else:
+                    engine.abort_request(rid)
+                if cache is not None:
+                    cache.forget(rid)
+
+        # Cache decisions, in arrival order. A hit pre-allocates the sequence's blocks and fills them from the
+        # host; if the GPU cannot allocate yet, wait (while something is running that will free memory).
+        while pending:
+            rid, seq, prompt_ids, use_cache = pending[0]
+            hit = cache.try_resume(seq, rid, prompt_ids) if use_cache else 0
+            if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
+                break
+            pending.popleft()
+            enqueue(rid, seq, max(hit, 0))
+        if cache is not None:
+            cache.poll()
+            if time.monotonic() - last_stats > 5:
+                last_stats = time.monotonic()
+                st = cache.stats_dict()
+                out_q.put(("stats", dp_rank, st))
+                if st["lookups"] != last_logged_lookups:
+                    last_logged_lookups = st["lookups"]
+                    print(
+                        f"[hicache rank {dp_rank}] lookups {st['lookups']} hits {st['hits']} "
+                        f"token hit rate {st['hit_rate_tokens']:.1%} deferred {st['deferred']} stores {st['stores']} "
+                        f"pages {st['pages_used']}/{st['pages_total']} states {st['states_used']}/{st['states_total']} "
+                        f"evicted {st['evicted_pages']}",
+                        flush=True,
+                    )
+
+        if engine.has_unfinished_requests():
+            outputs = engine.step()
+            inflight = True
+        elif inflight:
+            outputs = engine.flush_inflight_reqs()
+            inflight = False
+        else:
+            time.sleep(0.001)
+            continue
+
+        if len(outputs):
+            # Pointers are recycled on the next step, so read the fields now.
+            batch = [
+                (
+                    RequestOutput.get_request_id(p),
+                    RequestOutput.get_token_id(p),
+                    RequestOutput.get_finished(p),
+                    RequestOutput.get_finish_reason(p),
+                    RequestOutput.get_prompt_len(p),
+                )
+                for p in outputs
+            ]
+            if cache is not None:
+                for rid, _, finished, _, _ in batch:
+                    if finished:
+                        cache.forget(rid)
+            out_q.put(("tokens", None, batch))
+
+
+# --------------------------------------------------------------------------- HTTP side
+
+
+class IncrementalDetokenizer:
+    """vLLM-style prefix/read offsets so each delta decodes only a short tail."""
+
+    def __init__(self, tokenizer: Any) -> None:
+        self.tok = tokenizer
+        self.ids: list[int] = []
+        self.num_tokens = 0
+        self.prefix_offset = 0
+        self.read_offset = 0
+
+    def add(self, token_id: int) -> str:
+        self.num_tokens += 1
+        if token_id == self.tok.eos_token_id:
+            return ""
+        self.ids.append(token_id)
+        prefix = self.tok.decode(self.ids[self.prefix_offset : self.read_offset], skip_special_tokens=False)
+        full = self.tok.decode(self.ids[self.prefix_offset :], skip_special_tokens=False)
+        if len(full) > len(prefix) and not full.endswith("<EFBFBD>"):
+            self.prefix_offset, self.read_offset = self.read_offset, len(self.ids)
+            return full[len(prefix) :]
+        return ""
+
+
+class Frontend:
+    def __init__(self, args: argparse.Namespace) -> None:
+        from transformers import AutoTokenizer
+
+        sys.path.insert(0, str(Path(args.model) / "encoding"))
+        from encoding import encode_messages  # the checkpoint's own prompt-format reference
+
+        self.args = args
+        self.model_name = args.served_model_name or args.model
+        self.encode_messages = encode_messages
+        self.tokenizer = AutoTokenizer.from_pretrained(args.model)
+        probe = [
+            {"role": "system", "content": "sys", "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}]},
+            {"role": "user", "content": "hi 你好 123"},
+            {"role": "assistant", "content": "ok", "tool_calls": [{"type": "function", "function": {"name": "f", "arguments": "{}"}}]},
+            {"role": "tool", "content": "  out\n"},
+        ]
+        try:
+            probes = [encode_messages(probe, thinking_mode=args.thinking_mode)]
+        except Exception:  # noqa: BLE001 - the probe only adds coverage
+            probes = []
+        # getattr: frontends built on this one with their own argument parser (serve/pd_server.py) get the defaults
+        backend, verify_rate = getattr(args, "tokenizer_backend", "auto"), getattr(args, "tokenizer_verify_rate", 0.005)
+        self.prompt_tokenizer = PromptTokenizer(self.tokenizer, backend, verify_rate, extra_probes=probes)
+        print(f"[server] {self.prompt_tokenizer.describe()}", flush=True)
+        ctx = mp.get_context("spawn")
+        self.out_q: mp.Queue = ctx.Queue()
+        self.in_qs: list[mp.Queue] = [ctx.Queue() for _ in range(args.dp)]
+        # Not daemonic: the engine forks its own helper processes (as examples/_launch.py allows).
+        self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)]
+        self.load = [0] * args.dp
+        self.owner: dict[str, int] = {}
+        self.streams: dict[str, asyncio.Queue] = {}
+        self.ready = 0
+        self.stats = dict(requests=0, finished=0, aborted=0, errors=0, prompt_tokens=0, completion_tokens=0, encode_fallbacks=0)
+        self.cache_stats: dict[int, dict] = {}
+        self.loop: asyncio.AbstractEventLoop | None = None
+
+    def start(self) -> None:
+        for p in self.procs:
+            p.start()
+
+    def dispatch_forever(self) -> None:
+        """Thread: fan engine outputs out to the per-request asyncio queues."""
+        while True:
+            try:
+                kind, rid, payload = self.out_q.get(timeout=1.0)
+            except queue.Empty:
+                dead = [r for r, p in enumerate(self.procs) if not p.is_alive()]
+                if dead:
+                    print(f"[server] engine rank(s) {dead} died; failing open requests and exiting", flush=True)
+                    for rid in list(self.streams):
+                        self.loop.call_soon_threadsafe(self._deliver, rid, ("error", f"engine rank died: {dead}"))
+                    time.sleep(2)
+                    os._exit(1)
+                continue
+            if kind == "ready":
+                self.ready += 1
+                print(f"[server] DP rank {rid} ready ({self.ready}/{self.args.dp})", flush=True)
+            elif kind == "error":
+                self.loop.call_soon_threadsafe(self._deliver, rid, ("error", payload))
+            elif kind == "stats":
+                self.cache_stats[rid] = payload
+            else:
+                self.loop.call_soon_threadsafe(self._deliver_batch, payload)
+
+    def _deliver(self, rid: str, item: tuple) -> None:
+        q = self.streams.get(rid)
+        if q is not None:
+            q.put_nowait(item)
+
+    def _deliver_batch(self, batch: list[tuple]) -> None:
+        for rid, token_id, finished, reason, prompt_len in batch:
+            q = self.streams.get(rid)
+            if q is not None:
+                q.put_nowait(("token", token_id, finished, reason, prompt_len))
+
+    def render_prompt(self, body: dict) -> list[int]:
+        messages = [dict(m) for m in body["messages"]]
+        if body.get("tools"):
+            if not messages or messages[0].get("role") != "system":
+                messages.insert(0, {"role": "system", "content": ""})
+            messages[0]["tools"] = body["tools"]
+        try:
+            prompt = self.encode_messages(messages, thinking_mode=self.args.thinking_mode)
+        except Exception:  # noqa: BLE001 - trace messages the reference encoder rejects
+            # Keep the request (and its prompt length) instead of failing the benchmark.
+            self.stats["encode_fallbacks"] += 1
+            parts = ["<｜begin▁of▁sentence｜>"]
+            for m in messages:
+                content = m.get("content") or ""
+                if isinstance(content, list):
+                    content = "".join(b.get("text", "") for b in content if isinstance(b, dict))
+                if m.get("tool_calls"):
+                    content += json.dumps(m["tool_calls"], ensure_ascii=False)
+                tag = {"user": "<｜User｜>", "assistant": "<｜Assistant｜>", "tool": "<｜User｜>"}.get(m.get("role"), "")
+                parts.append(tag + content)
+            if body.get("tools"):
+                parts.insert(1, json.dumps(body["tools"], ensure_ascii=False))
+            prompt = "".join(parts) + "<｜Assistant｜></think>"
+        return self.prompt_tokenizer.encode(prompt)
+
+    @staticmethod
+    def affinity_key(body: dict) -> bytes | None:
+        """Stable per-conversation key: the first user message (AgentX's cache-bust marker lives there and is
+        constant across a trajectory's turns, and forked subagents inherit it)."""
+        for m in body.get("messages", []):
+            if m.get("role") == "user":
+                content = m.get("content") or ""
+                if isinstance(content, list):
+                    content = "".join(b.get("text", "") for b in content if isinstance(b, dict))
+                return hashlib.sha1(content[:4096].encode("utf-8", "ignore")).digest()
+        return None
+
+    def submit(self, prompt_ids: list[int], sampling: dict, key: bytes | None = None, use_cache: bool = True) -> tuple[str, asyncio.Queue]:
+        rid = uuid.uuid4().hex
+        if self.args.routing == "affinity" and key is not None:
+            rank = int.from_bytes(key[:4], "little") % self.args.dp
+        else:
+            rank = min(range(self.args.dp), key=lambda r: self.load[r])
+        self.load[rank] += 1
+        self.owner[rid] = rank
+        q: asyncio.Queue = asyncio.Queue()
+        self.streams[rid] = q
+        self.stats["requests"] += 1
+        self.in_qs[rank].put(("add", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling, "use_cache": use_cache}))
+        return rid, q
+
+    def release(self, rid: str, aborted: bool) -> None:
+        rank = self.owner.pop(rid, None)
+        self.streams.pop(rid, None)
+        if rank is None:
+            return
+        self.load[rank] -= 1
+        if aborted:
+            self.stats["aborted"] += 1
+            self.in_qs[rank].put(("abort", rid, None))
+        else:
+            self.stats["finished"] += 1
+
+
+def build_app(fe: Frontend) -> FastAPI:
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
+
+    app = FastAPI()
+
+    @app.on_event("startup")
+    async def _startup() -> None:
+        fe.loop = asyncio.get_running_loop()
+        threading.Thread(target=fe.dispatch_forever, daemon=True).start()
+
+    @app.get("/health")
+    async def health():
+        ok = fe.ready == fe.args.dp
+        return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready}, status_code=200 if ok else 503)
+
+    @app.get("/v1/models")
+    async def models():
+        return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]}
+
+    # Not at /metrics: aiperf scrapes that path automatically and InferenceX's aggregator
+    # rejects metric families it has no backend adapter for.
+    @app.get("/deepinfer/metrics")
+    async def metrics():
+        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
+        lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)]
+        for r, st in sorted(fe.cache_stats.items()):
+            lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()]
+        pt = fe.prompt_tokenizer
+        lines += [f'deepinfer_tokenizer:{k}{{backend="{pt.name}"}} {v}' for k, v in pt.metrics().items()]
+        return PlainTextResponse("\n".join(lines) + "\n")
+
+    @app.post("/v1/chat/completions")
+    async def chat(request: Request):
+        if fe.ready != fe.args.dp:
+            return JSONResponse({"error": {"message": "engine loading"}}, status_code=503)
+        body = await request.json()
+        prompt_ids = await asyncio.to_thread(fe.render_prompt, body)
+        max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096
+        if len(prompt_ids) + 1 > fe.args.max_model_len:
+            return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400)
+        max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids)))
+        sampling = dict(
+            temperature=float(body.get("temperature", 1.0)),
+            top_p=float(body.get("top_p", 1.0)),
+            ignore_eos=bool(body.get("ignore_eos", False)),
+            max_tokens=max_tokens,
+        )
+        if body.get("seed") is not None:
+            sampling["seq_seed"] = int(body["seed"])
+        fe.stats["prompt_tokens"] += len(prompt_ids)
+        include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
+        # Non-standard request field: "prefix_cache": false bypasses the host prefix cache lookup (A/B checks).
+        rid, q = fe.submit(prompt_ids, sampling, fe.affinity_key(body), use_cache=bool(body.get("prefix_cache", True)))
+        created = int(time.time())
+        cid = f"chatcmpl-{rid}"
+        detok = IncrementalDetokenizer(fe.tokenizer)
+
+        def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str:
+            obj = {
+                "id": cid,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": fe.model_name,
+                "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}],
+            }
+            if usage:
+                obj["usage"] = usage
+            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"
+
+        async def tokens():
+            """Yield (text_delta, finish_reason) until the request finishes."""
+            while True:
+                item = await q.get()
+                if item[0] == "error":
+                    raise RuntimeError(item[1])
+                _, token_id, finished, reason, _ = item
+                text = detok.add(token_id)
+                fe.stats["completion_tokens"] += 1
+                if finished:
+                    yield text, ("length" if detok.num_tokens >= max_tokens else "stop")
+                    return
+                yield text, None
+
+        usage = lambda: {
+            "prompt_tokens": len(prompt_ids),
+            "completion_tokens": detok.num_tokens,  # noqa: E731
+            "total_tokens": len(prompt_ids) + detok.num_tokens,
+        }
+
+        if body.get("stream"):
+
+            async def sse():
+                done = False
+                try:
+                    yield chunk({"role": "assistant", "content": ""})
+                    async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
+                        if text or finish:
+                            yield chunk({"content": text} if text else {}, finish)
+                    done = True
+                    if include_usage:
+                        yield chunk({}, usage=usage())
+                    yield "data: [DONE]\n\n"
+                finally:
+                    fe.release(rid, aborted=not done)
+
+            return StreamingResponse(sse(), media_type="text/event-stream")
+
+        done = False
+        try:
+            parts, finish = [], None
+            async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
+                parts.append(text)
+            done = True
+        finally:
+            fe.release(rid, aborted=not done)
+        return {
+            "id": cid,
+            "object": "chat.completion",
+            "created": created,
+            "model": fe.model_name,
+            "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}],
+            "usage": usage(),
+        }
+
+    return app
+
+
+def main() -> None:
+    args = parse_args()
+    args.model = str(Path(args.model).expanduser().resolve())
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("RANK", "0")
+    fe = Frontend(args)
+    fe.start()
+
+    import uvicorn
+
+    uvicorn.run(build_app(fe), host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900)
+
+
+if __name__ == "__main__":
+    main()

serve/pd_run.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+# Start/stop the prefill (PE) and decode (CE) trays of serve/pd_server.py over ssh.
+#
+#   serve/pd_run.sh start <pe-host> <ce-host> [extra pd_server args...]
+#   serve/pd_run.sh stop  <pe-host> <ce-host>
+#   serve/pd_run.sh log   <host>              # tail the tray's log
+#
+# Ports: PE HTTP 8889, CE HTTP (user-facing) 8888, CE KV receivers 8900+rank. Logs: ~/pd_<role>.log on each tray.
+set -euo pipefail
+CMD="${1:?start|stop|log}"; shift
+MODEL="${DEEPINFER_HF_MODEL_DIR:-$HOME/models/DeepSeek-V4.1-Flash-1017}"
+SERVED=deepseek-ai/DeepSeek-V4.1-Flash-1017
+ENV='cd ~/DeepInfer && source env.sh && export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}'
+
+launch() {  # host role args...
+    local host=$1 role=$2; shift 2
+    # ssh -n + nohup with every fd redirected, otherwise the remote session stays attached to the server.
+    ssh -n "$host" "$ENV && rm -f ~/pd_$role.pgid && nohup setsid python -u serve/pd_server.py --role $role --model $MODEL --served-model-name $SERVED $* > ~/pd_$role.log 2>&1 < /dev/null & echo \$! > ~/pd_$role.pgid; echo started $role on \$(hostname) pgid \$(cat ~/pd_$role.pgid)"
+}
+wait_ready() {  # host port
+    local i
+    for i in $(seq 200); do
+        if curl -fs "http://$1:$2/health" >/dev/null 2>&1; then echo "$1:$2 ready"; return 0; fi
+        sleep 3
+    done
+    echo "$1:$2 did not become ready" >&2; return 1
+}
+
+case "$CMD" in
+    start)
+        PE=${1:?pe-host}; CE=${2:?ce-host}; shift 2
+        launch "$PE" pe --port 8889 "$@"
+        launch "$CE" ce --port 8888 --kv-port 8900 --kv-host "$CE" --pe-url "http://$PE:8889" "$@"
+        wait_ready "$PE" 8889; wait_ready "$CE" 8888
+        ;;
+    stop)
+        for spec in "${1:?pe-host}:pe" "${2:?ce-host}:ce"; do
+            h=${spec%%:*}; r=${spec##*:}
+            # Kill by pid list, excluding the remote shell itself: a pkill -f pattern would match this very
+            # command line (it names pd_server.py) and kill the shell before anything else runs.
+            # Frontend + its spawned engine workers, SIGKILL (uvicorn's graceful shutdown otherwise keeps the
+            # children, and the GPUs, alive for minutes); then wait until the GPUs are actually free.
+            ssh -n "$h" "for p in \$(pgrep -f '[p]d_server.py --role $r|multiprocessing.spawn|multiprocessing.resource_tracker'); do [ \"\$p\" != \"\$\$\" ] && kill -9 \"\$p\" 2>/dev/null; done; rm -f ~/pd_$r.pgid; for i in \$(seq 60); do [ \"\$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | sort -n | tail -1)\" -lt 1000 ] && break; sleep 3; done; echo stopped $r on \$(hostname), gpu mem \$(nvidia-smi --query-gpu=memory.used --format=csv,noheader | tr '\\n' ' ')" || true
+        done
+        ;;
+    log)
+        ssh "${1:?host}" 'tail -n 40 ~/pd_*.log'
+        ;;
+esac

serve/pd_server.py

@@ -0,0 +1,826 @@
+"""Prefill/decode-disaggregated serving: a prompt-only engine tray (PE) and a decode-only engine tray (CE).
+
+    # decode tray (user-facing OpenAI API on --port, KV receivers on --kv-port + rank):
+    python serve/pd_server.py --role ce --model ... --dp 4 --port 8888 --kv-port 8900 --pe-url http://<pe-host>:8889
+    # prefill tray:
+    python serve/pd_server.py --role pe --model ... --dp 4 --port 8889
+
+Flow per request: CE frontend renders the prompt, picks a CE rank, POSTs /prefill to the PE with the KV target
+(ce-host, kv-port + rank). The PE prefills (host prefix cache on), exports the sequence's KV blocks + full SWA
+state (deepinfer.hicache.pd_transfer) and the first sampled token, streams the first token back in the HTTP
+response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence straight into decode
+(Scheduler.add_prompted_sequence); its tokens stream to the client through the CE frontend.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import collections
+import json
+import multiprocessing as mp
+import os
+import queue
+import socket
+import struct
+import sys
+import threading
+import time
+import uuid
+from collections.abc import Callable
+from pathlib import Path
+from types import ModuleType
+from typing import TYPE_CHECKING
+
+from fastapi import FastAPI, Request
+
+if TYPE_CHECKING:
+    from deepinfer._C.engine import RequestOutput, Sequence
+    from deepinfer.engine.llm_engine import LLMEngine
+    from deepinfer.hicache.pd_transfer import PDTransfer
+    from deepinfer.sampling_params import SamplingParams
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+sys.path.insert(0, str(HERE.parent / "examples"))
+
+import openai_server as oas  # noqa: E402  (the single-engine server: prompt rendering, detokenizer, frontend plumbing)
+
+HDR = struct.Struct("<16sqqqqq")  # rid(16 bytes hex) | first_token | prompt_len | num_pages | page_bytes | state_bytes
+# --kv-transport mooncake: MC_MAGIC + HDR; CE replies one JSON line {"status": ok|busy|error, "addr", "segment"};
+# after the TransferEngine write the PE sends b"go" (or b"xx" on failure) and the CE answers b"ok".
+MC_MAGIC = b"MCK1"  # cannot start a TCP-path header (rids are hex)
+MC_RETRY_S = 0.005
+MC_WRITE_ATTEMPTS = 3  # then fall back to TCP for that transfer, so a failing peer never strands a request
+
+
+def _mooncake_link(args: argparse.Namespace, device):  # noqa: ANN001, ANN202
+    from deepinfer.hicache.pd_mooncake import GpuArena, MooncakeLink
+
+    arena = GpuArena(int(args.pd_arena_gb * 2**30), device)
+    host = args.mooncake_host or (args.kv_host if args.role == "ce" and args.kv_host else socket.gethostbyname(socket.gethostname()))
+    return MooncakeLink(host, args.mooncake_protocol, arena, args.mooncake_device)
+
+
+def _recv_line(sock: socket.socket) -> bytes:
+    out = bytearray()
+    while not out.endswith(b"\n"):
+        chunk = sock.recv(1)
+        if not chunk:
+            raise ConnectionError("connection closed")
+        out += chunk
+    return bytes(out)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--role", required=True, choices=["pe", "ce"])
+    p.add_argument("--model", required=True)
+    p.add_argument("--served-model-name", default=None)
+    p.add_argument("--dp", type=int, default=4)
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8888)
+    p.add_argument("--kv-port", type=int, default=8900, help="ce: TCP receiver base port (+ rank)")
+    p.add_argument("--kv-host", default=None, help="ce: address the PE should push KV to (default: this host's name)")
+    p.add_argument("--pe-url", default=None, help="ce: prefill tray URL")
+    p.add_argument("--max-model-len", type=int, default=1 << 20)
+    p.add_argument("--max-num-seqs", type=int, default=128)
+    p.add_argument("--max-num-batched-tokens", type=int, default=32768)
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    p.add_argument("--yoco-keep-tokens", type=int, default=140)
+    p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"])
+    p.add_argument("--log-interval-steps", type=int, default=200)
+    p.add_argument("--hicache-size-gb", type=float, default=32, help="pe: host prefix cache per rank (0 disables)")
+    p.add_argument("--hicache-state-slots", type=int, default=4096)
+    p.add_argument("--routing", default="affinity", choices=["affinity", "least-loaded"])
+    p.add_argument("--prefill-any-interval", type=int, default=2)
+    p.add_argument("--eplb-interval", type=int, default=0, help="see openai_server.py; 0 (default) disables EPLB")
+    p.add_argument("--eplb-after-prefill", type=int, default=0, choices=[0, 1], help="see openai_server.py")
+    p.add_argument("--tokenizer-backend", default="auto", choices=["auto", "fastokens", "segment", "hf"], help="see openai_server.py")
+    p.add_argument("--tokenizer-verify-rate", type=float, default=0.005, help="see openai_server.py")
+    p.add_argument(
+        "--kv-transport",
+        default="tcp",
+        choices=["tcp", "mooncake"],
+        help="KV/state data plane PE -> CE. tcp: pinned host buffers over the KV socket. mooncake: GPU arena to GPU "
+        "arena through a Mooncake TransferEngine (nvlink_intra within a node, rdma/efa across hosts); the KV socket "
+        "carries only control messages",
+    )
+    p.add_argument("--mooncake-protocol", default="nvlink_intra", help="TransferEngine protocol (nvlink_intra, rdma, efa, tcp)")
+    p.add_argument("--mooncake-device", default="", help="TransferEngine device name(s) for rdma/efa (empty: auto)")
+    p.add_argument("--mooncake-host", default=None, help="address peers use for this rank's TransferEngine handshake")
+    p.add_argument("--pd-arena-gb", type=float, default=4, help="mooncake: GPU staging arena per rank (PE send / CE receive)")
+    return p.parse_args()
+
+
+def _recv_exact(sock: socket.socket, buf: memoryview) -> None:
+    view = buf
+    while len(view):
+        n = sock.recv_into(view, min(len(view), 1 << 24))
+        if n == 0:
+            raise ConnectionError("peer closed during transfer")
+        view = view[n:]
+
+
+# ============================================================================ engine processes
+
+
+def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None:
+    import ctypes
+    import signal
+
+    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL)  # PR_SET_PDEATHSIG
+
+    import numpy as np
+    from _config import create_llm, setup_common_env
+
+    from deepinfer.engine.engine_type import EngineType
+
+    if args.yoco_keep_tokens > 0:
+        os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact"
+        os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens)
+        os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1"
+    os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512")  # identical state layout on PE and CE
+    os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval))
+    if args.eplb_interval == 0:
+        # Same as openai_server.py: without EPLB the redundant expert slots only add expert weight per rank.
+        os.environ.setdefault("DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS", "0")
+    setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False)
+    if args.eplb_interval > 0:
+        os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval)
+        os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill)
+
+    engine_type = EngineType.PROMPT if args.role == "pe" else EngineType.COMPLETION
+    llm = create_llm(
+        args.model,
+        dp_rank,
+        ep_size=args.dp,
+        env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)},
+        data_parallel_size=args.dp,
+        max_num_seqs=args.max_num_seqs,
+        max_num_batched_tokens=args.max_num_batched_tokens,
+        max_model_len=args.max_model_len,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
+        engine_type=engine_type,
+        stop_preempted_request=True,
+    )
+
+    from deepinfer._C.engine import RequestOutput  # pyright: ignore[reportMissingModuleSource]
+    from deepinfer.hicache.pd_transfer import PDTransfer
+    from deepinfer.sampling_params import SamplingParams
+
+    engine = llm.llm_engine
+    transfer = PDTransfer(engine)
+    mm = lambda: dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None  # noqa: E731
+
+    def enqueue(rid: str, seq: Sequence, num_cached: int) -> None:
+        engine.scheduler.add_sequence(seq)
+        if engine.multi_modal_processor is not None:
+            engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached)
+
+    if args.role == "pe":
+        _pe_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm, enqueue)
+    else:
+        _ce_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm)
+
+
+def _pe_loop(
+    args: argparse.Namespace,
+    dp_rank: int,
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    engine: LLMEngine,
+    transfer: PDTransfer,
+    SamplingParams: type[SamplingParams],
+    RequestOutput: type[RequestOutput],
+    np: ModuleType,
+    mm: Callable[[], dict | None],
+    enqueue: Callable[[str, Sequence, int], None],
+) -> None:
+    cache = None
+    if args.hicache_size_gb > 0:
+        from deepinfer.hicache import HostPrefixCache
+
+        cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots)
+
+    link = _mooncake_link(args, transfer.device) if args.kv_transport == "mooncake" else None
+    targets: dict[str, tuple[str, int]] = {}  # rid -> (host, port) of the CE rank
+    awaiting_token: dict[str, dict] = {}  # rid -> transfer waiting for the step's sampled first token
+    send_q: queue.Queue = queue.Queue()
+
+    def on_prefill_step(scheduled: list[Sequence], finished: list[Sequence]) -> None:
+        # Engine hook at the start of post-processing for a prompt step: the prefill's KV is complete and the
+        # sequence still owns its blocks (with max_tokens=1 it finishes, and is freed, once outputs are processed).
+        if cache is not None:
+            cache.on_prefill_step(scheduled, finished)
+        for seq in scheduled:
+            rid = seq.request_id
+            if seq.chunked_rest_len != 0 or rid not in targets:
+                continue
+            prompt_len = seq.get_len()  # no output appended yet
+            need = (prompt_len + transfer.block_size - 1) // transfer.block_size
+            if len(seq.block_table) < need:
+                out_q.put(("error", rid, f"prefilled sequence has {len(seq.block_table)} blocks, needs {need}"))
+                targets.pop(rid, None)
+                continue
+            region = None
+            if link is not None:
+                from deepinfer.hicache.pd_mooncake import region_layout
+
+                nb = transfer.num_blocks(prompt_len)
+                state_off, total = region_layout(nb, transfer.page_bytes, transfer.state_bytes)
+                off = link.arena.alloc(total)
+                if off is not None:  # else: arena full, this one goes over TCP
+                    region = (off, total)
+                    pages = link.arena.view(off, nb * transfer.page_bytes).view(nb, transfer.page_bytes)
+                    state = link.arena.view(off + state_off, transfer.state_bytes)
+                    ev = transfer.export_into(seq, prompt_len, pages, state)
+            if region is None:
+                pages, state, ev = transfer.export(seq, prompt_len)
+            item = {
+                "pages": pages,
+                "state": state,
+                "region": region,
+                "event": ev,
+                "prompt_len": prompt_len,
+                "target": targets.pop(rid),
+                "first_token": None,
+                "ready": threading.Event(),
+            }
+            awaiting_token[rid] = item
+            send_q.put((rid, item))
+
+    engine.pe_do_save_kvcache = on_prefill_step
+
+    def send_mooncake(rid: str, item: dict, host: str, port: int) -> None:
+        pages, state = item["pages"], item["state"]
+        hdr = MC_MAGIC + HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel())
+        off, total = item["region"]
+        deadline = time.monotonic() + 600
+        while True:
+            with socket.create_connection((host, port), timeout=60) as s:
+                s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+                s.sendall(hdr)
+                resv = json.loads(_recv_line(s))
+                if resv["status"] == "ok":
+                    for attempt in range(MC_WRITE_ATTEMPTS):
+                        try:
+                            link.write(resv["segment"], off, resv["addr"], total)
+                            break
+                        except Exception as e:  # noqa: BLE001
+                            print(
+                                f"[pe rank {dp_rank}] mooncake write to {resv['segment']} failed (attempt {attempt + 1}): {e!r}", flush=True
+                            )
+                            time.sleep(0.01 * (attempt + 1))
+                    else:
+                        s.sendall(b"xx")
+                        return False  # caller falls back to TCP
+                    s.sendall(b"go")
+                    ack = s.recv(2)
+                    if ack != b"ok":
+                        raise RuntimeError(f"CE rejected transfer: {ack!r}")
+                    return True
+                if resv["status"] != "busy" or time.monotonic() > deadline:
+                    raise RuntimeError(f"CE reservation failed: {resv}")
+            time.sleep(MC_RETRY_S)  # CE arena full: retry once it has admitted something
+
+    def sender():
+        if link is not None:
+            import torch
+
+            torch.cuda.set_device(transfer.device)  # TransferEngine copies use the calling thread's device
+        while True:
+            rid, item = send_q.get()
+            host, port = item["target"]
+            try:
+                item["ready"].wait(timeout=600)
+                if item["first_token"] is None:
+                    raise RuntimeError("no first token produced for the prefilled sequence")
+                item["event"].synchronize()
+                if item["region"] is not None:
+                    try:
+                        sent = send_mooncake(rid, item, host, port)
+                        if not sent:  # copy the region to pinned host buffers and use the TCP path below
+                            item["pages"] = item["pages"].to("cpu").pin_memory()
+                            item["state"] = item["state"].to("cpu").pin_memory()
+                    finally:
+                        link.arena.release(*item["region"])
+                    if sent:
+                        continue
+                pages, state = item["pages"], item["state"]
+                with socket.create_connection((host, port), timeout=60) as s:
+                    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+                    s.sendall(
+                        HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel())
+                    )
+                    s.sendall(memoryview(pages.numpy()).cast("B"))
+                    s.sendall(memoryview(state.numpy()).cast("B"))
+                    ack = s.recv(2)
+                if ack != b"ok":
+                    out_q.put(("error", rid, f"CE rejected transfer: {ack!r}"))
+            except Exception as e:  # noqa: BLE001
+                out_q.put(("error", rid, f"KV transfer to {host}:{port} failed: {e!r}"))
+
+    threading.Thread(target=sender, daemon=True).start()
+    out_q.put(("ready", dp_rank, None))
+
+    pending = collections.deque()
+    inflight = False
+    last_stats = time.monotonic()
+    while True:
+        while True:
+            try:
+                kind, rid, payload = in_q.get_nowait()
+            except queue.Empty:
+                break
+            if kind == "add":
+                try:
+                    prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32)
+                    sp = SamplingParams(**payload["sampling"])
+                    sp.max_tokens = 1  # prefill + first token only; the engine then finishes and frees the sequence
+                    seq = engine.add_request(
+                        user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=prompt_ids, multi_modal_data=mm(), enqueue=False
+                    )
+                    if seq is None:
+                        out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
+                    else:
+                        targets[rid] = tuple(payload["kv_target"])
+                        pending.append((rid, seq, prompt_ids))
+                except Exception as e:  # noqa: BLE001
+                    out_q.put(("error", rid, repr(e)))
+            elif kind == "abort":
+                targets.pop(rid, None)
+                if any(p[0] == rid for p in pending):
+                    pending = collections.deque(p for p in pending if p[0] != rid)
+                else:
+                    engine.abort_request(rid)
+                if cache is not None:
+                    cache.forget(rid)
+        while pending:
+            rid, seq, prompt_ids = pending[0]
+            hit = cache.try_resume(seq, rid, prompt_ids) if cache is not None else 0
+            if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
+                break
+            pending.popleft()
+            enqueue(rid, seq, max(hit, 0))
+        if cache is not None:
+            cache.poll()
+            if time.monotonic() - last_stats > 5:
+                last_stats = time.monotonic()
+                out_q.put(("stats", dp_rank, cache.stats_dict()))
+
+        if engine.has_unfinished_requests():
+            outputs = engine.step()
+            inflight = True
+        elif inflight:
+            outputs = engine.flush_inflight_reqs()
+            inflight = False
+        else:
+            time.sleep(0.001)
+            continue
+        if len(outputs):
+            batch = [
+                (
+                    RequestOutput.get_request_id(p),
+                    RequestOutput.get_token_id(p),
+                    RequestOutput.get_finished(p),
+                    RequestOutput.get_finish_reason(p),
+                    RequestOutput.get_prompt_len(p),
+                )
+                for p in outputs
+            ]
+            for rid, token_id, finished, _, _ in batch:
+                item = awaiting_token.pop(rid, None)
+                if item is not None:
+                    item["first_token"] = int(token_id)
+                    item["ready"].set()
+                if finished and cache is not None:
+                    cache.forget(rid)
+            out_q.put(("tokens", None, batch))
+        # A prompt-only engine never reclaims finished sequences by itself (that happens in the decode
+        # scheduler): drop them from `running` and free their blocks/state. The export and cache copies were
+        # enqueued on the stream before any kernel that could reuse those blocks.
+        for seq in engine.scheduler.free_finished_seqs():
+            engine.scheduler.free_seq(seq)
+
+
+def _ce_loop(
+    args: argparse.Namespace,
+    dp_rank: int,
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    engine: LLMEngine,
+    transfer: PDTransfer,
+    SamplingParams: type[SamplingParams],
+    RequestOutput: type[RequestOutput],
+    np: ModuleType,
+    mm: Callable[[], dict | None],
+) -> None:
+    import torch
+
+    recv_q: queue.Queue = queue.Queue()  # (rid, first_token, prompt_len, pages, state, arena region or None)
+    link = _mooncake_link(args, transfer.device) if args.kv_transport == "mooncake" else None
+
+    def receiver():
+        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        srv.bind((args.host, args.kv_port + dp_rank))
+        srv.listen(64)
+        while True:
+            conn, _ = srv.accept()
+            threading.Thread(target=receive_one, args=(conn,), daemon=True).start()
+
+    def receive_mooncake(conn: socket.socket, rid: str, first_token: int, prompt_len: int, num_pages: int) -> None:
+        from deepinfer.hicache.pd_mooncake import region_layout
+
+        state_off, total = region_layout(num_pages, transfer.page_bytes, transfer.state_bytes)
+        off = link.arena.alloc(total) if link is not None else None
+        if off is None:
+            status = "busy" if link is not None and total <= link.arena.nbytes else "error"
+            conn.sendall(json.dumps({"status": status}).encode() + b"\n")
+            return
+        try:
+            conn.sendall(json.dumps({"status": "ok", "addr": link.arena.base + off, "segment": link.segment}).encode() + b"\n")
+            go = bytearray(2)
+            _recv_exact(conn, memoryview(go))
+            if go != b"go":
+                raise RuntimeError(f"PE aborted the transfer: {bytes(go)!r}")
+        except BaseException:
+            link.arena.release(off, total)
+            raise
+        pages = link.arena.view(off, num_pages * transfer.page_bytes).view(num_pages, transfer.page_bytes)
+        state = link.arena.view(off + state_off, transfer.state_bytes)
+        recv_q.put((rid, first_token, prompt_len, pages, state, (off, total)))
+        conn.sendall(b"ok")
+
+    def receive_one(conn: socket.socket):
+        try:
+            with conn:
+                hdr = bytearray(HDR.size)
+                _recv_exact(conn, memoryview(hdr)[: len(MC_MAGIC)])
+                mooncake = bytes(hdr[: len(MC_MAGIC)]) == MC_MAGIC
+                _recv_exact(conn, memoryview(hdr)[0 if mooncake else len(MC_MAGIC) :])
+                rid_b, first_token, prompt_len, num_pages, page_bytes, state_bytes = HDR.unpack(hdr)
+                rid = rid_b.decode()
+                assert page_bytes == transfer.page_bytes and state_bytes == transfer.state_bytes, "layout mismatch between PE and CE"
+                if mooncake:
+                    receive_mooncake(conn, rid, first_token, prompt_len, num_pages)
+                    return
+                pages = torch.empty((num_pages, page_bytes), dtype=torch.uint8, pin_memory=True)
+                state = torch.empty(state_bytes, dtype=torch.uint8, pin_memory=True)
+                _recv_exact(conn, memoryview(pages.numpy()).cast("B"))
+                _recv_exact(conn, memoryview(state.numpy()).cast("B"))
+                recv_q.put((rid, first_token, prompt_len, pages, state, None))
+                conn.sendall(b"ok")
+        except Exception as e:  # noqa: BLE001
+            print(f"[ce rank {dp_rank}] transfer failed: {e!r}", flush=True)
+
+    threading.Thread(target=receiver, daemon=True).start()
+    out_q.put(("ready", dp_rank, None))
+
+    expected: dict[str, dict] = {}  # rid -> {"prompt_token_ids", "sampling"} from the frontend
+    arrived: dict[str, tuple] = {}  # rid -> transfer
+    pending: collections.deque = collections.deque()  # rids with both parts, waiting for GPU allocation
+    keep: dict[str, tuple] = {}  # rid -> (admit event, pages, state, arena region or None)
+
+    def release(region: tuple[int, int] | None) -> None:
+        if region is not None:
+            link.arena.release(*region)
+
+    inflight = False
+    while True:
+        while True:
+            try:
+                kind, rid, payload = in_q.get_nowait()
+            except queue.Empty:
+                break
+            if kind == "expect":
+                expected[rid] = payload
+                if rid in arrived:
+                    pending.append(rid)
+            elif kind == "abort":
+                expected.pop(rid, None)
+                gone = arrived.pop(rid, None)
+                if gone is not None:
+                    release(gone[4])
+                if rid in pending:
+                    pending.remove(rid)
+                else:
+                    engine.abort_request(rid)
+        while True:
+            try:
+                rid, first_token, prompt_len, pages, state, region = recv_q.get_nowait()
+            except queue.Empty:
+                break
+            arrived[rid] = (first_token, prompt_len, pages, state, region)
+            if rid in expected:
+                pending.append(rid)
+        while pending:
+            rid = pending[0]
+            spec, (first_token, prompt_len, pages, state, region) = expected[rid], arrived[rid]
+            tokens = np.asarray(list(spec["prompt_token_ids"]) + [first_token], dtype=np.int32)
+            assert len(tokens) - 1 == prompt_len, (len(tokens), prompt_len)
+            sp = SamplingParams(**spec["sampling"])
+            sp.max_tokens = max(1, sp.max_tokens - 1)  # the first token was produced on the PE
+            seq = engine.add_request(
+                user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=tokens, multi_modal_data=mm(), enqueue=False
+            )
+            if seq is None:
+                out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
+                pending.popleft()
+                expected.pop(rid, None)
+                arrived.pop(rid, None)
+                release(region)
+                continue
+            try:
+                allocated = engine.scheduler.try_allocate_seq(seq)
+            except RuntimeError as e:
+                out_q.put(("error", rid, f"cannot allocate: {e!r}"))
+                allocated = None
+            if allocated is False and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
+                break  # retry once something finishes
+            pending.popleft()
+            expected.pop(rid, None)
+            arrived.pop(rid, None)
+            if not allocated:
+                if allocated is False:
+                    out_q.put(("error", rid, "cannot allocate KV blocks for the transferred sequence"))
+                release(region)
+                continue
+            # Keep the pinned receive buffers alive until the H2D page-copy kernels have consumed them: the kernel
+            # reads them by raw pointer, so the caching host allocator would otherwise hand the memory to the next
+            # transfer (or unmap it) while the copy is still queued -> corrupted KV / cudaErrorLaunchFailure.
+            keep[rid] = (transfer.admit(seq, pages, state, prompt_len), pages, state, region)
+            if engine.multi_modal_processor is not None:
+                engine.multi_modal_processor.add(rid, seq.multi_modal_data, prompt_len)
+        for rid in [r for r, k in keep.items() if k[0].query()]:
+            release(keep.pop(rid)[3])
+
+        if engine.has_unfinished_requests():
+            outputs = engine.step()
+            inflight = True
+        elif inflight:
+            outputs = engine.flush_inflight_reqs()
+            inflight = False
+        else:
+            time.sleep(0.001)
+            continue
+        if len(outputs):
+            batch = [
+                (
+                    RequestOutput.get_request_id(p),
+                    RequestOutput.get_token_id(p),
+                    RequestOutput.get_finished(p),
+                    RequestOutput.get_finish_reason(p),
+                    RequestOutput.get_prompt_len(p),
+                )
+                for p in outputs
+            ]
+            out_q.put(("tokens", None, batch))
+
+
+# ============================================================================ frontends
+
+
+class PDFrontend(oas.Frontend):
+    def __init__(self, args: argparse.Namespace) -> None:
+        super().__init__(args)
+        ctx = mp.get_context("spawn")
+        self.out_q = ctx.Queue()
+        self.in_qs = [ctx.Queue() for _ in range(args.dp)]
+        self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)]
+
+
+def build_pe_app(fe: PDFrontend) -> FastAPI:
+    """Prefill tray: POST /prefill {rid, prompt_token_ids, sampling, kv_target:[host,port]} -> {first_token}."""
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse
+
+    app = FastAPI()
+
+    @app.on_event("startup")
+    async def _startup() -> None:
+        fe.loop = asyncio.get_running_loop()
+        threading.Thread(target=fe.dispatch_forever, daemon=True).start()
+
+    @app.get("/health")
+    async def health():
+        ok = fe.ready == fe.args.dp
+        return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready, "role": "pe"}, status_code=200 if ok else 503)
+
+    @app.get("/deepinfer/metrics")
+    async def metrics():
+        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
+        for r, st in sorted(fe.cache_stats.items()):
+            lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()]
+        return "\n".join(lines) + "\n"
+
+    @app.post("/prefill")
+    async def prefill(request: Request):
+        if fe.ready != fe.args.dp:
+            return JSONResponse({"error": "engine loading"}, status_code=503)
+        body = await request.json()
+        rid = body["rid"]
+        rank = min(range(fe.args.dp), key=lambda r: fe.load[r])
+        key = bytes.fromhex(body["affinity"]) if body.get("affinity") else None
+        if fe.args.routing == "affinity" and key is not None:
+            rank = int.from_bytes(key[:4], "little") % fe.args.dp
+        fe.load[rank] += 1
+        fe.owner[rid] = rank
+        q: asyncio.Queue = asyncio.Queue()
+        fe.streams[rid] = q
+        fe.stats["requests"] += 1
+        fe.in_qs[rank].put(
+            ("add", rid, {"prompt_token_ids": body["prompt_token_ids"], "sampling": body["sampling"], "kv_target": body["kv_target"]})
+        )
+        try:
+            item = await asyncio.wait_for(q.get(), timeout=1800)
+        finally:
+            fe.release(rid, aborted=False)
+        if item[0] == "error":
+            return JSONResponse({"error": item[1]}, status_code=500)
+        return {"first_token": item[1], "prompt_len": item[4]}
+
+    return app
+
+
+def build_ce_app(fe: PDFrontend, pe_url: str, kv_host: str) -> FastAPI:
+    """Decode tray: the user-facing OpenAI API."""
+    import httpx
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse, StreamingResponse
+
+    app = FastAPI()
+    client = httpx.AsyncClient(timeout=1800)
+
+    @app.on_event("startup")
+    async def _startup() -> None:
+        fe.loop = asyncio.get_running_loop()
+        threading.Thread(target=fe.dispatch_forever, daemon=True).start()
+
+    @app.get("/health")
+    async def health():
+        ok = fe.ready == fe.args.dp
+        return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready, "role": "ce"}, status_code=200 if ok else 503)
+
+    @app.get("/v1/models")
+    async def models():
+        return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]}
+
+    @app.get("/deepinfer/metrics")
+    async def metrics():
+        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
+        lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)]
+        return "\n".join(lines) + "\n"
+
+    @app.post("/v1/chat/completions")
+    async def chat(request: Request):
+        if fe.ready != fe.args.dp:
+            return JSONResponse({"error": {"message": "engine loading"}}, status_code=503)
+        body = await request.json()
+        prompt_ids = await asyncio.to_thread(fe.render_prompt, body)
+        max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096
+        if len(prompt_ids) + 1 > fe.args.max_model_len:
+            return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400)
+        max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids)))
+        sampling = dict(
+            temperature=float(body.get("temperature", 1.0)),
+            top_p=float(body.get("top_p", 1.0)),
+            ignore_eos=bool(body.get("ignore_eos", False)),
+            max_tokens=max_tokens,
+        )
+        if body.get("seed") is not None:
+            sampling["seq_seed"] = int(body["seed"])
+        fe.stats["prompt_tokens"] += len(prompt_ids)
+        include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
+        key = fe.affinity_key(body)
+
+        # Register with the CE rank first so the transfer can be admitted as soon as it lands.
+        rid = uuid.uuid4().hex[:16]
+        rank = (
+            int.from_bytes(key[:4], "little") % fe.args.dp
+            if (fe.args.routing == "affinity" and key)
+            else min(range(fe.args.dp), key=lambda r: fe.load[r])
+        )
+        fe.load[rank] += 1
+        fe.owner[rid] = rank
+        q: asyncio.Queue = asyncio.Queue()
+        fe.streams[rid] = q
+        fe.stats["requests"] += 1
+        fe.in_qs[rank].put(("expect", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling}))
+
+        created = int(time.time())
+        cid = f"chatcmpl-{rid}"
+        detok = oas.IncrementalDetokenizer(fe.tokenizer)
+
+        def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str:
+            obj = {
+                "id": cid,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": fe.model_name,
+                "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}],
+            }
+            if usage:
+                obj["usage"] = usage
+            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"
+
+        async def prefill_first():
+            r = await client.post(
+                f"{pe_url}/prefill",
+                json={
+                    "rid": rid,
+                    "prompt_token_ids": prompt_ids,
+                    "sampling": sampling,
+                    "kv_target": [kv_host, fe.args.kv_port + rank],
+                    "affinity": key.hex() if key else None,
+                },
+            )
+            if r.status_code != 200:
+                raise RuntimeError(f"PE: {r.text[:200]}")
+            return int(r.json()["first_token"])
+
+        async def tokens():
+            first = await prefill_first()
+            text = detok.add(first)
+            fe.stats["completion_tokens"] += 1
+            if max_tokens == 1 or first == fe.tokenizer.eos_token_id and not sampling["ignore_eos"]:
+                yield text, ("length" if max_tokens == 1 else "stop")
+                return
+            yield text, None
+            while True:
+                item = await q.get()
+                if item[0] == "error":
+                    raise RuntimeError(item[1])
+                _, token_id, finished, reason, _ = item
+                text = detok.add(token_id)
+                fe.stats["completion_tokens"] += 1
+                if finished:
+                    yield text, ("length" if detok.num_tokens >= max_tokens else "stop")
+                    return
+                yield text, None
+
+        usage = lambda: {
+            "prompt_tokens": len(prompt_ids),
+            "completion_tokens": detok.num_tokens,  # noqa: E731
+            "total_tokens": len(prompt_ids) + detok.num_tokens,
+        }
+
+        if body.get("stream"):
+
+            async def sse():
+                done = False
+                try:
+                    yield chunk({"role": "assistant", "content": ""})
+                    async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
+                        if text or finish:
+                            yield chunk({"content": text} if text else {}, finish)
+                    done = True
+                    if include_usage:
+                        yield chunk({}, usage=usage())
+                    yield "data: [DONE]\n\n"
+                finally:
+                    fe.release(rid, aborted=not done)
+
+            return StreamingResponse(sse(), media_type="text/event-stream")
+
+        done = False
+        try:
+            parts, finish = [], None
+            async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
+                parts.append(text)
+            done = True
+        finally:
+            fe.release(rid, aborted=not done)
+        return {
+            "id": cid,
+            "object": "chat.completion",
+            "created": created,
+            "model": fe.model_name,
+            "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}],
+            "usage": usage(),
+        }
+
+    return app
+
+
+def main() -> None:
+    args = parse_args()
+    args.model = str(Path(args.model).expanduser().resolve())
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("RANK", "0")
+    fe = PDFrontend(args)
+    fe.start()
+    import uvicorn
+
+    if args.role == "pe":
+        app = build_pe_app(fe)
+    else:
+        assert args.pe_url, "--pe-url is required for the decode tray"
+        app = build_ce_app(fe, args.pe_url.rstrip("/"), args.kv_host or socket.gethostname())
+    uvicorn.run(app, host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900)
+
+
+if __name__ == "__main__":
+    main()

serve/prompt_tokenizer.py

@@ -0,0 +1,264 @@
+"""Prompt tokenization for the serving frontends: exact token ids, much faster than `tokenizer.encode` on long prompts.
+
+Agentic traffic re-sends the whole conversation every turn, so the frontend re-tokenizes 100K+ token prompts on
+nearly every request. HF `tokenizer.encode` takes ~1.3 ms per 1K tokens (GB300 Grace core), which becomes the largest
+part of TTFT once the prefix cache hits. Backends, all required to produce the ids of
+`tokenizer.encode(text, add_special_tokens=False)`:
+
+- `fastokens` (optional `pip install fastokens`): a faster Rust implementation of the HF pipeline, built from the
+  tokenizer exactly as transformers configured it (`backend_tokenizer.to_str()`, not the file on disk).
+- `segment`: no extra dependency. HF tokenizers first split the input at non-normalized added tokens and encode the
+  text between them independently. Splitting there ourselves and caching each piece's ids by its text is exact, and
+  a new turn only pays for the pieces it adds (the new tool result, the new user message). It is enabled only when
+  the tokenizer's config shows the split is exact (see `_segment_split_tokens`).
+- `hf`: `tokenizer.encode`.
+
+`auto` takes the first of fastokens, segment, hf that is available and matches `tokenizer.encode` on a set of probes
+built from the tokenizer's own added tokens. At runtime a sample of prompts is re-encoded with `tokenizer.encode`
+on a background thread; the first mismatch switches the frontend to `hf`.
+"""
+
+from __future__ import annotations
+
+import array
+import collections
+import contextlib
+import json
+import queue
+import random
+import re
+import threading
+import time
+from typing import Any
+
+BACKENDS = ("auto", "fastokens", "segment", "hf")
+
+
+class PromptTokenizer:
+    def __init__(
+        self,
+        tokenizer: Any,
+        backend: str = "auto",
+        verify_rate: float = 0.005,
+        cache_bytes: int = 512 << 20,
+        extra_probes: list[str] | None = None,
+    ) -> None:
+        assert backend in BACKENDS, backend
+        self.tokenizer = tokenizer
+        self.verify_rate = verify_rate
+        self.stats = dict(calls=0, errors=0, verified=0, verify_mismatches=0)
+        self.notes: list[str] = []
+        self._times: collections.deque[float] = collections.deque(maxlen=4096)
+        self._probes = _probes(tokenizer) + list(extra_probes or [])
+
+        candidates = ("fastokens", "segment") if backend == "auto" else (() if backend == "hf" else (backend,))
+        self.name, self._encode = "hf", self._hf_encode
+        for name in candidates:
+            try:
+                encode = _FastokensBackend(tokenizer).encode if name == "fastokens" else _SegmentBackend(tokenizer, cache_bytes).encode
+                if (bad := self._first_mismatch(encode)) is not None:
+                    raise ValueError(f"ids differ from tokenizer.encode on probe {bad!r}")
+            except Exception as e:  # noqa: BLE001 - a backend that cannot be built or fails parity is skipped
+                self.notes.append(f"{name}: {type(e).__name__}: {e}"[:300])
+                if backend != "auto":
+                    raise RuntimeError(f"tokenizer backend {name} unavailable: {self.notes[-1]}") from e
+                continue
+            self.name, self._encode = name, encode
+            break
+
+        self._verify_q: queue.Queue[tuple[str, list[int]]] = queue.Queue(maxsize=8)
+        if self.name != "hf" and verify_rate > 0:
+            threading.Thread(target=self._verify_forever, name="prompt-tokenizer-verify", daemon=True).start()
+
+    def encode(self, text: str) -> list[int]:
+        t0 = time.perf_counter()
+        name = self.name
+        try:
+            ids = self._encode(text)
+        except Exception:  # noqa: BLE001 - e.g. a regex engine limit on pathological input
+            self.stats["errors"] += 1
+            name, ids = "hf", self._hf_encode(text)
+        self._times.append(time.perf_counter() - t0)
+        self.stats["calls"] += 1
+        if name != "hf" and self.verify_rate > 0 and random.random() < self.verify_rate:
+            with contextlib.suppress(queue.Full):  # the checker is behind: skip this sample
+                self._verify_q.put_nowait((text, ids))
+        return ids
+
+    def metrics(self) -> dict[str, float]:
+        """Counters plus encode latency percentiles (ms) over the last 4096 calls."""
+        out: dict[str, float] = dict(self.stats)
+        times = sorted(self._times)
+        for q in (50, 90, 99):
+            out[f"encode_ms_p{q}"] = round(times[min(len(times) - 1, len(times) * q // 100)] * 1e3, 3) if times else 0.0
+        return out
+
+    def describe(self) -> str:
+        skipped = "; ".join(self.notes)
+        return f"prompt tokenizer backend: {self.name}" + (f" (skipped {skipped})" if skipped else "")
+
+    def _hf_encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    def _first_mismatch(self, encode: Any) -> str | None:
+        for p in self._probes:
+            if list(encode(p)) != self._hf_encode(p):
+                return p[:80]
+        return None
+
+    def _verify_forever(self) -> None:
+        while True:
+            text, ids = self._verify_q.get()
+            ref = self._hf_encode(text)
+            self.stats["verified"] += 1
+            if ids != ref:
+                self.stats["verify_mismatches"] += 1
+                at = next((i for i, (a, b) in enumerate(zip(ids, ref, strict=False)) if a != b), min(len(ids), len(ref)))
+                print(
+                    f"[server] prompt tokenizer {self.name} differs from tokenizer.encode at token {at} of {len(ref)} "
+                    f"(text {text[:60]!r}...); switching to hf",
+                    flush=True,
+                )
+                self.name, self._encode = "hf", self._hf_encode
+
+
+class _FastokensBackend:
+    def __init__(self, tokenizer: Any) -> None:
+        from fastokens._native import Tokenizer  # optional dependency
+
+        if not getattr(tokenizer, "is_fast", False):
+            raise ValueError("needs a fast (tokenizer.json) tokenizer")
+        self._tok = Tokenizer.from_json_str(tokenizer.backend_tokenizer.to_str())
+
+    def encode(self, text: str) -> list[int]:
+        return self._tok.encode(text, add_special_tokens=False).ids
+
+
+class _SegmentBackend:
+    """Split at added tokens, encode only the pieces not seen before (one batched call), LRU-cache pieces by text."""
+
+    def __init__(self, tokenizer: Any, cache_bytes: int) -> None:
+        if not getattr(tokenizer, "is_fast", False):
+            raise ValueError("needs a fast (tokenizer.json) tokenizer")
+        self._backend = tokenizer.backend_tokenizer
+        split = _segment_split_tokens(self._backend)
+        # Longest first: Python's alternation is leftmost-first, HF's matcher is leftmost-longest. lstrip/rstrip
+        # tokens swallow the whitespace next to them, as in HF (Unicode White_Space, Rust's char::is_whitespace).
+        alts = [
+            f"{_WS}*" * t.lstrip + re.escape(t.content) + f"{_WS}*" * t.rstrip
+            for t in sorted(split.values(), key=lambda t: -len(t.content))
+        ]
+        self._split = re.compile("(" + "|".join(alts) + ")")
+        self._ids = {c: t.id for c, t in split.items()}
+        self._stripped = any(t.lstrip or t.rstrip for t in split.values())
+        # ids as uint32 arrays: 4 B/token instead of ~40 B for a list of Python ints
+        self._cache: collections.OrderedDict[str, array.array] = collections.OrderedDict()
+        self._cache_bytes = 0
+        self._max_bytes = cache_bytes
+        self._lock = threading.Lock()
+
+    def encode(self, text: str) -> list[int]:
+        parts = self._split.split(text)  # odd indices are the split tokens
+        pieces: list[Any] = [None] * len(parts)
+        missing: list[int] = []
+        with self._lock:
+            for i, p in enumerate(parts):
+                if i % 2:
+                    pieces[i] = [self._ids[_strip_ws(p) if self._stripped else p]]
+                elif p:
+                    ids = self._cache.get(p)
+                    if ids is None:
+                        missing.append(i)
+                    else:
+                        self._cache.move_to_end(p)
+                        pieces[i] = ids
+        if missing:
+            encs = self._backend.encode_batch([parts[i] for i in missing], add_special_tokens=False)
+            with self._lock:
+                for i, e in zip(missing, encs, strict=True):
+                    pieces[i] = e.ids
+                    if parts[i] not in self._cache:
+                        self._cache[parts[i]] = array.array("I", e.ids)
+                        self._cache_bytes += _cost(parts[i], e.ids)
+                while self._cache_bytes > self._max_bytes and self._cache:
+                    k, v = self._cache.popitem(last=False)
+                    self._cache_bytes -= _cost(k, v)
+        out: list[int] = []
+        for ids in pieces:
+            if ids:
+                out.extend(ids)
+        return out
+
+
+_WS = "[\\t\\n\\x0b\\x0c\\r \\x85\\xa0\\u1680\\u2000-\\u200a\\u2028\\u2029\\u202f\\u205f\\u3000]"
+_WS_CHARS = "\t\n\x0b\x0c\r \x85\xa0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000"
+
+
+def _cost(text: str, ids: Any) -> int:
+    return len(text) + 4 * len(ids)
+
+
+def _strip_ws(s: str) -> str:
+    return s.strip(_WS_CHARS)
+
+
+_SplitToken = collections.namedtuple("_SplitToken", "content id lstrip rstrip")
+
+
+def _segment_split_tokens(backend: Any) -> dict[str, _SplitToken]:
+    """Added tokens at which splitting the input and encoding the pieces separately gives tokenizer.encode's ids.
+
+    HF's AddedVocabulary first matches the non-normalized added tokens on the raw input (leftmost-longest, an
+    lstrip/rstrip token also taking the whitespace next to it), and the rest of the pipeline (normalizer,
+    pre-tokenizer, model) then runs on each piece in between on its own. Raises when the config can make a piece
+    depend on its neighbours.
+    """
+    cfg = json.loads(backend.to_str())
+    added = cfg.get("added_tokens") or []
+    if any(t["single_word"] for t in added):
+        raise ValueError("single_word added tokens depend on the characters around them")
+    if _has_first_only_prefix(cfg.get("pre_tokenizer")) or _has_first_only_prefix(cfg.get("normalizer")):
+        raise ValueError("prepend_scheme='first' prefixes only the first piece of the whole input")
+    split = {t["content"]: _SplitToken(t["content"], t["id"], t["lstrip"], t["rstrip"]) for t in added if not t["normalized"]}
+    if not split:
+        raise ValueError("no non-normalized added tokens to split at")
+    if any(_strip_ws(c) != c for c in split):
+        raise ValueError("an added token starts or ends with whitespace")
+    return split
+
+
+def _has_first_only_prefix(node: Any) -> bool:
+    if isinstance(node, dict):
+        if node.get("prepend_scheme") == "first":
+            return True
+        return any(_has_first_only_prefix(v) for v in node.values())
+    if isinstance(node, list):
+        return any(_has_first_only_prefix(v) for v in node)
+    return False
+
+
+def _probes(tokenizer: Any) -> list[str]:
+    """Inputs that exercise added-token boundaries, scripts, digits and whitespace for this tokenizer."""
+    added = [t.content for t in tokenizer.added_tokens_decoder.values()]
+    rng = random.Random(0)
+    sample = added if len(added) <= 64 else added[:16] + rng.sample(added, 48)
+    texts = [
+        "",
+        "hello world",
+        " leading space",
+        "trailing space ",
+        "tabs\tand\r\nwindows\nlines\n\n\n",
+        "    indented code:\n        return x  # comment\n",
+        "数字 12345678 和中文、日本語のテキスト、カタカナ",
+        "emoji 😀👍🏽 and combining é é ñ",
+        '{"name": "read_file", "arguments": {"path": "/tmp/a b.py", "n": 3.14}}',
+        "a" * 3000,
+        " " * 500 + "x",
+        "-" * 2000,
+    ]
+    probes = list(texts)
+    for i, tok in enumerate(sample):
+        a, b = texts[(i + 1) % len(texts)][:120], texts[(i + 5) % len(texts)][:120]
+        probes += [tok, f"{a}{tok}{b}", f"{a} {tok} {b}", f"{tok}{tok}", f"\n{tok}\n"]
+    probes.append("".join(f"{t}{texts[i % len(texts)][:40]}" for i, t in enumerate(sample)))
+    return probes

serve/server.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Start/stop serve/openai_server.py as one process group, so stopping it also reaps the
+# engine's helper processes (otherwise they outlive the frontend and keep the GPUs).
+#
+#   serve/server.sh start [openai_server.py args...]
+#   serve/server.sh stop
+set -euo pipefail
+ROOT=$(cd "$(dirname "$0")/.." && pwd)
+PIDFILE="$ROOT/serve/server.pgid"
+LOG="${LOG:-$ROOT/serve.log}"
+
+case "${1:-}" in
+    start)
+        shift
+        [ -s "$PIDFILE" ] && kill -0 -- "-$(cat "$PIDFILE")" 2>/dev/null && { echo "already running (pgid $(cat "$PIDFILE"))" >&2; exit 1; }
+        [ -f "$ROOT/env.sh" ] && source "$ROOT/env.sh"  # optional local environment (venv, DEEPINFER_* vars)
+        : "${DEEPINFER_HF_MODEL_DIR:?set DEEPINFER_HF_MODEL_DIR to the local checkpoint directory}"
+        export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
+        setsid python -u "$ROOT/serve/openai_server.py" --model "$DEEPINFER_HF_MODEL_DIR" \
+            --served-model-name deepseek-ai/DeepSeek-V4.1-Flash-1017 "$@" > "$LOG" 2>&1 < /dev/null &
+        echo $! > "$PIDFILE"
+        port=8888 prev=
+        for a in "$@"; do [[ "$prev" == --port ]] && port=$a; prev=$a; done
+        until curl -fs "localhost:$port/health" >/dev/null 2>&1; do
+            kill -0 "$(cat "$PIDFILE")" 2>/dev/null || { echo "server died; see $LOG" >&2; exit 1; }
+            sleep 3
+        done
+        echo "ready (pgid $(cat "$PIDFILE"), log $LOG)"
+        ;;
+    stop)
+        [ -s "$PIDFILE" ] || { echo "not running"; exit 0; }
+        pgid=$(cat "$PIDFILE")
+        kill -TERM -- "-$pgid" 2>/dev/null || true
+        for _ in $(seq 30); do kill -0 -- "-$pgid" 2>/dev/null || break; sleep 1; done
+        kill -KILL -- "-$pgid" 2>/dev/null || true
+        rm -f "$PIDFILE"
+        echo "stopped"
+        ;;
+    *)
+        echo "usage: $0 start|stop" >&2; exit 2 ;;
+esac

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/serve/test_prompt_tokenizer.py

@@ -0,0 +1,136 @@
+"""serve/prompt_tokenizer.py: every backend must give tokenizer.encode's ids (CPU only, needs the checkpoint's tokenizer)."""
+
+from __future__ import annotations
+
+import copy
+import json
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(REPO_ROOT / "serve"))
+from prompt_tokenizer import PromptTokenizer  # noqa: E402
+
+MODEL = os.environ.get("DEEPINFER_HF_MODEL_DIR")
+pytestmark = pytest.mark.skipif(not MODEL or not Path(MODEL, "tokenizer.json").exists(), reason="DEEPINFER_HF_MODEL_DIR not set")
+
+
+@pytest.fixture(scope="module")
+def tokenizer() -> Any:
+    from transformers import AutoTokenizer
+
+    return AutoTokenizer.from_pretrained(MODEL)
+
+
+@pytest.fixture(scope="module")
+def turns() -> list[str]:
+    """A growing tool-use conversation rendered with the checkpoint's encode_messages, one prompt per turn."""
+    sys.path.insert(0, str(Path(MODEL, "encoding")))
+    from encoding import encode_messages
+
+    files = sorted((REPO_ROOT / "serve").glob("*.py")) + sorted((REPO_ROOT / "deepinfer" / "engine").glob("*.py"))[:8]
+    tools = [
+        {
+            "type": "function",
+            "function": {"name": "read_file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}},
+        }
+    ]
+    msgs = [
+        {"role": "system", "content": "You are a coding agent. 你是一个编程助手。", "tools": tools},
+        {"role": "user", "content": "Summarise serve/. 12345"},
+    ]
+    out = []
+    for i, f in enumerate(files):
+        call = {"type": "function", "function": {"name": "read_file", "arguments": json.dumps({"path": str(f.name)})}}
+        msgs.append({"role": "assistant", "content": f"Step {i}: reading.", "tool_calls": [call]})
+        msgs.append({"role": "tool", "content": f.read_text()})
+        out.append(encode_messages(copy.deepcopy(msgs), thinking_mode="chat"))
+    return out
+
+
+def available(tokenizer: Any) -> list[str]:
+    names = ["segment", "hf"]
+    try:
+        import fastokens  # noqa: F401
+
+        names.insert(0, "fastokens")
+    except ImportError:
+        pass
+    return names
+
+
+def test_backends_match_hf(tokenizer: Any, turns: list[str]) -> None:
+    for name in available(tokenizer):
+        pt = PromptTokenizer(tokenizer, backend=name, verify_rate=0)
+        assert pt.name == name
+        for prompt in turns:
+            assert pt.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False), name
+
+
+def test_auto_picks_a_fast_backend(tokenizer: Any) -> None:
+    assert PromptTokenizer(tokenizer, verify_rate=0).name in ("fastokens", "segment")
+
+
+def test_segment_cache_reuses_earlier_turns(tokenizer: Any, turns: list[str]) -> None:
+    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0)
+    hf_s = seg_s = 0.0
+    for prompt in turns:
+        t0 = time.perf_counter()
+        ref = tokenizer.encode(prompt, add_special_tokens=False)
+        t1 = time.perf_counter()
+        assert pt.encode(prompt) == ref
+        hf_s, seg_s = hf_s + t1 - t0, seg_s + time.perf_counter() - t1
+    assert seg_s < hf_s
+
+
+def test_segment_eviction_stays_exact(tokenizer: Any, turns: list[str]) -> None:
+    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0, cache_bytes=4096)
+    for prompt in turns + turns[::-1]:
+        assert pt.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False)
+
+
+def test_segment_with_strip_tokens(tokenizer: Any) -> None:
+    """Added tokens that swallow neighbouring whitespace (lstrip/rstrip, as Phi-3 and others use) stay exact."""
+    from tokenizers import AddedToken
+
+    tok = copy.deepcopy(tokenizer)
+    tok.add_tokens([AddedToken("<|l|>", lstrip=True, normalized=False), AddedToken("<|r|>", rstrip=True, normalized=False)])
+    pt = PromptTokenizer(tok, backend="segment", verify_rate=0)
+    for text in ["a  <|l|>  b", "a<|r|> \n\t b", " <|l|><|r|> ", "x　<|l|> y<|r|> z", "<|r|>   ", "   <|l|>"]:
+        assert pt.encode(text) == tok.encode(text, add_special_tokens=False), repr(text)
+
+
+def test_runtime_mismatch_switches_to_hf(tokenizer: Any) -> None:
+    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=1.0)
+    pt._encode = lambda text: [0]  # a backend gone wrong after startup
+    pt.encode("hello world")
+    for _ in range(200):
+        if pt.name == "hf":
+            break
+        time.sleep(0.01)
+    assert pt.name == "hf" and pt.stats["verify_mismatches"] == 1
+    assert pt.encode("hello world") == tokenizer.encode("hello world", add_special_tokens=False)
+
+
+def test_backend_error_falls_back_to_hf(tokenizer: Any) -> None:
+    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0)
+
+    def broken(text: str) -> list[int]:
+        raise RuntimeError("regex limit")
+
+    pt._encode = broken
+    assert pt.encode("hello") == tokenizer.encode("hello", add_special_tokens=False)
+    assert pt.stats["errors"] == 1
+
+
+def test_forced_unavailable_backend_raises(tokenizer: Any, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(sys.modules, "fastokens", None)
+    monkeypatch.setitem(sys.modules, "fastokens._native", None)
+    with pytest.raises(RuntimeError, match="fastokens"):
+        PromptTokenizer(tokenizer, backend="fastokens")
+    assert PromptTokenizer(tokenizer, verify_rate=0).name == "segment"