From 88e70d1975bad2b7544888dabb7cea406afd0654 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:10:50 +0000 Subject: [PATCH 01/19] Add a host-memory prefix cache (HiCache-style) for the offline engine The open-source engine has no KV reuse across requests, so every turn of a multi-turn conversation prefills its whole prompt. This adds deepinfer/hicache: - Pinned host pools organized in pages (one page = one 128-token KV block) plus slots for the SWA state that the production KV store persists (select_state_slice_to_store / load_state_from_store_to_prefill_cache). Pages are indexed by a hash chain over the prompt, LRU evicted, written through at the end of every prefill (engine hook pe_do_save_kvcache). - Transfers use one small CUDA kernel copying pages between GPU and pinned host memory through UVA pointers on the engine stream ("kernel" I/O method, ~180 GiB/s on GB300); completion is tracked with CUDA events. - A hit is capped so that at least yoco keep_tokens remain to recompute (bounded replay rebuilds the decoder SWA state from that suffix), then snapped to the nearest 128-aligned position with a stored state; states are stored at the last few boundaries of a prompt so short follow-ups hit. The sequence is pre-allocated, filled from host and scheduled with num_cached_tokens set (HostPrefixCache.try_resume). - Engine: read-only Sequence.block_table binding (to copy a finished prefill's blocks out) and its stub. Validated on 4x GB300 with DeepSeek-V4.1-Flash through an OpenAI-compatible server (separate PR): greedy outputs byte-identical with and without the cache over 8 multi-turn requests; InferenceX AgentX replay at concurrency 1 served 224 requests per hour vs 99, TTFT p50 0.59 s vs 8.9 s, with 89-96% of prompt tokens served from the cache. tests/hicache/test_host_cache.py covers the cache logic against a fake engine. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- csrc/engine/pybind_sequence.cpp | 2 + deepinfer/_C/engine.pyi | 3 + deepinfer/hicache/__init__.py | 3 + deepinfer/hicache/host_cache.py | 432 +++++++++++++++++++++++++++++++ deepinfer/hicache/kernels.py | 25 ++ deepinfer/hicache/page_copy.cpp | 9 + deepinfer/hicache/page_copy.cu | 45 ++++ tests/hicache/test_host_cache.py | 190 ++++++++++++++ 8 files changed, 709 insertions(+) create mode 100644 deepinfer/hicache/__init__.py create mode 100644 deepinfer/hicache/host_cache.py create mode 100644 deepinfer/hicache/kernels.py create mode 100644 deepinfer/hicache/page_copy.cpp create mode 100644 deepinfer/hicache/page_copy.cu create mode 100644 tests/hicache/test_host_cache.py diff --git a/csrc/engine/pybind_sequence.cpp b/csrc/engine/pybind_sequence.cpp index 7d628dc..cd183ec 100644 --- a/csrc/engine/pybind_sequence.cpp +++ b/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; }); diff --git a/deepinfer/_C/engine.pyi b/deepinfer/_C/engine.pyi index 91aa482..f500ad7 100644 --- a/deepinfer/_C/engine.pyi +++ b/deepinfer/_C/engine.pyi @@ -455,6 +455,9 @@ class Sequence: def cached_hit_tokens(self, arg1: typing.SupportsInt | typing.SupportsIndex) -> None: ... @property + def block_table(self) -> list[int]: + ... + @property def chunk_checkpoints(self) -> list[int]: ... @chunk_checkpoints.setter diff --git a/deepinfer/hicache/__init__.py b/deepinfer/hicache/__init__.py new file mode 100644 index 0000000..9461a6f --- /dev/null +++ b/deepinfer/hicache/__init__.py @@ -0,0 +1,3 @@ +from deepinfer.hicache.host_cache import HostPrefixCache + +__all__ = ["HostPrefixCache"] diff --git a/deepinfer/hicache/host_cache.py b/deepinfer/hicache/host_cache.py new file mode 100644 index 0000000..e073fc5 --- /dev/null +++ b/deepinfer/hicache/host_cache.py @@ -0,0 +1,432 @@ +"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache. + +Layout + Two pinned host pools, both organized in pages: + - KV pages: one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block. + - State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a + 128-token-aligned position with GPUCacheInterface.select_state_slice_to_store. + Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)), + so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool. + +Transfer + "kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through + UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never + blocks the CPU. Completion is tracked with CUDA events polled from the engine loop. + +Policy + Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last + 128-aligned position are stored (pages already present are only touched). On a new request, the longest + page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to + be recomputed, then snapped down to the nearest position that has a stored state; the sequence is + pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with + num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is + rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store). + +Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's +prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank). +""" + +from __future__ import annotations + +import hashlib +import os +import time +from collections import OrderedDict, deque +from dataclasses import dataclass, field + +import numpy as np +import torch + +from deepinfer import envs +from deepinfer.logger import init_logger +from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface + +logger = init_logger(__name__) + + +class _Entry: + __slots__ = ("key", "slot", "refs", "ready") + + def __init__(self, key: bytes, slot: int) -> None: + self.key = key + self.slot = slot + self.refs = 0 # in-flight transfers (store or load) touching this slot + self.ready = False # host copy complete and hit-able + + +@dataclass +class _Chain: + tokens: np.ndarray + hashes: list[bytes] = field(default_factory=list) + counted: bool = False # lookup stats recorded (a deferred request is looked up repeatedly) + + +@dataclass +class Stats: + lookups: int = 0 + hits: int = 0 + hit_tokens: int = 0 + lookup_tokens: int = 0 + deferred: int = 0 # hit found but GPU allocation not possible yet + stores: int = 0 + store_skipped: int = 0 + pages_stored: int = 0 + pages_loaded: int = 0 + states_stored: int = 0 + states_loaded: int = 0 + evicted_pages: int = 0 + evicted_states: int = 0 + errors: int = 0 + + def as_dict(self) -> dict[str, int]: + return dict(self.__dict__) + + +class HostPrefixCache: + def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None: + assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)" + assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1" + assert GPUCacheInterface.cache_info_ready, "cache engine not initialized" + + self.engine = engine + self.scheduler = engine.scheduler + cache_engine = engine.driver_worker.cache_engine + self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks + self.gpu_states: torch.Tensor = cache_engine.gpu_states + assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8 + self.device = self.gpu_blocks.device + self.block_size: int = engine.cache_config.block_size + assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN) + self.page_bytes: int = self.gpu_blocks.size(1) + self.state_bytes: int = GPUCacheInterface.state_size_to_disk + assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes) + self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS + # The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative. + self.min_hit_tokens = 2 * self.block_size + # A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only + # at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a + # block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA + # ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore. + self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4")) + self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference + + self.num_pages = max(0, size_bytes // self.page_bytes) + self.num_state_slots = num_state_slots + self.staging_slots = staging_slots + assert self.num_pages > 0 and num_state_slots > 0 + self.gpu_blocks_ptr = self.gpu_blocks.data_ptr() + # Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be + # cached before a prefill completes, and an unused cache then costs no pinned memory). + self._allocated = False + self.host_pages = self.host_states = None + self.host_pages_ptr = self.host_states_ptr = 0 + self._staging: list[torch.Tensor] = [] + self._staging_idx = 0 + self.kernel = None + + self.pages: OrderedDict[bytes, _Entry] = OrderedDict() + self.states: OrderedDict[bytes, _Entry] = OrderedDict() + self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1)) + self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1)) + self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque() + self._chains: dict[str, _Chain] = {} + self.stats = Stats() + self.enabled = True + # DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel). + self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1" + logger.info( + "HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d", + self.num_pages, + self.page_bytes, + self.num_pages * self.page_bytes / 2**30, + num_state_slots, + self.state_bytes, + num_state_slots * self.state_bytes / 2**30, + self.keep_tokens, + ) + + def _allocate(self) -> None: + if self._allocated: + return + from deepinfer.hicache.kernels import copy_pages_module + + t0 = time.perf_counter() + self.kernel = copy_pages_module() + self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True) + self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True) + self.host_pages_ptr = self.host_pages.data_ptr() + self.host_states_ptr = self.host_states.data_ptr() + # Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream. + self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)] + self._allocated = True + logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0) + + # ------------------------------------------------------------------ hashing + + def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None: + chain = self._chains.get(request_id) + if chain is None and tokens is not None: + chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32)) + return chain + + def _extend(self, chain: _Chain, num_pages: int) -> None: + hashes, toks, bs = chain.hashes, chain.tokens, self.block_size + prev = hashes[-1] if hashes else b"" + for i in range(len(hashes), num_pages): + prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest() + hashes.append(prev) + + def forget(self, request_id: str) -> None: + self._chains.pop(request_id, None) + + # ------------------------------------------------------------------ transfers + + def _ids(self, ids: list[int]) -> torch.Tensor: + return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True) + + def _next_staging(self) -> torch.Tensor: + t = self._staging[self._staging_idx] + self._staging_idx = (self._staging_idx + 1) % len(self._staging) + return t + + def _track(self, entries: list[_Entry]) -> None: + ev = torch.cuda.Event() + ev.record() + self._inflight.append((ev, entries)) + + def _checkpoint(self, what: str) -> None: + if self._debug_sync: + from deepinfer.utils import allow_cuda_sync + + with allow_cuda_sync(): + torch.cuda.synchronize() + logger.info("hicache debug: ok after %s", what) + + def poll(self) -> None: + """Retire completed transfers (events on one stream complete in order, so only the head is checked).""" + while self._inflight and self._inflight[0][0].query(): + _, entries = self._inflight.popleft() + for e in entries: + e.refs -= 1 + e.ready = True + + def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int: + """Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found.""" + victims, busy = [], [] + for key, e in pool.items(): + if len(victims) >= need: + break + if e.refs or not e.ready: + busy.append(key) + continue + victims.append(key) + for key in busy: + pool.move_to_end(key) + if len(victims) < need: + return 0 + for key in victims: + free.append(pool.pop(key).slot) + return len(victims) + + def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool: + if need > len(free) + len(pool): + return False + if len(free) < need: + evicted = self._evict(pool, free, need - len(free)) + if pool is self.pages: + self.stats.evicted_pages += evicted + else: + self.stats.evicted_states += evicted + return len(free) >= need + + # ------------------------------------------------------------------ load (hit) + + def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + """Fill `seq` (not yet in the scheduler) from the host cache. + + Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not + allocatable right now, sequence untouched, retry later). + """ + if not self.enabled: + return 0 + try: + return self._try_resume(seq, request_id, prompt_ids) + except Exception: # noqa: BLE001 - never take the engine down for a cache problem + logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id) + self._fail() + return 0 + + def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + bs = self.block_size + n = len(prompt_ids) + chain = self._chain(request_id, prompt_ids) + if not chain.counted: + chain.counted = True + self.stats.lookups += 1 + self.stats.lookup_tokens += n + limit = ((n - self.keep_tokens) // bs) * bs # the replay rectangle [n - keep, n) must be recomputed + if limit < self.min_hit_tokens: + return 0 + max_pages = limit // bs + self._extend(chain, max_pages) + matched = 0 + for h in chain.hashes[:max_pages]: + e = self.pages.get(h) + if e is None or not e.ready: + break + matched += 1 + b = min(matched * bs, limit) + while b >= self.min_hit_tokens: + st = self.states.get(chain.hashes[b // bs - 1]) + if st is not None and st.ready: + break + b -= bs + else: + return 0 + + try: + allocated = self.scheduler.try_allocate_seq(seq) + except RuntimeError: # can never fit (prompt longer than the whole KV pool): let the scheduler reject it + return 0 + if not allocated: + self.stats.deferred += 1 + return -1 + num_pages = b // bs + page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)] + for e in page_entries: + e.refs += 1 + self.pages.move_to_end(e.key) + st.refs += 1 + self.states.move_to_end(st.key) + + block_table = seq.block_table + self.kernel.copy_pages( + self.host_pages_ptr, + self.gpu_blocks_ptr, + self._ids([e.slot for e in page_entries]), + self._ids(list(block_table[:num_pages])), + self.page_bytes, + ) + self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})") + staging = self._next_staging() + self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes) + self._checkpoint("load: state copy H2D") + GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states) + self._checkpoint("load: load_state_to_cache") + self._track(page_entries + [st]) + + seq.num_cached_tokens = b + seq.cached_hit_tokens = b + self.stats.hits += 1 + self.stats.hit_tokens += b + self.stats.pages_loaded += num_pages + self.stats.states_loaded += 1 + return b + + # ------------------------------------------------------------------ store (write-through at prefill end) + + def on_prefill_step(self, scheduled, finished) -> None: + """Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed.""" + if not self.enabled: + return + try: + for seq in scheduled: + if seq.chunked_rest_len == 0: # last chunk of this prompt + self._store(seq) + except Exception: # noqa: BLE001 + logger.exception("host prefix cache: store failed") + self._fail() + + def _store(self, seq) -> None: + chain = self._chains.get(seq.request_id) + if chain is None: # request did not go through try_resume: hash the prompt from the sequence itself + chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32)) + self._allocate() + bs = self.block_size + b = (seq.get_len() // bs) * bs + if b < self.min_hit_tokens: + return + num_pages = b // bs + self._extend(chain, num_pages) + missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages] + for i in range(num_pages): + if i not in missing: + self.pages.move_to_end(chain.hashes[i]) + # State boundaries: b, b - bs, ... while the window ending there is still in the ring. + seq_len = seq.get_len() + new_states = [] + for k in range(self.state_lookback_blocks): + pos = b - k * bs + if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack: + break + key = chain.hashes[pos // bs - 1] + if key in self.states: + self.states.move_to_end(key) + else: + new_states.append((pos, key)) + if not missing and not new_states: + return + if not self._reserve(self.pages, self.free_pages, len(missing)): + self.stats.store_skipped += 1 + return + # States are best effort: keep the newest boundaries that fit. + while new_states and not self._reserve(self.states, self.free_states, len(new_states)): + new_states.pop() + if not missing and not new_states: + self.stats.store_skipped += 1 + return + + block_table = seq.block_table + entries: list[_Entry] = [] + if missing: + new = [] + for i in missing: + e = _Entry(chain.hashes[i], self.free_pages.pop()) + e.refs = 1 + self.pages[e.key] = e + new.append(e) + self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})") + self.kernel.copy_pages( + self.gpu_blocks_ptr, + self.host_pages_ptr, + self._ids([block_table[i] for i in missing]), + self._ids([e.slot for e in new]), + self.page_bytes, + ) + self._checkpoint("store: page copy D2H") + entries += new + self.stats.pages_stored += len(new) + for pos, key in new_states: + st = _Entry(key, self.free_states.pop()) + st.refs = 1 + self.states[key] = st + staging = self._next_staging() + GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging) + self._checkpoint(f"store: select_state_slice at {pos}") + self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes) + self._checkpoint("store: state copy D2H") + entries.append(st) + self.stats.states_stored += 1 + self._track(entries) + self.stats.stores += 1 + + # ------------------------------------------------------------------ misc + + def _fail(self) -> None: + self.stats.errors += 1 + if self.stats.errors >= 3: + logger.error("host prefix cache: disabled after %d errors", self.stats.errors) + self.enabled = False + + def stats_dict(self) -> dict[str, int | float]: + d = self.stats.as_dict() + d.update( + enabled=int(self.enabled), + pages_used=len(self.pages), + pages_total=self.num_pages, + states_used=len(self.states), + states_total=len(self.free_states) + len(self.states), + inflight=len(self._inflight), + hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0, + ) + return d diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py new file mode 100644 index 0000000..0384de7 --- /dev/null +++ b/deepinfer/hicache/kernels.py @@ -0,0 +1,25 @@ +"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions).""" + +from __future__ import annotations + +from pathlib import Path + +_module = None + + +def copy_pages_module(): + global _module + if _module is None: + from torch.utils.cpp_extension import load + + # Do not touch TORCH_CUDA_ARCH_LIST here: the engine's own JIT kernels (deep_jit / TileLang) run + # in this process and must keep their arch-specific (sm_1xxa) targets. torch's loader derives the + # arch for this trivial kernel from the visible device. + here = Path(__file__).resolve().parent + _module = load( + name="deepinfer_hicache_page_copy", + sources=[str(here / "page_copy.cpp"), str(here / "page_copy.cu")], + extra_cuda_cflags=["-O3"], + verbose=False, + ) + return _module diff --git a/deepinfer/hicache/page_copy.cpp b/deepinfer/hicache/page_copy.cpp new file mode 100644 index 0000000..4a20e78 --- /dev/null +++ b/deepinfer/hicache/page_copy.cpp @@ -0,0 +1,9 @@ +#include + +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", ©_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."); +} diff --git a/deepinfer/hicache/page_copy.cu b/deepinfer/hicache/page_copy.cu new file mode 100644 index 0000000..a445421 --- /dev/null +++ b/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 +#include +#include +#include + +#include +#include + +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(src_base + static_cast(src_ids[blockIdx.x]) * page_bytes); + int4 *dst = reinterpret_cast(dst_base + static_cast(dst_ids[blockIdx.x]) * page_bytes); + const int64_t stride = static_cast(gridDim.y) * blockDim.x; + for (int64_t i = static_cast(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(std::min((n16 + threads - 1) / threads, 16)); + const dim3 grid(static_cast(n), blocks_y); + auto stream = c10::cuda::getCurrentCUDAStream(); + copy_pages_kernel<<>>(reinterpret_cast(src_base), + reinterpret_cast(dst_base), src_ids.data_ptr(), dst_ids.data_ptr(), page_bytes); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tests/hicache/test_host_cache.py b/tests/hicache/test_host_cache.py new file mode 100644 index 0000000..77ffc4e --- /dev/null +++ b/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 From eb1c7a51ab1cafecaf4e74006649924be7eb129e Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:11:38 +0000 Subject: [PATCH 02/19] Add a minimal OpenAI-compatible chat server over the offline engine serve/openai_server.py exposes /v1/chat/completions (streaming, tools, usage) on top of LLMEngine: one engine process per GPU (attention DP + MoE EP, the layout of examples/benchmark.py), prompts rendered with the checkpoint's own encoding.py, conversation-affinity routing, a watchdog that fails open requests if an engine rank dies, and /deepinfer/metrics. serve/server.sh runs it as a process group so a stop also reaps the engine's helper processes, which otherwise keep the GPUs. It integrates the host prefix cache from deepinfer/hicache when present (--hicache-size-gb, --hicache-state-slots, per-request "prefix_cache": false bypass; serve/check_prefix_cache.py is the A/B check) and serves without it otherwise. It also sets DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL to 2 by default (--prefill-any-interval): with the engine default of 100, a rank whose peers are idle admits a waiting prompt only every 100 decode steps, a ~1.4 s TTFT floor on every request in a serving setting. Used to drive the InferenceX AgentX replay (aiperf) on 4x GB300. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/README.md | 45 ++++ serve/check_prefix_cache.py | 87 ++++++ serve/openai_server.py | 518 ++++++++++++++++++++++++++++++++++++ serve/server.sh | 41 +++ 4 files changed, 691 insertions(+) create mode 100644 serve/README.md create mode 100644 serve/check_prefix_cache.py create mode 100644 serve/openai_server.py create mode 100755 serve/server.sh diff --git a/serve/README.md b/serve/README.md new file mode 100644 index 0000000..9f204eb --- /dev/null +++ b/serve/README.md @@ -0,0 +1,45 @@ +# 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). + +## 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; `prefill-any-interval` matters there, since a single lane keeps most ranks idle. diff --git a/serve/check_prefix_cache.py b/serve/check_prefix_cache.py new file mode 100644 index 0000000..bb8fae1 --- /dev/null +++ b/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) diff --git a/serve/openai_server.py b/serve/openai_server.py new file mode 100644 index 0000000..d34acdc --- /dev/null +++ b/serve/openai_server.py @@ -0,0 +1,518 @@ +"""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 + +# Module level so FastAPI can resolve the string annotation `request: Request`. +from fastapi import Request + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "examples")) + + +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=2, + 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 (engine default 100 => ~1.4 s TTFT floor)", + ) + 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)) + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + + 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, seq, num_cached_tokens): + # 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) -> 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("�"): + 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) + 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) -> None: + q = self.streams.get(rid) + if q is not None: + q.put_nowait(item) + + def _deliver_batch(self, batch) -> 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｜>" + return self.tokenizer.encode(prompt, add_special_tokens=False) + + @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): + 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()] + 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(): + 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(): + 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() diff --git a/serve/server.sh b/serve/server.sh new file mode 100755 index 0000000..b612fd2 --- /dev/null +++ b/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 From aaf1e62c9eab1ce5bd260a594dfaea613760e89f Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:26:07 +0000 Subject: [PATCH 03/19] serve: type annotations and B007 for ruff Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- serve/openai_server.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/serve/openai_server.py b/serve/openai_server.py index d34acdc..1e41bf2 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -26,9 +26,13 @@ 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 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")) @@ -123,7 +127,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m def mm_data(): return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None - def enqueue(rid, seq, num_cached_tokens): + 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) @@ -228,7 +232,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m class IncrementalDetokenizer: """vLLM-style prefix/read offsets so each delta decodes only a short tail.""" - def __init__(self, tokenizer) -> None: + def __init__(self, tokenizer: Any) -> None: self.tok = tokenizer self.ids: list[int] = [] self.num_tokens = 0 @@ -300,12 +304,12 @@ class Frontend: else: self.loop.call_soon_threadsafe(self._deliver_batch, payload) - def _deliver(self, rid: str, item) -> None: + 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) -> None: + 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: @@ -375,7 +379,7 @@ class Frontend: self.stats["finished"] += 1 -def build_app(fe: Frontend): +def build_app(fe: Frontend) -> FastAPI: from fastapi import FastAPI from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse @@ -469,7 +473,7 @@ def build_app(fe: Frontend): done = False try: yield chunk({"role": "assistant", "content": ""}) - async for text, finish in tokens(): + 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 @@ -484,7 +488,7 @@ def build_app(fe: Frontend): done = False try: parts, finish = [], None - async for text, finish in tokens(): + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) parts.append(text) done = True finally: From e1f71ecdf8f1e53224bf44eee6989dee9de5f28a Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:26:07 +0000 Subject: [PATCH 04/19] hicache: type annotations for ruff ANN rules Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- deepinfer/hicache/host_cache.py | 15 ++++++++++----- deepinfer/hicache/kernels.py | 3 ++- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/deepinfer/hicache/host_cache.py b/deepinfer/hicache/host_cache.py index e073fc5..40e5ff2 100644 --- a/deepinfer/hicache/host_cache.py +++ b/deepinfer/hicache/host_cache.py @@ -33,6 +33,7 @@ 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 @@ -41,6 +42,10 @@ 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__) @@ -83,7 +88,7 @@ class Stats: class HostPrefixCache: - def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None: + 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" @@ -241,7 +246,7 @@ class HostPrefixCache: # ------------------------------------------------------------------ load (hit) - def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + 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 @@ -256,7 +261,7 @@ class HostPrefixCache: self._fail() return 0 - def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + 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) @@ -325,7 +330,7 @@ class HostPrefixCache: # ------------------------------------------------------------------ store (write-through at prefill end) - def on_prefill_step(self, scheduled, finished) -> None: + 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 @@ -337,7 +342,7 @@ class HostPrefixCache: logger.exception("host prefix cache: store failed") self._fail() - def _store(self, seq) -> None: + 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)) diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py index 0384de7..9cc372c 100644 --- a/deepinfer/hicache/kernels.py +++ b/deepinfer/hicache/kernels.py @@ -3,11 +3,12 @@ from __future__ import annotations from pathlib import Path +from types import ModuleType _module = None -def copy_pages_module(): +def copy_pages_module() -> ModuleType: global _module if _module is None: from torch.utils.cpp_extension import load From ee2fe708ef30107e10e01dc0e3bf2b1498f03e30 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:10:50 +0000 Subject: [PATCH 05/19] Add a host-memory prefix cache (HiCache-style) for the offline engine The open-source engine has no KV reuse across requests, so every turn of a multi-turn conversation prefills its whole prompt. This adds deepinfer/hicache: - Pinned host pools organized in pages (one page = one 128-token KV block) plus slots for the SWA state that the production KV store persists (select_state_slice_to_store / load_state_from_store_to_prefill_cache). Pages are indexed by a hash chain over the prompt, LRU evicted, written through at the end of every prefill (engine hook pe_do_save_kvcache). - Transfers use one small CUDA kernel copying pages between GPU and pinned host memory through UVA pointers on the engine stream ("kernel" I/O method, ~180 GiB/s on GB300); completion is tracked with CUDA events. - A hit is capped so that at least yoco keep_tokens remain to recompute (bounded replay rebuilds the decoder SWA state from that suffix), then snapped to the nearest 128-aligned position with a stored state; states are stored at the last few boundaries of a prompt so short follow-ups hit. The sequence is pre-allocated, filled from host and scheduled with num_cached_tokens set (HostPrefixCache.try_resume). - Engine: read-only Sequence.block_table binding (to copy a finished prefill's blocks out) and its stub. Validated on 4x GB300 with DeepSeek-V4.1-Flash through an OpenAI-compatible server (separate PR): greedy outputs byte-identical with and without the cache over 8 multi-turn requests; InferenceX AgentX replay at concurrency 1 served 224 requests per hour vs 99, TTFT p50 0.59 s vs 8.9 s, with 89-96% of prompt tokens served from the cache. tests/hicache/test_host_cache.py covers the cache logic against a fake engine. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- csrc/engine/pybind_sequence.cpp | 2 + deepinfer/_C/engine.pyi | 3 + deepinfer/hicache/__init__.py | 3 + deepinfer/hicache/host_cache.py | 432 +++++++++++++++++++++++++++++++ deepinfer/hicache/kernels.py | 25 ++ deepinfer/hicache/page_copy.cpp | 9 + deepinfer/hicache/page_copy.cu | 45 ++++ tests/hicache/test_host_cache.py | 190 ++++++++++++++ 8 files changed, 709 insertions(+) create mode 100644 deepinfer/hicache/__init__.py create mode 100644 deepinfer/hicache/host_cache.py create mode 100644 deepinfer/hicache/kernels.py create mode 100644 deepinfer/hicache/page_copy.cpp create mode 100644 deepinfer/hicache/page_copy.cu create mode 100644 tests/hicache/test_host_cache.py diff --git a/csrc/engine/pybind_sequence.cpp b/csrc/engine/pybind_sequence.cpp index 7d628dc..cd183ec 100644 --- a/csrc/engine/pybind_sequence.cpp +++ b/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; }); diff --git a/deepinfer/_C/engine.pyi b/deepinfer/_C/engine.pyi index 91aa482..f500ad7 100644 --- a/deepinfer/_C/engine.pyi +++ b/deepinfer/_C/engine.pyi @@ -455,6 +455,9 @@ class Sequence: def cached_hit_tokens(self, arg1: typing.SupportsInt | typing.SupportsIndex) -> None: ... @property + def block_table(self) -> list[int]: + ... + @property def chunk_checkpoints(self) -> list[int]: ... @chunk_checkpoints.setter diff --git a/deepinfer/hicache/__init__.py b/deepinfer/hicache/__init__.py new file mode 100644 index 0000000..9461a6f --- /dev/null +++ b/deepinfer/hicache/__init__.py @@ -0,0 +1,3 @@ +from deepinfer.hicache.host_cache import HostPrefixCache + +__all__ = ["HostPrefixCache"] diff --git a/deepinfer/hicache/host_cache.py b/deepinfer/hicache/host_cache.py new file mode 100644 index 0000000..e073fc5 --- /dev/null +++ b/deepinfer/hicache/host_cache.py @@ -0,0 +1,432 @@ +"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache. + +Layout + Two pinned host pools, both organized in pages: + - KV pages: one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block. + - State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a + 128-token-aligned position with GPUCacheInterface.select_state_slice_to_store. + Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)), + so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool. + +Transfer + "kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through + UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never + blocks the CPU. Completion is tracked with CUDA events polled from the engine loop. + +Policy + Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last + 128-aligned position are stored (pages already present are only touched). On a new request, the longest + page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to + be recomputed, then snapped down to the nearest position that has a stored state; the sequence is + pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with + num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is + rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store). + +Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's +prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank). +""" + +from __future__ import annotations + +import hashlib +import os +import time +from collections import OrderedDict, deque +from dataclasses import dataclass, field + +import numpy as np +import torch + +from deepinfer import envs +from deepinfer.logger import init_logger +from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface + +logger = init_logger(__name__) + + +class _Entry: + __slots__ = ("key", "slot", "refs", "ready") + + def __init__(self, key: bytes, slot: int) -> None: + self.key = key + self.slot = slot + self.refs = 0 # in-flight transfers (store or load) touching this slot + self.ready = False # host copy complete and hit-able + + +@dataclass +class _Chain: + tokens: np.ndarray + hashes: list[bytes] = field(default_factory=list) + counted: bool = False # lookup stats recorded (a deferred request is looked up repeatedly) + + +@dataclass +class Stats: + lookups: int = 0 + hits: int = 0 + hit_tokens: int = 0 + lookup_tokens: int = 0 + deferred: int = 0 # hit found but GPU allocation not possible yet + stores: int = 0 + store_skipped: int = 0 + pages_stored: int = 0 + pages_loaded: int = 0 + states_stored: int = 0 + states_loaded: int = 0 + evicted_pages: int = 0 + evicted_states: int = 0 + errors: int = 0 + + def as_dict(self) -> dict[str, int]: + return dict(self.__dict__) + + +class HostPrefixCache: + def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None: + assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)" + assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1" + assert GPUCacheInterface.cache_info_ready, "cache engine not initialized" + + self.engine = engine + self.scheduler = engine.scheduler + cache_engine = engine.driver_worker.cache_engine + self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks + self.gpu_states: torch.Tensor = cache_engine.gpu_states + assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8 + self.device = self.gpu_blocks.device + self.block_size: int = engine.cache_config.block_size + assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN) + self.page_bytes: int = self.gpu_blocks.size(1) + self.state_bytes: int = GPUCacheInterface.state_size_to_disk + assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes) + self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS + # The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative. + self.min_hit_tokens = 2 * self.block_size + # A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only + # at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a + # block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA + # ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore. + self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4")) + self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference + + self.num_pages = max(0, size_bytes // self.page_bytes) + self.num_state_slots = num_state_slots + self.staging_slots = staging_slots + assert self.num_pages > 0 and num_state_slots > 0 + self.gpu_blocks_ptr = self.gpu_blocks.data_ptr() + # Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be + # cached before a prefill completes, and an unused cache then costs no pinned memory). + self._allocated = False + self.host_pages = self.host_states = None + self.host_pages_ptr = self.host_states_ptr = 0 + self._staging: list[torch.Tensor] = [] + self._staging_idx = 0 + self.kernel = None + + self.pages: OrderedDict[bytes, _Entry] = OrderedDict() + self.states: OrderedDict[bytes, _Entry] = OrderedDict() + self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1)) + self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1)) + self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque() + self._chains: dict[str, _Chain] = {} + self.stats = Stats() + self.enabled = True + # DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel). + self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1" + logger.info( + "HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d", + self.num_pages, + self.page_bytes, + self.num_pages * self.page_bytes / 2**30, + num_state_slots, + self.state_bytes, + num_state_slots * self.state_bytes / 2**30, + self.keep_tokens, + ) + + def _allocate(self) -> None: + if self._allocated: + return + from deepinfer.hicache.kernels import copy_pages_module + + t0 = time.perf_counter() + self.kernel = copy_pages_module() + self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True) + self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True) + self.host_pages_ptr = self.host_pages.data_ptr() + self.host_states_ptr = self.host_states.data_ptr() + # Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream. + self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)] + self._allocated = True + logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0) + + # ------------------------------------------------------------------ hashing + + def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None: + chain = self._chains.get(request_id) + if chain is None and tokens is not None: + chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32)) + return chain + + def _extend(self, chain: _Chain, num_pages: int) -> None: + hashes, toks, bs = chain.hashes, chain.tokens, self.block_size + prev = hashes[-1] if hashes else b"" + for i in range(len(hashes), num_pages): + prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest() + hashes.append(prev) + + def forget(self, request_id: str) -> None: + self._chains.pop(request_id, None) + + # ------------------------------------------------------------------ transfers + + def _ids(self, ids: list[int]) -> torch.Tensor: + return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True) + + def _next_staging(self) -> torch.Tensor: + t = self._staging[self._staging_idx] + self._staging_idx = (self._staging_idx + 1) % len(self._staging) + return t + + def _track(self, entries: list[_Entry]) -> None: + ev = torch.cuda.Event() + ev.record() + self._inflight.append((ev, entries)) + + def _checkpoint(self, what: str) -> None: + if self._debug_sync: + from deepinfer.utils import allow_cuda_sync + + with allow_cuda_sync(): + torch.cuda.synchronize() + logger.info("hicache debug: ok after %s", what) + + def poll(self) -> None: + """Retire completed transfers (events on one stream complete in order, so only the head is checked).""" + while self._inflight and self._inflight[0][0].query(): + _, entries = self._inflight.popleft() + for e in entries: + e.refs -= 1 + e.ready = True + + def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int: + """Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found.""" + victims, busy = [], [] + for key, e in pool.items(): + if len(victims) >= need: + break + if e.refs or not e.ready: + busy.append(key) + continue + victims.append(key) + for key in busy: + pool.move_to_end(key) + if len(victims) < need: + return 0 + for key in victims: + free.append(pool.pop(key).slot) + return len(victims) + + def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool: + if need > len(free) + len(pool): + return False + if len(free) < need: + evicted = self._evict(pool, free, need - len(free)) + if pool is self.pages: + self.stats.evicted_pages += evicted + else: + self.stats.evicted_states += evicted + return len(free) >= need + + # ------------------------------------------------------------------ load (hit) + + def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + """Fill `seq` (not yet in the scheduler) from the host cache. + + Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not + allocatable right now, sequence untouched, retry later). + """ + if not self.enabled: + return 0 + try: + return self._try_resume(seq, request_id, prompt_ids) + except Exception: # noqa: BLE001 - never take the engine down for a cache problem + logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id) + self._fail() + return 0 + + def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + bs = self.block_size + n = len(prompt_ids) + chain = self._chain(request_id, prompt_ids) + if not chain.counted: + chain.counted = True + self.stats.lookups += 1 + self.stats.lookup_tokens += n + limit = ((n - self.keep_tokens) // bs) * bs # the replay rectangle [n - keep, n) must be recomputed + if limit < self.min_hit_tokens: + return 0 + max_pages = limit // bs + self._extend(chain, max_pages) + matched = 0 + for h in chain.hashes[:max_pages]: + e = self.pages.get(h) + if e is None or not e.ready: + break + matched += 1 + b = min(matched * bs, limit) + while b >= self.min_hit_tokens: + st = self.states.get(chain.hashes[b // bs - 1]) + if st is not None and st.ready: + break + b -= bs + else: + return 0 + + try: + allocated = self.scheduler.try_allocate_seq(seq) + except RuntimeError: # can never fit (prompt longer than the whole KV pool): let the scheduler reject it + return 0 + if not allocated: + self.stats.deferred += 1 + return -1 + num_pages = b // bs + page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)] + for e in page_entries: + e.refs += 1 + self.pages.move_to_end(e.key) + st.refs += 1 + self.states.move_to_end(st.key) + + block_table = seq.block_table + self.kernel.copy_pages( + self.host_pages_ptr, + self.gpu_blocks_ptr, + self._ids([e.slot for e in page_entries]), + self._ids(list(block_table[:num_pages])), + self.page_bytes, + ) + self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})") + staging = self._next_staging() + self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes) + self._checkpoint("load: state copy H2D") + GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states) + self._checkpoint("load: load_state_to_cache") + self._track(page_entries + [st]) + + seq.num_cached_tokens = b + seq.cached_hit_tokens = b + self.stats.hits += 1 + self.stats.hit_tokens += b + self.stats.pages_loaded += num_pages + self.stats.states_loaded += 1 + return b + + # ------------------------------------------------------------------ store (write-through at prefill end) + + def on_prefill_step(self, scheduled, finished) -> None: + """Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed.""" + if not self.enabled: + return + try: + for seq in scheduled: + if seq.chunked_rest_len == 0: # last chunk of this prompt + self._store(seq) + except Exception: # noqa: BLE001 + logger.exception("host prefix cache: store failed") + self._fail() + + def _store(self, seq) -> None: + chain = self._chains.get(seq.request_id) + if chain is None: # request did not go through try_resume: hash the prompt from the sequence itself + chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32)) + self._allocate() + bs = self.block_size + b = (seq.get_len() // bs) * bs + if b < self.min_hit_tokens: + return + num_pages = b // bs + self._extend(chain, num_pages) + missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages] + for i in range(num_pages): + if i not in missing: + self.pages.move_to_end(chain.hashes[i]) + # State boundaries: b, b - bs, ... while the window ending there is still in the ring. + seq_len = seq.get_len() + new_states = [] + for k in range(self.state_lookback_blocks): + pos = b - k * bs + if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack: + break + key = chain.hashes[pos // bs - 1] + if key in self.states: + self.states.move_to_end(key) + else: + new_states.append((pos, key)) + if not missing and not new_states: + return + if not self._reserve(self.pages, self.free_pages, len(missing)): + self.stats.store_skipped += 1 + return + # States are best effort: keep the newest boundaries that fit. + while new_states and not self._reserve(self.states, self.free_states, len(new_states)): + new_states.pop() + if not missing and not new_states: + self.stats.store_skipped += 1 + return + + block_table = seq.block_table + entries: list[_Entry] = [] + if missing: + new = [] + for i in missing: + e = _Entry(chain.hashes[i], self.free_pages.pop()) + e.refs = 1 + self.pages[e.key] = e + new.append(e) + self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})") + self.kernel.copy_pages( + self.gpu_blocks_ptr, + self.host_pages_ptr, + self._ids([block_table[i] for i in missing]), + self._ids([e.slot for e in new]), + self.page_bytes, + ) + self._checkpoint("store: page copy D2H") + entries += new + self.stats.pages_stored += len(new) + for pos, key in new_states: + st = _Entry(key, self.free_states.pop()) + st.refs = 1 + self.states[key] = st + staging = self._next_staging() + GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging) + self._checkpoint(f"store: select_state_slice at {pos}") + self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes) + self._checkpoint("store: state copy D2H") + entries.append(st) + self.stats.states_stored += 1 + self._track(entries) + self.stats.stores += 1 + + # ------------------------------------------------------------------ misc + + def _fail(self) -> None: + self.stats.errors += 1 + if self.stats.errors >= 3: + logger.error("host prefix cache: disabled after %d errors", self.stats.errors) + self.enabled = False + + def stats_dict(self) -> dict[str, int | float]: + d = self.stats.as_dict() + d.update( + enabled=int(self.enabled), + pages_used=len(self.pages), + pages_total=self.num_pages, + states_used=len(self.states), + states_total=len(self.free_states) + len(self.states), + inflight=len(self._inflight), + hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0, + ) + return d diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py new file mode 100644 index 0000000..0384de7 --- /dev/null +++ b/deepinfer/hicache/kernels.py @@ -0,0 +1,25 @@ +"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions).""" + +from __future__ import annotations + +from pathlib import Path + +_module = None + + +def copy_pages_module(): + global _module + if _module is None: + from torch.utils.cpp_extension import load + + # Do not touch TORCH_CUDA_ARCH_LIST here: the engine's own JIT kernels (deep_jit / TileLang) run + # in this process and must keep their arch-specific (sm_1xxa) targets. torch's loader derives the + # arch for this trivial kernel from the visible device. + here = Path(__file__).resolve().parent + _module = load( + name="deepinfer_hicache_page_copy", + sources=[str(here / "page_copy.cpp"), str(here / "page_copy.cu")], + extra_cuda_cflags=["-O3"], + verbose=False, + ) + return _module diff --git a/deepinfer/hicache/page_copy.cpp b/deepinfer/hicache/page_copy.cpp new file mode 100644 index 0000000..4a20e78 --- /dev/null +++ b/deepinfer/hicache/page_copy.cpp @@ -0,0 +1,9 @@ +#include + +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", ©_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."); +} diff --git a/deepinfer/hicache/page_copy.cu b/deepinfer/hicache/page_copy.cu new file mode 100644 index 0000000..a445421 --- /dev/null +++ b/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 +#include +#include +#include + +#include +#include + +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(src_base + static_cast(src_ids[blockIdx.x]) * page_bytes); + int4 *dst = reinterpret_cast(dst_base + static_cast(dst_ids[blockIdx.x]) * page_bytes); + const int64_t stride = static_cast(gridDim.y) * blockDim.x; + for (int64_t i = static_cast(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(std::min((n16 + threads - 1) / threads, 16)); + const dim3 grid(static_cast(n), blocks_y); + auto stream = c10::cuda::getCurrentCUDAStream(); + copy_pages_kernel<<>>(reinterpret_cast(src_base), + reinterpret_cast(dst_base), src_ids.data_ptr(), dst_ids.data_ptr(), page_bytes); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tests/hicache/test_host_cache.py b/tests/hicache/test_host_cache.py new file mode 100644 index 0000000..77ffc4e --- /dev/null +++ b/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 From 306e49e8d9d460461cb1e96dd6d9b930ac4b7362 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:26:07 +0000 Subject: [PATCH 06/19] hicache: type annotations for ruff ANN rules Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- deepinfer/hicache/host_cache.py | 15 ++++++++++----- deepinfer/hicache/kernels.py | 3 ++- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/deepinfer/hicache/host_cache.py b/deepinfer/hicache/host_cache.py index e073fc5..40e5ff2 100644 --- a/deepinfer/hicache/host_cache.py +++ b/deepinfer/hicache/host_cache.py @@ -33,6 +33,7 @@ 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 @@ -41,6 +42,10 @@ 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__) @@ -83,7 +88,7 @@ class Stats: class HostPrefixCache: - def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None: + 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" @@ -241,7 +246,7 @@ class HostPrefixCache: # ------------------------------------------------------------------ load (hit) - def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + 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 @@ -256,7 +261,7 @@ class HostPrefixCache: self._fail() return 0 - def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int: + 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) @@ -325,7 +330,7 @@ class HostPrefixCache: # ------------------------------------------------------------------ store (write-through at prefill end) - def on_prefill_step(self, scheduled, finished) -> None: + 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 @@ -337,7 +342,7 @@ class HostPrefixCache: logger.exception("host prefix cache: store failed") self._fail() - def _store(self, seq) -> None: + 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)) diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py index 0384de7..9cc372c 100644 --- a/deepinfer/hicache/kernels.py +++ b/deepinfer/hicache/kernels.py @@ -3,11 +3,12 @@ from __future__ import annotations from pathlib import Path +from types import ModuleType _module = None -def copy_pages_module(): +def copy_pages_module() -> ModuleType: global _module if _module is None: from torch.utils.cpp_extension import load From c6faea3bc2d0b46942ab16df3887cc044612c7e2 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:11:38 +0000 Subject: [PATCH 07/19] Add a minimal OpenAI-compatible chat server over the offline engine serve/openai_server.py exposes /v1/chat/completions (streaming, tools, usage) on top of LLMEngine: one engine process per GPU (attention DP + MoE EP, the layout of examples/benchmark.py), prompts rendered with the checkpoint's own encoding.py, conversation-affinity routing, a watchdog that fails open requests if an engine rank dies, and /deepinfer/metrics. serve/server.sh runs it as a process group so a stop also reaps the engine's helper processes, which otherwise keep the GPUs. It integrates the host prefix cache from deepinfer/hicache when present (--hicache-size-gb, --hicache-state-slots, per-request "prefix_cache": false bypass; serve/check_prefix_cache.py is the A/B check) and serves without it otherwise. It also sets DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL to 2 by default (--prefill-any-interval): with the engine default of 100, a rank whose peers are idle admits a waiting prompt only every 100 decode steps, a ~1.4 s TTFT floor on every request in a serving setting. Used to drive the InferenceX AgentX replay (aiperf) on 4x GB300. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/README.md | 45 ++++ serve/check_prefix_cache.py | 87 ++++++ serve/openai_server.py | 518 ++++++++++++++++++++++++++++++++++++ serve/server.sh | 41 +++ 4 files changed, 691 insertions(+) create mode 100644 serve/README.md create mode 100644 serve/check_prefix_cache.py create mode 100644 serve/openai_server.py create mode 100755 serve/server.sh diff --git a/serve/README.md b/serve/README.md new file mode 100644 index 0000000..9f204eb --- /dev/null +++ b/serve/README.md @@ -0,0 +1,45 @@ +# 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). + +## 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; `prefill-any-interval` matters there, since a single lane keeps most ranks idle. diff --git a/serve/check_prefix_cache.py b/serve/check_prefix_cache.py new file mode 100644 index 0000000..bb8fae1 --- /dev/null +++ b/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) diff --git a/serve/openai_server.py b/serve/openai_server.py new file mode 100644 index 0000000..d34acdc --- /dev/null +++ b/serve/openai_server.py @@ -0,0 +1,518 @@ +"""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 + +# Module level so FastAPI can resolve the string annotation `request: Request`. +from fastapi import Request + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "examples")) + + +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=2, + 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 (engine default 100 => ~1.4 s TTFT floor)", + ) + 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)) + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + + 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, seq, num_cached_tokens): + # 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) -> 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("�"): + 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) + 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) -> None: + q = self.streams.get(rid) + if q is not None: + q.put_nowait(item) + + def _deliver_batch(self, batch) -> 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｜>" + return self.tokenizer.encode(prompt, add_special_tokens=False) + + @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): + 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()] + 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(): + 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(): + 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() diff --git a/serve/server.sh b/serve/server.sh new file mode 100755 index 0000000..b612fd2 --- /dev/null +++ b/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 From 77892b7ae423bba8cd0e3b4e0c1356c195f96368 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:26:07 +0000 Subject: [PATCH 08/19] serve: type annotations and B007 for ruff Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- serve/openai_server.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/serve/openai_server.py b/serve/openai_server.py index d34acdc..1e41bf2 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -26,9 +26,13 @@ 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 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")) @@ -123,7 +127,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m def mm_data(): return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None - def enqueue(rid, seq, num_cached_tokens): + 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) @@ -228,7 +232,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m class IncrementalDetokenizer: """vLLM-style prefix/read offsets so each delta decodes only a short tail.""" - def __init__(self, tokenizer) -> None: + def __init__(self, tokenizer: Any) -> None: self.tok = tokenizer self.ids: list[int] = [] self.num_tokens = 0 @@ -300,12 +304,12 @@ class Frontend: else: self.loop.call_soon_threadsafe(self._deliver_batch, payload) - def _deliver(self, rid: str, item) -> None: + 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) -> None: + 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: @@ -375,7 +379,7 @@ class Frontend: self.stats["finished"] += 1 -def build_app(fe: Frontend): +def build_app(fe: Frontend) -> FastAPI: from fastapi import FastAPI from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse @@ -469,7 +473,7 @@ def build_app(fe: Frontend): done = False try: yield chunk({"role": "assistant", "content": ""}) - async for text, finish in tokens(): + 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 @@ -484,7 +488,7 @@ def build_app(fe: Frontend): done = False try: parts, finish = [], None - async for text, finish in tokens(): + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) parts.append(text) done = True finally: From 57f206a8d53af3510353c65ab0561a42e9fb55dd Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:19:09 +0000 Subject: [PATCH 09/19] engine: expose Scheduler.add_prompted_sequence and free_finished_seqs Two bindings a prefill/decode split needs from the scheduler: - add_prompted_sequence(seq): put a sequence whose KV blocks and state were filled by the caller (received from a prompt-only engine) straight into the prompted queue, without cloning. The sequence holds prompt + first generated token, with everything but the last token computed; the checks assert that layout. - free_finished_seqs(): expose the existing helper that moves finished sequences out of running and returns them without freeing blocks or state, so a prompt-only engine can export the prefilled KV first and free the sequence itself afterwards. No scheduling behaviour changes for engines that do not call them. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- csrc/engine/pybind_scheduler.cpp | 4 ++++ csrc/engine/scheduler.hpp | 13 +++++++++++++ deepinfer/_C/engine.pyi | 4 ++++ 3 files changed, 21 insertions(+) diff --git a/csrc/engine/pybind_scheduler.cpp b/csrc/engine/pybind_scheduler.cpp index a379e85..62fabe1 100644 --- a/csrc/engine/pybind_scheduler.cpp +++ b/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) diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index ed6e897..88e1109 100644 --- a/csrc/engine/scheduler.hpp +++ b/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 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"); diff --git a/deepinfer/_C/engine.pyi b/deepinfer/_C/engine.pyi index f500ad7..08b5a78 100644 --- a/deepinfer/_C/engine.pyi +++ b/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: From 8e488bbd55f3a7270872fb3690ab40051f2b6d47 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:26:29 +0000 Subject: [PATCH 10/19] serve: prefill/decode split prototype over two hosts (pd_server.py) A prompt-only engine tray (PE) and a decode-only engine tray (CE) behind the OpenAI-compatible frontend from serve/openai_server.py: - serve/pd_server.py: --role pe / --role ce. The CE frontend renders the prompt, picks a CE rank (affinity or least-loaded routing) and POSTs /prefill to the PE. The PE prefills with the host prefix cache on, exports the prompt's KV blocks and the full SWA state in the pe_do_save_kvcache hook, streams the first sampled token back in the HTTP response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence into decode via Scheduler.add_prompted_sequence and reclaims PE sequences through Scheduler.free_finished_seqs. - deepinfer/hicache/pd_transfer.py: export/admit of KV pages and state through the hicache page-copy kernels. The CE keeps the pinned receive buffers alive until the admit copy kernels finish (the kernels read them by raw pointer; the caching host allocator otherwise reused the memory while the copy was queued -> cudaErrorLaunchFailure). - serve/pd_run.sh start/stop/logs over ssh; serve/check_pd.py compares outputs against a single-engine server token by token; serve/nccl_xtray_test.py reports the cross-host NCCL transport. Measured on 4 GB300 (PE) + 4 GB300 (CE) against one 4-GPU tray with the AgentX replay, outputs identical: ITL p90 stays within 8.6-15 ms from conc 8 to 256 (single tray 12-98 ms), P90 interactivity 66-117 tok/s/user where the single tray falls to 10; per total chip PD is behind up to conc 64 and ahead at 256 (doc/benchmark_gb300.md). Depends on the host prefix cache (#144), the OpenAI server (#147) and the scheduler bindings (add_prompted_sequence / free_finished_seqs), which are included as the first three commits. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- deepinfer/hicache/pd_transfer.py | 125 ++++++ serve/README.md | 18 + serve/check_pd.py | 81 ++++ serve/nccl_xtray_test.py | 30 ++ serve/pd_run.sh | 49 +++ serve/pd_server.py | 682 +++++++++++++++++++++++++++++++ 6 files changed, 985 insertions(+) create mode 100644 deepinfer/hicache/pd_transfer.py create mode 100644 serve/check_pd.py create mode 100644 serve/nccl_xtray_test.py create mode 100755 serve/pd_run.sh create mode 100644 serve/pd_server.py diff --git a/deepinfer/hicache/pd_transfer.py b/deepinfer/hicache/pd_transfer.py new file mode 100644 index 0000000..abc69e1 --- /dev/null +++ b/deepinfer/hicache/pd_transfer.py @@ -0,0 +1,125 @@ +"""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 + num_blocks = (prompt_len + self.block_size - 1) // self.block_size + block_table = list(seq.block_table)[:num_blocks] + pages = torch.empty((num_blocks, self.page_bytes), dtype=torch.uint8, pin_memory=True) + 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 + ) + state = torch.empty(self.state_bytes, dtype=torch.uint8, pin_memory=True) + 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 = torch.cuda.Event() + 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 pages, state, 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 buffers and put it in the scheduler's prompted queue. Caller must have called try_allocate_seq. + + 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, + ) + 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 diff --git a/serve/README.md b/serve/README.md index 9f204eb..a362759 100644 --- a/serve/README.md +++ b/serve/README.md @@ -43,3 +43,21 @@ stored when they come back as part of the next turn's prompt). The server speaks what InferenceX's AgentX replay (aiperf, `--endpoint /v1/chat/completions --streaming --use-server-token-count`) needs; `prefill-any-interval` matters there, since a single lane keeps most ranks idle. + +## 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 HTTP 8889, CE HTTP 8888, CE KV receivers 8900+rank +python serve/check_pd.py --ref http://:8888 --pd http://:8888 # token-identical? +serve/pd_run.sh stop +``` + +`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`. diff --git a/serve/check_pd.py b/serve/check_pd.py new file mode 100644 index 0000000..2203325 --- /dev/null +++ b/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) diff --git a/serve/nccl_xtray_test.py b/serve/nccl_xtray_test.py new file mode 100644 index 0000000..543bf87 --- /dev/null +++ b/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= 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() diff --git a/serve/pd_run.sh b/serve/pd_run.sh new file mode 100755 index 0000000..f61a451 --- /dev/null +++ b/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 [extra pd_server args...] +# serve/pd_run.sh stop +# serve/pd_run.sh log # tail the tray's log +# +# Ports: PE HTTP 8889, CE HTTP (user-facing) 8888, CE KV receivers 8900+rank. Logs: ~/pd_.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 diff --git a/serve/pd_server.py b/serve/pd_server.py new file mode 100644 index 0000000..d94672e --- /dev/null +++ b/serve/pd_server.py @@ -0,0 +1,682 @@ +"""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://: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 + + +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) + 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)) + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + + 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) + + 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 + pages, state, ev = transfer.export(seq, prompt_len) + item = { + "pages": pages, + "state": state, + "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 sender(): + 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() + 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) + + 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_one(conn: socket.socket): + try: + with conn: + hdr = bytearray(HDR.size) + _recv_exact(conn, memoryview(hdr)) + 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" + 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)) + 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, torch.cuda.Event] = {} + 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) + arrived.pop(rid, None) + if rid in pending: + pending.remove(rid) + else: + engine.abort_request(rid) + while True: + try: + rid, first_token, prompt_len, pages, state = recv_q.get_nowait() + except queue.Empty: + break + arrived[rid] = (first_token, prompt_len, pages, state) + if rid in expected: + pending.append(rid) + while pending: + rid = pending[0] + spec, (first_token, prompt_len, pages, state) = 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) + 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")) + 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) + 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, (ev, _, _) in keep.items() if ev.query()]: + del keep[rid] + + 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() From d7844a221d7c7e8ec0694795a68dbfcad3bc9068 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:17:22 +0000 Subject: [PATCH 11/19] serve: EPLB rearrangement off by default, --eplb-interval / --eplb-after-prefill flags Hard-wired enable_eplb=True rearranged experts every 30 steps and after every prefill step; under multi-turn traffic that fires almost continuously and each rearrangement stalls all ranks for ~0.3 s. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/README.md | 4 ++++ serve/openai_server.py | 21 ++++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index 9f204eb..1f30b79 100644 --- a/serve/README.md +++ b/serve/README.md @@ -30,6 +30,10 @@ Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-l 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. + ## 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` diff --git a/serve/openai_server.py b/serve/openai_server.py index 1e41bf2..e117e1b 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -67,6 +67,22 @@ def parse_args() -> argparse.Namespace: 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 (engine default 100 => ~1.4 s TTFT floor)", ) + 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", + ) return p.parse_args() @@ -92,7 +108,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m # 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)) - setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + 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, From 4159328eb56efc3eef11676d475034ce25a976d2 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:17:22 +0000 Subject: [PATCH 12/19] serve: EPLB rearrangement off by default, --eplb-interval / --eplb-after-prefill flags Hard-wired enable_eplb=True rearranged experts every 30 steps and after every prefill step; under multi-turn traffic that fires almost continuously and each rearrangement stalls all ranks for ~0.3 s. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/README.md | 4 ++++ serve/openai_server.py | 21 ++++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index a362759..edcf9fa 100644 --- a/serve/README.md +++ b/serve/README.md @@ -30,6 +30,10 @@ Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-l 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. + ## 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` diff --git a/serve/openai_server.py b/serve/openai_server.py index 1e41bf2..e117e1b 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -67,6 +67,22 @@ def parse_args() -> argparse.Namespace: 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 (engine default 100 => ~1.4 s TTFT floor)", ) + 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", + ) return p.parse_args() @@ -92,7 +108,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m # 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)) - setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + 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, From 94d032c2f0dfd7dd95099b5f253ae86d757574eb Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:17:44 +0000 Subject: [PATCH 13/19] serve/pd_server: EPLB off by default, same --eplb-interval / --eplb-after-prefill flags Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/pd_server.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/serve/pd_server.py b/serve/pd_server.py index d94672e..94bd4cb 100644 --- a/serve/pd_server.py +++ b/serve/pd_server.py @@ -71,6 +71,8 @@ def parse_args() -> argparse.Namespace: 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") return p.parse_args() @@ -103,7 +105,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m 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)) - setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + 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( From 269a61f7c54a2d56145bee46c566808541cd3467 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:49:52 +0000 Subject: [PATCH 14/19] serve: exact fast prompt tokenization (fastokens or added-token segment cache), checked against tokenizer.encode Agentic prompts are re-tokenized in full every turn; tokenizer.encode takes ~280 ms p50 on AgentX prompts (p50 163K tokens) and dominates TTFT once the prefix cache hits. serve/prompt_tokenizer.py picks fastokens if installed, else a dependency-free segment cache that splits at added tokens (where HF already encodes pieces independently) and reuses earlier turns' pieces; both are enabled only after matching tokenizer.encode on startup probes, and a sampled background re-check switches to hf on any mismatch. Suggested by whn09 in review. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 18 ++ serve/openai_server.py | 33 +++- serve/prompt_tokenizer.py | 264 +++++++++++++++++++++++++++ tests/serve/__init__.py | 0 tests/serve/test_prompt_tokenizer.py | 136 ++++++++++++++ 5 files changed, 450 insertions(+), 1 deletion(-) create mode 100644 serve/prompt_tokenizer.py create mode 100644 tests/serve/__init__.py create mode 100644 tests/serve/test_prompt_tokenizer.py diff --git a/serve/README.md b/serve/README.md index 1f30b79..20b2f1f 100644 --- a/serve/README.md +++ b/serve/README.md @@ -34,6 +34,24 @@ EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearran 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. +## 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` diff --git a/serve/openai_server.py b/serve/openai_server.py index e117e1b..772390d 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -36,6 +36,8 @@ if TYPE_CHECKING: 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: @@ -83,6 +85,19 @@ def parse_args() -> argparse.Namespace: 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() @@ -282,6 +297,20 @@ class Frontend: 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)] @@ -357,7 +386,7 @@ class Frontend: if body.get("tools"): parts.insert(1, json.dumps(body["tools"], ensure_ascii=False)) prompt = "".join(parts) + "<｜Assistant｜>" - return self.tokenizer.encode(prompt, add_special_tokens=False) + return self.prompt_tokenizer.encode(prompt) @staticmethod def affinity_key(body: dict) -> bytes | None: @@ -426,6 +455,8 @@ def build_app(fe: Frontend) -> FastAPI: 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") diff --git a/serve/prompt_tokenizer.py b/serve/prompt_tokenizer.py new file mode 100644 index 0000000..feda0f0 --- /dev/null +++ b/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 diff --git a/tests/serve/__init__.py b/tests/serve/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/serve/test_prompt_tokenizer.py b/tests/serve/test_prompt_tokenizer.py new file mode 100644 index 0000000..cf18997 --- /dev/null +++ b/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" From 83be1fa44c34231911cfc223b9212d6364254462 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Sun, 25 Oct 2026 21:35:59 +0000 Subject: [PATCH 15/19] serve: drop the 32 redundant expert slots when EPLB is off Without rearrangement they are fixed copies of experts 0-31 and only add expert weight per rank. Decode SPS +4.8% / +3.9% at 4 / 16 requests per rank on GB300 (examples/benchmark.py, 128K context), matching whn09's B200 numbers; 6 GiB less model memory per GPU. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 4 +++- serve/openai_server.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index 20b2f1f..3a7a33a 100644 --- a/serve/README.md +++ b/serve/README.md @@ -32,7 +32,9 @@ per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, 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. +`--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%. ## Prompt tokenization diff --git a/serve/openai_server.py b/serve/openai_server.py index 772390d..ca6c557 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -123,6 +123,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m # 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) From 9e455f55d9bb480fa774067206b5d1b8083d5c54 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Sun, 25 Oct 2026 21:36:15 +0000 Subject: [PATCH 16/19] serve/pd_server: drop the redundant expert slots when EPLB is off, as openai_server.py does Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/pd_server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/serve/pd_server.py b/serve/pd_server.py index 3c7813f..342c223 100644 --- a/serve/pd_server.py +++ b/serve/pd_server.py @@ -107,6 +107,9 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m 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) From 8b1ed3f26aefbd956178f5a3577931b34874d8fe Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Sun, 25 Oct 2026 21:37:45 +0000 Subject: [PATCH 17/19] serve: document the --prefill-any-interval TTFT / decode-speed trade-off Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 5 +++++ serve/openai_server.py | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index 3a7a33a..44ce939 100644 --- a/serve/README.md +++ b/serve/README.md @@ -36,6 +36,11 @@ EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearran 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 2) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and +with 2 almost every arriving request triggers one. That keeps TTFT low; the engine default of 100 adds a ~1.4 s +floor. At high concurrency, 30-100 gives faster decode and more throughput at the cost of TTFT: on B200 AgentX c256, +interval 100 gave 255K vs 211K tok/s/GPU and interactivity p50 119 vs 55, with TTFT p50 2.5 vs 0.7 s. + ## Prompt tokenization Agentic clients resend the whole conversation every turn, and `tokenizer.encode` needs ~1.7 ms per 1K tokens on a diff --git a/serve/openai_server.py b/serve/openai_server.py index ca6c557..b11f94e 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -67,7 +67,8 @@ def parse_args() -> argparse.Namespace: type=int, default=2, 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 (engine default 100 => ~1.4 s TTFT floor)", + "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor). A prompt step " + "pauses decode on every rank, so at high concurrency a larger value (30-100) trades TTFT for decode speed", ) p.add_argument( "--eplb-interval", From f6388940003e9c962337bcd1f8d0ba693d977f86 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Mon, 26 Oct 2026 04:43:27 +0000 Subject: [PATCH 18/19] serve: default --prefill-any-interval to 30 (throughput over TTFT) At AgentX c256 on one GB300 tray, 30 vs 2 gives +9% tok/s/GPU and 2.2x P90 interactivity, for +0.4 s TTFT p50. Use 2 for low TTFT. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 19 ++++++++++++++----- serve/openai_server.py | 6 +++--- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/serve/README.md b/serve/README.md index 44ce939..e7b119f 100644 --- a/serve/README.md +++ b/serve/README.md @@ -36,10 +36,18 @@ EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearran 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 2) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and -with 2 almost every arriving request triggers one. That keeps TTFT low; the engine default of 100 adds a ~1.4 s -floor. At high concurrency, 30-100 gives faster decode and more throughput at the cost of TTFT: on B200 AgentX c256, -interval 100 gave 255K vs 211K tok/s/GPU and interactivity p50 119 vs 55, with TTFT p50 2.5 vs 0.7 s. +`--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 @@ -71,4 +79,5 @@ 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; `prefill-any-interval` matters there, since a single lane keeps most ranks idle. +--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. diff --git a/serve/openai_server.py b/serve/openai_server.py index b11f94e..4bfd939 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -65,10 +65,10 @@ def parse_args() -> argparse.Namespace: p.add_argument( "--prefill-any-interval", type=int, - default=2, + 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 (engine default 100 => ~1.4 s TTFT floor). A prompt step " - "pauses decode on every rank, so at high concurrency a larger value (30-100) trades TTFT for decode speed", + "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", From d54fe25993a44a7f09f5bbe6420fec6438c1ea79 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 10:31:23 +0800 Subject: [PATCH 19/19] serve: Mooncake Transfer Engine KV transport for the prefill/decode split (--kv-transport mooncake) Stacked on #197 (pd_server.py / pd_transfer.py), which moves the prefilled KV blocks and SWA state PE -> CE as pinned host buffers over TCP. --kv-transport mooncake keeps that flow and control socket, and replaces the data plane with a Mooncake TransferEngine: - deepinfer/hicache/pd_mooncake.py: GpuArena (one registered device buffer per rank, first-fit regions; cuMemAlloc like the KV pools, since expandable_segments memory has no cudaIpc handle) and MooncakeLink (TransferEngine with P2PHANDSHAKE metadata, register_memory, transfer_sync_write). - PDTransfer.export_into / admit accept device buffers: the PE exports a sequence straight into an arena region (state included, no extra staging) and the CE admits from its arena. - pd_server.py: MC_MAGIC header -> CE reserves an arena region and replies {addr, segment} (or busy -> retry) -> one contiguous write -> go/ok. Writes are retried, then the region is sent over the TCP path, so a failing peer cannot strand a request. TCP is unchanged and is also the fallback when the PE arena is full. - --mooncake-protocol nvlink_intra (default, one host) or rdma/efa across hosts. - serve/README.md: usage; containers must share IPC and PID namespaces (--pid=host). Co-Authored-By: Claude Opus 5.5 --- deepinfer/hicache/pd_mooncake.py | 106 +++++++++++++++++++++ deepinfer/hicache/pd_transfer.py | 32 +++++-- serve/README.md | 33 +++++++ serve/pd_server.py | 158 ++++++++++++++++++++++++++++--- 4 files changed, 309 insertions(+), 20 deletions(-) create mode 100644 deepinfer/hicache/pd_mooncake.py diff --git a/deepinfer/hicache/pd_mooncake.py b/deepinfer/hicache/pd_mooncake.py new file mode 100644 index 0000000..c69376e --- /dev/null +++ b/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}") diff --git a/deepinfer/hicache/pd_transfer.py b/deepinfer/hicache/pd_transfer.py index abc69e1..09f5a4e 100644 --- a/deepinfer/hicache/pd_transfer.py +++ b/deepinfer/hicache/pd_transfer.py @@ -74,27 +74,40 @@ class PDTransfer: """ if prompt_len is None: prompt_len = seq.get_len() - 1 - num_blocks = (prompt_len + self.block_size - 1) // self.block_size + 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] - pages = torch.empty((num_blocks, self.page_bytes), dtype=torch.uint8, pin_memory=True) 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 ) - state = torch.empty(self.state_bytes, dtype=torch.uint8, pin_memory=True) + 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 = torch.cuda.Event() 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 pages, state, ev + 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 buffers and put it in the scheduler's prompted queue. Caller must have called try_allocate_seq. + 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. """ @@ -108,8 +121,11 @@ class PDTransfer: self._ids(block_table[:num_blocks]), self.page_bytes, ) - 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) + 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 diff --git a/serve/README.md b/serve/README.md index ba9add5..9643e32 100644 --- a/serve/README.md +++ b/serve/README.md @@ -99,3 +99,36 @@ serve/pd_run.sh stop `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. diff --git a/serve/pd_server.py b/serve/pd_server.py index 342c223..e5900df 100644 --- a/serve/pd_server.py +++ b/serve/pd_server.py @@ -47,6 +47,29 @@ 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: @@ -75,6 +98,18 @@ def parse_args() -> argparse.Namespace: 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() @@ -169,6 +204,7 @@ def _pe_loop( 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() @@ -188,10 +224,24 @@ def _pe_loop( out_q.put(("error", rid, f"prefilled sequence has {len(seq.block_table)} blocks, needs {need}")) targets.pop(rid, None) continue - pages, state, ev = transfer.export(seq, prompt_len) + 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), @@ -203,7 +253,43 @@ def _pe_loop( 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"] @@ -212,6 +298,16 @@ def _pe_loop( 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) @@ -323,7 +419,8 @@ def _ce_loop( ) -> None: import torch - recv_q: queue.Queue = queue.Queue() # (rid, first_token, prompt_len, pages, state) + 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) @@ -334,19 +431,47 @@ def _ce_loop( 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)) + _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)) + 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) @@ -357,7 +482,12 @@ def _ce_loop( 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, torch.cuda.Event] = {} + 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: @@ -371,22 +501,24 @@ def _ce_loop( pending.append(rid) elif kind == "abort": expected.pop(rid, None) - arrived.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 = recv_q.get_nowait() + rid, first_token, prompt_len, pages, state, region = recv_q.get_nowait() except queue.Empty: break - arrived[rid] = (first_token, prompt_len, pages, state) + 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) = expected[rid], arrived[rid] + 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"]) @@ -399,6 +531,7 @@ def _ce_loop( pending.popleft() expected.pop(rid, None) arrived.pop(rid, None) + release(region) continue try: allocated = engine.scheduler.try_allocate_seq(seq) @@ -413,15 +546,16 @@ def _ce_loop( 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) + 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, (ev, _, _) in keep.items() if ev.query()]: - del keep[rid] + 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()