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/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/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 91aa482..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: @@ -455,6 +459,9 @@ class Sequence: def cached_hit_tokens(self, arg1: typing.SupportsInt | typing.SupportsIndex) -> None: ... @property + def block_table(self) -> list[int]: + ... + @property def chunk_checkpoints(self) -> list[int]: ... @chunk_checkpoints.setter 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..40e5ff2 --- /dev/null +++ b/deepinfer/hicache/host_cache.py @@ -0,0 +1,437 @@ +"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache. + +Layout + Two pinned host pools, both organized in pages: + - KV pages: one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block. + - State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a + 128-token-aligned position with GPUCacheInterface.select_state_slice_to_store. + Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)), + so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool. + +Transfer + "kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through + UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never + blocks the CPU. Completion is tracked with CUDA events polled from the engine loop. + +Policy + Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last + 128-aligned position are stored (pages already present are only touched). On a new request, the longest + page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to + be recomputed, then snapped down to the nearest position that has a stored state; the sequence is + pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with + num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is + rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store). + +Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's +prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank). +""" + +from __future__ import annotations + +import hashlib +import os +import time +from collections import OrderedDict, deque +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from deepinfer import envs +from deepinfer.logger import init_logger +from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface + +if TYPE_CHECKING: + from deepinfer._C.engine import Sequence + from deepinfer.engine.llm_engine import LLMEngine + +logger = init_logger(__name__) + + +class _Entry: + __slots__ = ("key", "slot", "refs", "ready") + + def __init__(self, key: bytes, slot: int) -> None: + self.key = key + self.slot = slot + self.refs = 0 # in-flight transfers (store or load) touching this slot + self.ready = False # host copy complete and hit-able + + +@dataclass +class _Chain: + tokens: np.ndarray + hashes: list[bytes] = field(default_factory=list) + counted: bool = False # lookup stats recorded (a deferred request is looked up repeatedly) + + +@dataclass +class Stats: + lookups: int = 0 + hits: int = 0 + hit_tokens: int = 0 + lookup_tokens: int = 0 + deferred: int = 0 # hit found but GPU allocation not possible yet + stores: int = 0 + store_skipped: int = 0 + pages_stored: int = 0 + pages_loaded: int = 0 + states_stored: int = 0 + states_loaded: int = 0 + evicted_pages: int = 0 + evicted_states: int = 0 + errors: int = 0 + + def as_dict(self) -> dict[str, int]: + return dict(self.__dict__) + + +class HostPrefixCache: + def __init__(self, engine: LLMEngine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None: + assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)" + assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1" + assert GPUCacheInterface.cache_info_ready, "cache engine not initialized" + + self.engine = engine + self.scheduler = engine.scheduler + cache_engine = engine.driver_worker.cache_engine + self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks + self.gpu_states: torch.Tensor = cache_engine.gpu_states + assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8 + self.device = self.gpu_blocks.device + self.block_size: int = engine.cache_config.block_size + assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN) + self.page_bytes: int = self.gpu_blocks.size(1) + self.state_bytes: int = GPUCacheInterface.state_size_to_disk + assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes) + self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS + # The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative. + self.min_hit_tokens = 2 * self.block_size + # A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only + # at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a + # block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA + # ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore. + self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4")) + self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference + + self.num_pages = max(0, size_bytes // self.page_bytes) + self.num_state_slots = num_state_slots + self.staging_slots = staging_slots + assert self.num_pages > 0 and num_state_slots > 0 + self.gpu_blocks_ptr = self.gpu_blocks.data_ptr() + # Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be + # cached before a prefill completes, and an unused cache then costs no pinned memory). + self._allocated = False + self.host_pages = self.host_states = None + self.host_pages_ptr = self.host_states_ptr = 0 + self._staging: list[torch.Tensor] = [] + self._staging_idx = 0 + self.kernel = None + + self.pages: OrderedDict[bytes, _Entry] = OrderedDict() + self.states: OrderedDict[bytes, _Entry] = OrderedDict() + self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1)) + self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1)) + self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque() + self._chains: dict[str, _Chain] = {} + self.stats = Stats() + self.enabled = True + # DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel). + self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1" + logger.info( + "HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d", + self.num_pages, + self.page_bytes, + self.num_pages * self.page_bytes / 2**30, + num_state_slots, + self.state_bytes, + num_state_slots * self.state_bytes / 2**30, + self.keep_tokens, + ) + + def _allocate(self) -> None: + if self._allocated: + return + from deepinfer.hicache.kernels import copy_pages_module + + t0 = time.perf_counter() + self.kernel = copy_pages_module() + self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True) + self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True) + self.host_pages_ptr = self.host_pages.data_ptr() + self.host_states_ptr = self.host_states.data_ptr() + # Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream. + self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)] + self._allocated = True + logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0) + + # ------------------------------------------------------------------ hashing + + def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None: + chain = self._chains.get(request_id) + if chain is None and tokens is not None: + chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32)) + return chain + + def _extend(self, chain: _Chain, num_pages: int) -> None: + hashes, toks, bs = chain.hashes, chain.tokens, self.block_size + prev = hashes[-1] if hashes else b"" + for i in range(len(hashes), num_pages): + prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest() + hashes.append(prev) + + def forget(self, request_id: str) -> None: + self._chains.pop(request_id, None) + + # ------------------------------------------------------------------ transfers + + def _ids(self, ids: list[int]) -> torch.Tensor: + return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True) + + def _next_staging(self) -> torch.Tensor: + t = self._staging[self._staging_idx] + self._staging_idx = (self._staging_idx + 1) % len(self._staging) + return t + + def _track(self, entries: list[_Entry]) -> None: + ev = torch.cuda.Event() + ev.record() + self._inflight.append((ev, entries)) + + def _checkpoint(self, what: str) -> None: + if self._debug_sync: + from deepinfer.utils import allow_cuda_sync + + with allow_cuda_sync(): + torch.cuda.synchronize() + logger.info("hicache debug: ok after %s", what) + + def poll(self) -> None: + """Retire completed transfers (events on one stream complete in order, so only the head is checked).""" + while self._inflight and self._inflight[0][0].query(): + _, entries = self._inflight.popleft() + for e in entries: + e.refs -= 1 + e.ready = True + + def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int: + """Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found.""" + victims, busy = [], [] + for key, e in pool.items(): + if len(victims) >= need: + break + if e.refs or not e.ready: + busy.append(key) + continue + victims.append(key) + for key in busy: + pool.move_to_end(key) + if len(victims) < need: + return 0 + for key in victims: + free.append(pool.pop(key).slot) + return len(victims) + + def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool: + if need > len(free) + len(pool): + return False + if len(free) < need: + evicted = self._evict(pool, free, need - len(free)) + if pool is self.pages: + self.stats.evicted_pages += evicted + else: + self.stats.evicted_states += evicted + return len(free) >= need + + # ------------------------------------------------------------------ load (hit) + + def try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int: + """Fill `seq` (not yet in the scheduler) from the host cache. + + Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not + allocatable right now, sequence untouched, retry later). + """ + if not self.enabled: + return 0 + try: + return self._try_resume(seq, request_id, prompt_ids) + except Exception: # noqa: BLE001 - never take the engine down for a cache problem + logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id) + self._fail() + return 0 + + def _try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int: + bs = self.block_size + n = len(prompt_ids) + chain = self._chain(request_id, prompt_ids) + if not chain.counted: + chain.counted = True + self.stats.lookups += 1 + self.stats.lookup_tokens += n + limit = ((n - self.keep_tokens) // bs) * bs # the replay rectangle [n - keep, n) must be recomputed + if limit < self.min_hit_tokens: + return 0 + max_pages = limit // bs + self._extend(chain, max_pages) + matched = 0 + for h in chain.hashes[:max_pages]: + e = self.pages.get(h) + if e is None or not e.ready: + break + matched += 1 + b = min(matched * bs, limit) + while b >= self.min_hit_tokens: + st = self.states.get(chain.hashes[b // bs - 1]) + if st is not None and st.ready: + break + b -= bs + else: + return 0 + + try: + allocated = self.scheduler.try_allocate_seq(seq) + except RuntimeError: # can never fit (prompt longer than the whole KV pool): let the scheduler reject it + return 0 + if not allocated: + self.stats.deferred += 1 + return -1 + num_pages = b // bs + page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)] + for e in page_entries: + e.refs += 1 + self.pages.move_to_end(e.key) + st.refs += 1 + self.states.move_to_end(st.key) + + block_table = seq.block_table + self.kernel.copy_pages( + self.host_pages_ptr, + self.gpu_blocks_ptr, + self._ids([e.slot for e in page_entries]), + self._ids(list(block_table[:num_pages])), + self.page_bytes, + ) + self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})") + staging = self._next_staging() + self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes) + self._checkpoint("load: state copy H2D") + GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states) + self._checkpoint("load: load_state_to_cache") + self._track(page_entries + [st]) + + seq.num_cached_tokens = b + seq.cached_hit_tokens = b + self.stats.hits += 1 + self.stats.hit_tokens += b + self.stats.pages_loaded += num_pages + self.stats.states_loaded += 1 + return b + + # ------------------------------------------------------------------ store (write-through at prefill end) + + def on_prefill_step(self, scheduled: list[Sequence], finished: list[Sequence]) -> None: + """Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed.""" + if not self.enabled: + return + try: + for seq in scheduled: + if seq.chunked_rest_len == 0: # last chunk of this prompt + self._store(seq) + except Exception: # noqa: BLE001 + logger.exception("host prefix cache: store failed") + self._fail() + + def _store(self, seq: Sequence) -> None: + chain = self._chains.get(seq.request_id) + if chain is None: # request did not go through try_resume: hash the prompt from the sequence itself + chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32)) + self._allocate() + bs = self.block_size + b = (seq.get_len() // bs) * bs + if b < self.min_hit_tokens: + return + num_pages = b // bs + self._extend(chain, num_pages) + missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages] + for i in range(num_pages): + if i not in missing: + self.pages.move_to_end(chain.hashes[i]) + # State boundaries: b, b - bs, ... while the window ending there is still in the ring. + seq_len = seq.get_len() + new_states = [] + for k in range(self.state_lookback_blocks): + pos = b - k * bs + if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack: + break + key = chain.hashes[pos // bs - 1] + if key in self.states: + self.states.move_to_end(key) + else: + new_states.append((pos, key)) + if not missing and not new_states: + return + if not self._reserve(self.pages, self.free_pages, len(missing)): + self.stats.store_skipped += 1 + return + # States are best effort: keep the newest boundaries that fit. + while new_states and not self._reserve(self.states, self.free_states, len(new_states)): + new_states.pop() + if not missing and not new_states: + self.stats.store_skipped += 1 + return + + block_table = seq.block_table + entries: list[_Entry] = [] + if missing: + new = [] + for i in missing: + e = _Entry(chain.hashes[i], self.free_pages.pop()) + e.refs = 1 + self.pages[e.key] = e + new.append(e) + self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})") + self.kernel.copy_pages( + self.gpu_blocks_ptr, + self.host_pages_ptr, + self._ids([block_table[i] for i in missing]), + self._ids([e.slot for e in new]), + self.page_bytes, + ) + self._checkpoint("store: page copy D2H") + entries += new + self.stats.pages_stored += len(new) + for pos, key in new_states: + st = _Entry(key, self.free_states.pop()) + st.refs = 1 + self.states[key] = st + staging = self._next_staging() + GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging) + self._checkpoint(f"store: select_state_slice at {pos}") + self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes) + self._checkpoint("store: state copy D2H") + entries.append(st) + self.stats.states_stored += 1 + self._track(entries) + self.stats.stores += 1 + + # ------------------------------------------------------------------ misc + + def _fail(self) -> None: + self.stats.errors += 1 + if self.stats.errors >= 3: + logger.error("host prefix cache: disabled after %d errors", self.stats.errors) + self.enabled = False + + def stats_dict(self) -> dict[str, int | float]: + d = self.stats.as_dict() + d.update( + enabled=int(self.enabled), + pages_used=len(self.pages), + pages_total=self.num_pages, + states_used=len(self.states), + states_total=len(self.free_states) + len(self.states), + inflight=len(self._inflight), + hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0, + ) + return d diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py new file mode 100644 index 0000000..9cc372c --- /dev/null +++ b/deepinfer/hicache/kernels.py @@ -0,0 +1,26 @@ +"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions).""" + +from __future__ import annotations + +from pathlib import Path +from types import ModuleType + +_module = None + + +def copy_pages_module() -> ModuleType: + global _module + if _module is None: + from torch.utils.cpp_extension import load + + # Do not touch TORCH_CUDA_ARCH_LIST here: the engine's own JIT kernels (deep_jit / TileLang) run + # in this process and must keep their arch-specific (sm_1xxa) targets. torch's loader derives the + # arch for this trivial kernel from the visible device. + here = Path(__file__).resolve().parent + _module = load( + name="deepinfer_hicache_page_copy", + sources=[str(here / "page_copy.cpp"), str(here / "page_copy.cu")], + extra_cuda_cflags=["-O3"], + verbose=False, + ) + return _module 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/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 new file mode 100644 index 0000000..09f5a4e --- /dev/null +++ b/deepinfer/hicache/pd_transfer.py @@ -0,0 +1,141 @@ +"""Prefill/decode split: move a prefilled sequence's KV blocks and SWA state between engines. + +The prefill engine (EngineType.PROMPT) exports, after the last prefill chunk, the sequence's full KV blocks +(all tokens of the prompt) and the complete SWA/compression state in the "to CE" layout +(GPUCacheInterface.select_state_slice_to_store(..., TO_CE=True)), both into pinned host buffers through the +page-copy kernel. The decode engine (EngineType.COMPLETION) receives those bytes into pinned buffers, allocates +blocks and a state slot for a new sequence (prompt + first sampled token), fills them with the kernel, and admits +the sequence into its `prompted` queue (Scheduler.add_prompted_sequence). Transport between the two is left to +the caller (plain TCP in serve/pd_server.py). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from deepinfer import envs +from deepinfer.hicache.kernels import copy_pages_module +from deepinfer.worker.cache_engine import GPUCacheInterface + +if TYPE_CHECKING: + from deepinfer._C.engine import Sequence + from deepinfer.engine.llm_engine import LLMEngine + + +def load_state_from_ce_transfer(state_id: int, last_position: int, state_tensor: torch.Tensor, state_cache: torch.Tensor) -> None: + """Inverse of select_state_slice_to_store(..., TO_CE=True): write a full inference-layout state (stored SWA, + dropped/consumer SWA and compression state) into state_cache[state_id] with the ring ending at last_position.""" + from deepinfer.model_executor.layers.tilelang_ops.state_cache import load_state_to_cache + + cls = GPUCacheInterface + cache_views, tensor_views = cls._carve_state_views(state_id, state_cache, state_tensor, TO_CE=True) + for tensor_view, cache_view in zip(tensor_views, cache_views, strict=True): + if tensor_view.numel(): + load_state_to_cache(tensor_view, cache_view, last_position) + dropped_tensor_offset = cls.swa_total_bytes_inference_pad + dropped_cache_offset = cls.swa_total_bytes_kvstore_pad + compress_bytes = cls.dropped_bytes_inference - cls.swa_drop_total_bytes_inference_pad + if compress_bytes: + state_cache[state_id, dropped_cache_offset : dropped_cache_offset + compress_bytes].copy_( + state_tensor[dropped_tensor_offset : dropped_tensor_offset + compress_bytes] + ) + + +class PDTransfer: + """Per-engine helper owning the GPU pools' pointers and small pinned staging for transfers.""" + + def __init__(self, engine: LLMEngine) -> None: + cache_engine = engine.driver_worker.cache_engine + self.engine = engine + self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks + self.gpu_states: torch.Tensor = cache_engine.gpu_states + self.device = self.gpu_blocks.device + self.block_size: int = engine.cache_config.block_size + self.page_bytes: int = self.gpu_blocks.size(1) + self.state_bytes: int = GPUCacheInterface.total_bytes_inference # full "to CE" state + assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0 + self.kernel = copy_pages_module() + + def _ids(self, ids: list[int]) -> torch.Tensor: + return torch.tensor(ids, dtype=torch.int32).pin_memory().to(self.device, non_blocking=True) + + # ---------------------------------------------------------------- prefill side + + def export(self, seq: Sequence, prompt_len: int | None = None) -> tuple[torch.Tensor, torch.Tensor, torch.cuda.Event]: + """Copy seq's KV blocks for its prompt and the full state to pinned host buffers. + + prompt_len: tokens with KV (default: seq.get_len() - 1, i.e. after the first output token was appended; + pass seq.get_len() when called before outputs are processed). KV exists for the prompt only, so a prompt + ending exactly on a block boundary must not count an extra block. + Returns (pages[num_blocks, page_bytes] pinned, state[state_bytes] pinned, event); the buffers are valid once + the event has completed. Enqueued on the current stream, so it is ordered after the prefill's kernels. + """ + if prompt_len is None: + prompt_len = seq.get_len() - 1 + pages = torch.empty((self.num_blocks(prompt_len), self.page_bytes), dtype=torch.uint8, pin_memory=True) + state = torch.empty(self.state_bytes, dtype=torch.uint8, pin_memory=True) + return pages, state, self.export_into(seq, prompt_len, pages, state) + + def num_blocks(self, prompt_len: int) -> int: + return (prompt_len + self.block_size - 1) // self.block_size + + def export_into(self, seq: Sequence, prompt_len: int, pages: torch.Tensor, state: torch.Tensor) -> torch.cuda.Event: + """export() into caller-provided buffers: pages [num_blocks(prompt_len), page_bytes] and state [state_bytes], + pinned host or device memory (e.g. a region of a Mooncake-registered arena, pd_mooncake.GpuArena).""" + num_blocks = pages.size(0) + block_table = list(seq.block_table)[:num_blocks] + self.kernel.copy_pages( + self.gpu_blocks.data_ptr(), pages.data_ptr(), self._ids(block_table), self._ids(list(range(num_blocks))), self.page_bytes + ) + ev = torch.cuda.Event() + if state.is_cuda: + GPUCacheInterface.select_state_slice_to_store(seq.state_id, prompt_len, self.gpu_states, state_tensor=state, TO_CE=True) + ev.record() + return ev + staging = torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) + GPUCacheInterface.select_state_slice_to_store(seq.state_id, prompt_len, self.gpu_states, state_tensor=staging, TO_CE=True) + self.kernel.copy_pages(staging.data_ptr(), state.data_ptr(), self._ids([0]), self._ids([0]), self.state_bytes) + ev.record() + # keep the staging alive until the copy ran (same-stream ordering makes later reuse safe) + ev.staging = staging # type: ignore[attr-defined] + return ev + + # ---------------------------------------------------------------- decode side + + def admit(self, seq: Sequence, pages: torch.Tensor, state: torch.Tensor, prompt_len: int) -> torch.cuda.Event: + """Fill a freshly created sequence (prompt + first token, created with add_request(enqueue=False)) from + pinned host or device buffers and put it in the scheduler's prompted queue. Caller must have called + try_allocate_seq, and must keep `pages`/`state` alive until the returned event has completed. + + prompt_len: tokens with KV in `pages` (= seq.get_len() - 1); the state ring ends there. + """ + num_blocks = pages.size(0) + block_table = list(seq.block_table) + assert len(block_table) >= num_blocks, (len(block_table), num_blocks) + self.kernel.copy_pages( + pages.data_ptr(), + self.gpu_blocks.data_ptr(), + self._ids(list(range(num_blocks))), + self._ids(block_table[:num_blocks]), + self.page_bytes, + ) + if state.is_cuda: + staging = state + else: + staging = torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) + self.kernel.copy_pages(state.data_ptr(), staging.data_ptr(), self._ids([0]), self._ids([0]), self.state_bytes) + load_state_from_ce_transfer(seq.state_id, prompt_len, staging, self.gpu_states) + seq.num_cached_tokens = prompt_len + seq.cached_hit_tokens = prompt_len + # A finished prefill leaves `nextn` draft tokens on the sequence (pad tokens when the prompt step produced + # none, as a prompt-only engine does); the decode sampler requires exactly that many. + nextn = int(envs.DEEPINFER_USE_NEXTN_DECODE) + if nextn > 0: + seq.last_next_n_token = self.engine.get_nextn_pad_token_ids().tolist()[:nextn] + self.engine.scheduler.add_prompted_sequence(seq) + ev = torch.cuda.Event() + ev.record() + ev.staging = staging # type: ignore[attr-defined] + return ev diff --git a/serve/README.md b/serve/README.md new file mode 100644 index 0000000..9643e32 --- /dev/null +++ b/serve/README.md @@ -0,0 +1,134 @@ +# Serving DeepInfer + host prefix cache (local additions, not part of upstream) + +## Pieces + +- `serve/openai_server.py` — OpenAI-compatible `/v1/chat/completions` (streaming, `tools`, usage) on top of the + offline engine: one engine process per GPU (attention DP + MoE EP), prompts rendered with the checkpoint's own + `encoding.py`. `serve/server.sh start|stop` runs it as a process group (stop reaps the engine helpers, which + otherwise keep the GPUs). +- `deepinfer/hicache/` — host-memory prefix cache (SGLang HiCache style): pinned host pools organized in pages + (one page = one 128-token KV block) plus SWA state slots, hash-chain index, LRU eviction, write-through at the + end of every prefill, "kernel" transfers (a CUDA kernel copying pages through UVA pointers on the engine + stream, `page_copy.cu`, JIT-built on first use). Needs the read-only `Sequence.block_table` binding added in + `csrc/engine/pybind_sequence.cpp` (rebuild `deepinfer._C` with `bash scripts/build.sh build`). +- `tests/hicache/test_host_cache.py` — cache logic against a fake engine (needs one GPU). +- `serve/check_prefix_cache.py` — A/B on a live server: each turn sent with `"prefix_cache": false` and with + the cache, greedy; compares outputs and TTFT. + +## Run + +```bash +export DEEPINFER_ACCELERATOR=nvidia DEEPINFER_HF_MODEL_DIR=/path/to/DeepSeek-V4.1-Flash-1017 +serve/server.sh start --dp 4 --port 8888 # cache on: 32 GiB pages + 4096 states per rank +serve/server.sh start --dp 4 --port 8888 --hicache-size-gb 0 # cache off +python serve/check_prefix_cache.py --port 8888 --turns 3 +curl -s localhost:8888/deepinfer/metrics | grep hicache # per-rank hit/store counters +serve/server.sh stop +``` + +Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-loaded` (affinity keeps a +conversation's turns on one rank, keyed by the first user message; required for hits since the cache is +per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, the default 140). + +EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearrangement stalls all ranks for about +0.3 s, and with `--eplb-after-prefill 1` it fires after nearly every step under multi-turn traffic; enable it with +`--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. With EPLB off the server also drops +the 32 redundant expert slots (`DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=0`): without rearrangement they are fixed copies +of experts 0-31, and removing them speeds decode up by 4-5%. + +`--prefill-any-interval` (default 30) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and +a rank with a waiting prompt forces one at most this many decode steps later. The default favours throughput. With 2, +almost every arriving request triggers a prompt step, which keeps TTFT low but slows decode at high concurrency. The +engine default of 100 adds a ~1.4 s TTFT floor. AgentX at c256 on one GB300 tray (20-minute runs): + +| interval | tok/s/GPU | TTFT p50 / p90 | ITL p50 / p90 | P90 interactivity | +|---|---|---|---|---| +| 2 | 258K | 1.26 / 4.53 s | 14.3 / 31.5 ms | 32 | +| 30 | 283K | 1.62 / 6.51 s | 10.0 / 14.1 ms | 71 | +| 100 | 276K | 3.00 / 14.95 s | 8.1 / 13.5 ms | 74 | + +B200 shows the same trend: at c256, interval 100 gave 255K vs 211K tok/s/GPU for interval 2. + +## Prompt tokenization + +Agentic clients resend the whole conversation every turn, and `tokenizer.encode` needs ~1.7 ms per 1K tokens on a +GB300 Grace core (p50 278 ms per request on replayed AgentX prompts of p50 163K tokens), which dominates TTFT once the +prefix cache hits. `serve/prompt_tokenizer.py` gives the same ids faster; `--tokenizer-backend` picks one: + +- `fastokens` (`pip install fastokens`; x86-64 and aarch64 wheels): Rust reimplementation of the HF pipeline, built + from the loaded tokenizer. 3.5 ms p50 on the same prompts. +- `segment` (no dependency): splits the prompt at added tokens, where HF tokenizers already encode the pieces + independently, and caches each piece's ids (LRU, 512 MiB), so a turn only encodes what it adds. 18 ms p50. + Enabled only when the tokenizer config makes the split exact (no `single_word` tokens, no `prepend_scheme="first"`). +- `hf`: `tokenizer.encode`. + +`auto` (default) takes the first of these that matches `tokenizer.encode` on startup probes built from the +tokenizer's added tokens. `--tokenizer-verify-rate` (default 0.005) re-encodes a sample of prompts with +`tokenizer.encode` in the background; a mismatch switches to `hf`. The backend, encode latency and check counters are +in `/deepinfer/metrics` (`deepinfer_tokenizer:*`). Tests: `tests/serve/test_prompt_tokenizer.py`. + +## How a hit works + +A new prompt's page hashes are matched from the start; the match is capped so that at least `keep_tokens` +tokens remain to be recomputed (the decoder SWA state is rebuilt from that suffix by bounded replay), then +snapped down to the nearest 128-aligned position that has a stored SWA state. The sequence is allocated up +front, pages and state are copied in on the engine stream, `num_cached_tokens` is set, and the sequence is +enqueued; the scheduler prefills only the suffix. Generated tokens are not stored at decode end (they are +stored when they come back as part of the next turn's prompt). + +## AgentX + +The server speaks what InferenceX's AgentX replay (aiperf, `--endpoint /v1/chat/completions --streaming +--use-server-token-count`) needs. At low concurrency, use `--prefill-any-interval 2`: a single lane keeps most +ranks idle, and the default adds a few hundred ms of TTFT. + +## Prefill/decode split (`pd_server.py`) + +`serve/pd_server.py` runs the same OpenAI-compatible frontend over two engines on two hosts: a prompt-only +engine (PE, `--role pe`, host prefix cache on) and a decode-only engine (CE, `--role ce`, user-facing). +Per request the CE frontend renders the prompt, picks a CE rank and POSTs `/prefill` to the PE; the PE +prefills, exports the sequence's KV blocks and SWA state (`deepinfer/hicache/pd_transfer.py`) plus the first +sampled token, pushes the bytes to the CE rank over TCP and the CE admits the sequence straight into decode +(`Scheduler.add_prompted_sequence`). + +```bash +serve/pd_run.sh start # PE 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`. + +### 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/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/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/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/openai_server.py b/serve/openai_server.py new file mode 100644 index 0000000..4bfd939 --- /dev/null +++ b/serve/openai_server.py @@ -0,0 +1,577 @@ +"""Minimal OpenAI-compatible chat server on top of the DeepInfer offline engine. + +DeepInfer ships offline inference only. This wraps one `LLM` per data-parallel rank +(attention DP + MoE EP, same layout as examples/benchmark.py) behind a FastAPI frontend +so request-replay benchmarks such as InferenceX AgentX (aiperf) can drive it. + + python serve/openai_server.py --model ~/models/DeepSeek-V4.1-Flash-1017 --dp 4 --port 8888 + +Limitations, by design of the open-source engine: +- No cross-request prefix cache: every request prefills its full prompt. +- Responses are streamed as plain `content`; DSML tool calls are not parsed into `tool_calls`. +""" + +from __future__ import annotations + +import argparse +import asyncio +import collections +import hashlib +import json +import multiprocessing as mp +import os +import queue +import sys +import threading +import time +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, Any + +# Module level so FastAPI can resolve the string annotation `request: Request`. +from fastapi import FastAPI, Request + +if TYPE_CHECKING: + from deepinfer._C.engine import Sequence + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "examples")) +sys.path.insert(0, str(REPO_ROOT / "serve")) +from prompt_tokenizer import PromptTokenizer # noqa: E402 + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument("--model", required=True, help="local HF checkpoint directory") + p.add_argument("--served-model-name", default=None) + p.add_argument("--dp", type=int, default=4, help="data-parallel ranks on this node (= EP size)") + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=8888) + p.add_argument("--max-model-len", type=int, default=1 << 20) + p.add_argument("--max-num-seqs", type=int, default=128, help="per DP rank") + p.add_argument("--max-num-batched-tokens", type=int, default=32768, help="prefill chunk per DP rank") + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--yoco-keep-tokens", type=int, default=140, help="0 disables decoder SWA bounded replay") + p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"]) + p.add_argument("--log-interval-steps", type=int, default=200) + p.add_argument("--hicache-size-gb", type=float, default=32, help="pinned host KV page pool per DP rank; 0 disables the prefix cache") + p.add_argument("--hicache-state-slots", type=int, default=4096, help="host SWA state slots per DP rank") + p.add_argument( + "--routing", + default="affinity", + choices=["affinity", "least-loaded"], + help="affinity: a conversation's turns go to one rank (needed for cache hits)", + ) + p.add_argument( + "--prefill-any-interval", + type=int, + default=30, + help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this " + "many decode steps later even if other ranks are idle. A prompt step pauses decode on every rank, so a larger " + "value gives faster decode and more throughput at the cost of TTFT. 30 favours throughput; use 2 for low TTFT", + ) + p.add_argument( + "--eplb-interval", + type=int, + default=0, + help="DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL: rearrange experts every N steps; 0 (default) disables EPLB. " + "Each rearrangement stalls every rank (~0.3 s on B200/GB300), which under online traffic costs far more than " + "the balance it buys", + ) + p.add_argument( + "--eplb-after-prefill", + type=int, + default=0, + choices=[0, 1], + help="DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE: also rearrange after every prefill step (only with " + "--eplb-interval > 0); with agentic traffic nearly every step batch contains a prefill", + ) + p.add_argument( + "--tokenizer-backend", + default="auto", + choices=["auto", "fastokens", "segment", "hf"], + help="prompt tokenization (serve/prompt_tokenizer.py); auto = the first of fastokens (if installed), segment, hf " + "that matches tokenizer.encode on startup probes. All produce tokenizer.encode's ids", + ) + p.add_argument( + "--tokenizer-verify-rate", + type=float, + default=0.005, + help="fraction of prompts re-encoded with tokenizer.encode in the background; a mismatch switches to hf", + ) + return p.parse_args() + + +# --------------------------------------------------------------------------- engine side + + +def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None: + """One DP rank: owns an engine and steps it in lockstep with the other ranks.""" + import ctypes + import signal + + # Die with the frontend instead of lingering as an orphan that holds the GPU. + ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL) # PR_SET_PDEATHSIG + + from _config import create_llm, setup_common_env + + if args.yoco_keep_tokens > 0: + os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact" + os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens) + os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1" + if args.hicache_size_gb > 0: + # Keep more SWA history in the state ring so the prefix cache can store states at several + # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks). + os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") + os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval)) + if args.eplb_interval == 0: + # The 32 redundant expert slots only pay off with EPLB rearrangement; without it they are fixed copies of + # experts 0-31 and only add expert weight to read per rank (decode +4-5% without them on B200 and GB300). + os.environ.setdefault("DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS", "0") + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False) + if args.eplb_interval > 0: + os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval) + os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill) + + llm = create_llm( + args.model, + dp_rank, + ep_size=args.dp, + env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)}, + data_parallel_size=args.dp, + max_num_seqs=args.max_num_seqs, + max_num_batched_tokens=args.max_num_batched_tokens, + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))), + ) + + import numpy as np + + from deepinfer._C.engine import RequestOutput # pyright: ignore[reportMissingModuleSource] + from deepinfer.sampling_params import SamplingParams + + engine = llm.llm_engine + cache = None + if args.hicache_size_gb > 0: + try: + from deepinfer.hicache import HostPrefixCache + except ImportError: # tree without the host prefix cache: serve without it + print(f"[server] deepinfer.hicache not available, prefix cache disabled (rank {dp_rank})", flush=True) + else: + cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots) + engine.pe_do_save_kvcache = cache.on_prefill_step # write-through at the end of every prefill + out_q.put(("ready", dp_rank, None)) + + def mm_data(): + return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None + + def enqueue(rid: str, seq: Sequence, num_cached_tokens: int) -> None: + # Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict + # (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length. + engine.scheduler.add_sequence(seq) + if engine.multi_modal_processor is not None: + engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached_tokens) + + pending = collections.deque() # (rid, seq, prompt_ids, use_cache): created, waiting for a cache decision + inflight = False + last_stats = time.monotonic() + last_logged_lookups = 0 + while True: + # Requests must be enqueued before the DP-wide has_unfinished_requests all-reduce, + # so every rank takes the same step / flush / idle branch in the same iteration. + while True: + try: + kind, rid, payload = in_q.get_nowait() + except queue.Empty: + break + if kind == "add": + try: + prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32) + seq = engine.add_request( + user_id="0", + request_id=rid, + sampling_params=SamplingParams(**payload["sampling"]), + prompt_token_ids=prompt_ids, + multi_modal_data=mm_data(), + enqueue=cache is None, + ) + if cache is not None: + if seq is None: + out_q.put(("error", rid, "request dropped: invalid prompt tokens")) + else: + pending.append((rid, seq, prompt_ids, payload.get("use_cache", True))) + except Exception as e: # noqa: BLE001 - surface to the client, keep the rank alive + out_q.put(("error", rid, repr(e))) + elif kind == "abort": + if any(p[0] == rid for p in pending): + pending = collections.deque(p for p in pending if p[0] != rid) + else: + engine.abort_request(rid) + if cache is not None: + cache.forget(rid) + + # Cache decisions, in arrival order. A hit pre-allocates the sequence's blocks and fills them from the + # host; if the GPU cannot allocate yet, wait (while something is running that will free memory). + while pending: + rid, seq, prompt_ids, use_cache = pending[0] + hit = cache.try_resume(seq, rid, prompt_ids) if use_cache else 0 + if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight): + break + pending.popleft() + enqueue(rid, seq, max(hit, 0)) + if cache is not None: + cache.poll() + if time.monotonic() - last_stats > 5: + last_stats = time.monotonic() + st = cache.stats_dict() + out_q.put(("stats", dp_rank, st)) + if st["lookups"] != last_logged_lookups: + last_logged_lookups = st["lookups"] + print( + f"[hicache rank {dp_rank}] lookups {st['lookups']} hits {st['hits']} " + f"token hit rate {st['hit_rate_tokens']:.1%} deferred {st['deferred']} stores {st['stores']} " + f"pages {st['pages_used']}/{st['pages_total']} states {st['states_used']}/{st['states_total']} " + f"evicted {st['evicted_pages']}", + flush=True, + ) + + if engine.has_unfinished_requests(): + outputs = engine.step() + inflight = True + elif inflight: + outputs = engine.flush_inflight_reqs() + inflight = False + else: + time.sleep(0.001) + continue + + if len(outputs): + # Pointers are recycled on the next step, so read the fields now. + batch = [ + ( + RequestOutput.get_request_id(p), + RequestOutput.get_token_id(p), + RequestOutput.get_finished(p), + RequestOutput.get_finish_reason(p), + RequestOutput.get_prompt_len(p), + ) + for p in outputs + ] + if cache is not None: + for rid, _, finished, _, _ in batch: + if finished: + cache.forget(rid) + out_q.put(("tokens", None, batch)) + + +# --------------------------------------------------------------------------- HTTP side + + +class IncrementalDetokenizer: + """vLLM-style prefix/read offsets so each delta decodes only a short tail.""" + + def __init__(self, tokenizer: Any) -> None: + self.tok = tokenizer + self.ids: list[int] = [] + self.num_tokens = 0 + self.prefix_offset = 0 + self.read_offset = 0 + + def add(self, token_id: int) -> str: + self.num_tokens += 1 + if token_id == self.tok.eos_token_id: + return "" + self.ids.append(token_id) + prefix = self.tok.decode(self.ids[self.prefix_offset : self.read_offset], skip_special_tokens=False) + full = self.tok.decode(self.ids[self.prefix_offset :], skip_special_tokens=False) + if len(full) > len(prefix) and not full.endswith("�"): + self.prefix_offset, self.read_offset = self.read_offset, len(self.ids) + return full[len(prefix) :] + return "" + + +class Frontend: + def __init__(self, args: argparse.Namespace) -> None: + from transformers import AutoTokenizer + + sys.path.insert(0, str(Path(args.model) / "encoding")) + from encoding import encode_messages # the checkpoint's own prompt-format reference + + self.args = args + self.model_name = args.served_model_name or args.model + self.encode_messages = encode_messages + self.tokenizer = AutoTokenizer.from_pretrained(args.model) + probe = [ + {"role": "system", "content": "sys", "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}]}, + {"role": "user", "content": "hi 你好 123"}, + {"role": "assistant", "content": "ok", "tool_calls": [{"type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "content": " out\n"}, + ] + try: + probes = [encode_messages(probe, thinking_mode=args.thinking_mode)] + except Exception: # noqa: BLE001 - the probe only adds coverage + probes = [] + # getattr: frontends built on this one with their own argument parser (serve/pd_server.py) get the defaults + backend, verify_rate = getattr(args, "tokenizer_backend", "auto"), getattr(args, "tokenizer_verify_rate", 0.005) + self.prompt_tokenizer = PromptTokenizer(self.tokenizer, backend, verify_rate, extra_probes=probes) + print(f"[server] {self.prompt_tokenizer.describe()}", flush=True) + ctx = mp.get_context("spawn") + self.out_q: mp.Queue = ctx.Queue() + self.in_qs: list[mp.Queue] = [ctx.Queue() for _ in range(args.dp)] + # Not daemonic: the engine forks its own helper processes (as examples/_launch.py allows). + self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)] + self.load = [0] * args.dp + self.owner: dict[str, int] = {} + self.streams: dict[str, asyncio.Queue] = {} + self.ready = 0 + self.stats = dict(requests=0, finished=0, aborted=0, errors=0, prompt_tokens=0, completion_tokens=0, encode_fallbacks=0) + self.cache_stats: dict[int, dict] = {} + self.loop: asyncio.AbstractEventLoop | None = None + + def start(self) -> None: + for p in self.procs: + p.start() + + def dispatch_forever(self) -> None: + """Thread: fan engine outputs out to the per-request asyncio queues.""" + while True: + try: + kind, rid, payload = self.out_q.get(timeout=1.0) + except queue.Empty: + dead = [r for r, p in enumerate(self.procs) if not p.is_alive()] + if dead: + print(f"[server] engine rank(s) {dead} died; failing open requests and exiting", flush=True) + for rid in list(self.streams): + self.loop.call_soon_threadsafe(self._deliver, rid, ("error", f"engine rank died: {dead}")) + time.sleep(2) + os._exit(1) + continue + if kind == "ready": + self.ready += 1 + print(f"[server] DP rank {rid} ready ({self.ready}/{self.args.dp})", flush=True) + elif kind == "error": + self.loop.call_soon_threadsafe(self._deliver, rid, ("error", payload)) + elif kind == "stats": + self.cache_stats[rid] = payload + else: + self.loop.call_soon_threadsafe(self._deliver_batch, payload) + + def _deliver(self, rid: str, item: tuple) -> None: + q = self.streams.get(rid) + if q is not None: + q.put_nowait(item) + + def _deliver_batch(self, batch: list[tuple]) -> None: + for rid, token_id, finished, reason, prompt_len in batch: + q = self.streams.get(rid) + if q is not None: + q.put_nowait(("token", token_id, finished, reason, prompt_len)) + + def render_prompt(self, body: dict) -> list[int]: + messages = [dict(m) for m in body["messages"]] + if body.get("tools"): + if not messages or messages[0].get("role") != "system": + messages.insert(0, {"role": "system", "content": ""}) + messages[0]["tools"] = body["tools"] + try: + prompt = self.encode_messages(messages, thinking_mode=self.args.thinking_mode) + except Exception: # noqa: BLE001 - trace messages the reference encoder rejects + # Keep the request (and its prompt length) instead of failing the benchmark. + self.stats["encode_fallbacks"] += 1 + parts = ["<｜begin▁of▁sentence｜>"] + for m in messages: + content = m.get("content") or "" + if isinstance(content, list): + content = "".join(b.get("text", "") for b in content if isinstance(b, dict)) + if m.get("tool_calls"): + content += json.dumps(m["tool_calls"], ensure_ascii=False) + tag = {"user": "<｜User｜>", "assistant": "<｜Assistant｜>", "tool": "<｜User｜>"}.get(m.get("role"), "") + parts.append(tag + content) + if body.get("tools"): + parts.insert(1, json.dumps(body["tools"], ensure_ascii=False)) + prompt = "".join(parts) + "<｜Assistant｜>" + return self.prompt_tokenizer.encode(prompt) + + @staticmethod + def affinity_key(body: dict) -> bytes | None: + """Stable per-conversation key: the first user message (AgentX's cache-bust marker lives there and is + constant across a trajectory's turns, and forked subagents inherit it).""" + for m in body.get("messages", []): + if m.get("role") == "user": + content = m.get("content") or "" + if isinstance(content, list): + content = "".join(b.get("text", "") for b in content if isinstance(b, dict)) + return hashlib.sha1(content[:4096].encode("utf-8", "ignore")).digest() + return None + + def submit(self, prompt_ids: list[int], sampling: dict, key: bytes | None = None, use_cache: bool = True) -> tuple[str, asyncio.Queue]: + rid = uuid.uuid4().hex + if self.args.routing == "affinity" and key is not None: + rank = int.from_bytes(key[:4], "little") % self.args.dp + else: + rank = min(range(self.args.dp), key=lambda r: self.load[r]) + self.load[rank] += 1 + self.owner[rid] = rank + q: asyncio.Queue = asyncio.Queue() + self.streams[rid] = q + self.stats["requests"] += 1 + self.in_qs[rank].put(("add", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling, "use_cache": use_cache})) + return rid, q + + def release(self, rid: str, aborted: bool) -> None: + rank = self.owner.pop(rid, None) + self.streams.pop(rid, None) + if rank is None: + return + self.load[rank] -= 1 + if aborted: + self.stats["aborted"] += 1 + self.in_qs[rank].put(("abort", rid, None)) + else: + self.stats["finished"] += 1 + + +def build_app(fe: Frontend) -> FastAPI: + from fastapi import FastAPI + from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse + + app = FastAPI() + + @app.on_event("startup") + async def _startup() -> None: + fe.loop = asyncio.get_running_loop() + threading.Thread(target=fe.dispatch_forever, daemon=True).start() + + @app.get("/health") + async def health(): + ok = fe.ready == fe.args.dp + return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready}, status_code=200 if ok else 503) + + @app.get("/v1/models") + async def models(): + return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]} + + # Not at /metrics: aiperf scrapes that path automatically and InferenceX's aggregator + # rejects metric families it has no backend adapter for. + @app.get("/deepinfer/metrics") + async def metrics(): + lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()] + lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)] + for r, st in sorted(fe.cache_stats.items()): + lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()] + pt = fe.prompt_tokenizer + lines += [f'deepinfer_tokenizer:{k}{{backend="{pt.name}"}} {v}' for k, v in pt.metrics().items()] + return PlainTextResponse("\n".join(lines) + "\n") + + @app.post("/v1/chat/completions") + async def chat(request: Request): + if fe.ready != fe.args.dp: + return JSONResponse({"error": {"message": "engine loading"}}, status_code=503) + body = await request.json() + prompt_ids = await asyncio.to_thread(fe.render_prompt, body) + max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096 + if len(prompt_ids) + 1 > fe.args.max_model_len: + return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400) + max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids))) + sampling = dict( + temperature=float(body.get("temperature", 1.0)), + top_p=float(body.get("top_p", 1.0)), + ignore_eos=bool(body.get("ignore_eos", False)), + max_tokens=max_tokens, + ) + if body.get("seed") is not None: + sampling["seq_seed"] = int(body["seed"]) + fe.stats["prompt_tokens"] += len(prompt_ids) + include_usage = bool((body.get("stream_options") or {}).get("include_usage")) + # Non-standard request field: "prefix_cache": false bypasses the host prefix cache lookup (A/B checks). + rid, q = fe.submit(prompt_ids, sampling, fe.affinity_key(body), use_cache=bool(body.get("prefix_cache", True))) + created = int(time.time()) + cid = f"chatcmpl-{rid}" + detok = IncrementalDetokenizer(fe.tokenizer) + + def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str: + obj = { + "id": cid, + "object": "chat.completion.chunk", + "created": created, + "model": fe.model_name, + "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}], + } + if usage: + obj["usage"] = usage + return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n" + + async def tokens(): + """Yield (text_delta, finish_reason) until the request finishes.""" + while True: + item = await q.get() + if item[0] == "error": + raise RuntimeError(item[1]) + _, token_id, finished, reason, _ = item + text = detok.add(token_id) + fe.stats["completion_tokens"] += 1 + if finished: + yield text, ("length" if detok.num_tokens >= max_tokens else "stop") + return + yield text, None + + usage = lambda: { + "prompt_tokens": len(prompt_ids), + "completion_tokens": detok.num_tokens, # noqa: E731 + "total_tokens": len(prompt_ids) + detok.num_tokens, + } + + if body.get("stream"): + + async def sse(): + done = False + try: + yield chunk({"role": "assistant", "content": ""}) + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) + if text or finish: + yield chunk({"content": text} if text else {}, finish) + done = True + if include_usage: + yield chunk({}, usage=usage()) + yield "data: [DONE]\n\n" + finally: + fe.release(rid, aborted=not done) + + return StreamingResponse(sse(), media_type="text/event-stream") + + done = False + try: + parts, finish = [], None + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) + parts.append(text) + done = True + finally: + fe.release(rid, aborted=not done) + return { + "id": cid, + "object": "chat.completion", + "created": created, + "model": fe.model_name, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}], + "usage": usage(), + } + + return app + + +def main() -> None: + args = parse_args() + args.model = str(Path(args.model).expanduser().resolve()) + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("RANK", "0") + fe = Frontend(args) + fe.start() + + import uvicorn + + uvicorn.run(build_app(fe), host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900) + + +if __name__ == "__main__": + main() 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..e5900df --- /dev/null +++ b/serve/pd_server.py @@ -0,0 +1,826 @@ +"""Prefill/decode-disaggregated serving: a prompt-only engine tray (PE) and a decode-only engine tray (CE). + + # decode tray (user-facing OpenAI API on --port, KV receivers on --kv-port + rank): + python serve/pd_server.py --role ce --model ... --dp 4 --port 8888 --kv-port 8900 --pe-url http://:8889 + # prefill tray: + python serve/pd_server.py --role pe --model ... --dp 4 --port 8889 + +Flow per request: CE frontend renders the prompt, picks a CE rank, POSTs /prefill to the PE with the KV target +(ce-host, kv-port + rank). The PE prefills (host prefix cache on), exports the sequence's KV blocks + full SWA +state (deepinfer.hicache.pd_transfer) and the first sampled token, streams the first token back in the HTTP +response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence straight into decode +(Scheduler.add_prompted_sequence); its tokens stream to the client through the CE frontend. +""" + +from __future__ import annotations + +import argparse +import asyncio +import collections +import json +import multiprocessing as mp +import os +import queue +import socket +import struct +import sys +import threading +import time +import uuid +from collections.abc import Callable +from pathlib import Path +from types import ModuleType +from typing import TYPE_CHECKING + +from fastapi import FastAPI, Request + +if TYPE_CHECKING: + from deepinfer._C.engine import RequestOutput, Sequence + from deepinfer.engine.llm_engine import LLMEngine + from deepinfer.hicache.pd_transfer import PDTransfer + from deepinfer.sampling_params import SamplingParams + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parent / "examples")) + +import openai_server as oas # noqa: E402 (the single-engine server: prompt rendering, detokenizer, frontend plumbing) + +HDR = struct.Struct("<16sqqqqq") # rid(16 bytes hex) | first_token | prompt_len | num_pages | page_bytes | state_bytes +# --kv-transport mooncake: MC_MAGIC + HDR; CE replies one JSON line {"status": ok|busy|error, "addr", "segment"}; +# after the TransferEngine write the PE sends b"go" (or b"xx" on failure) and the CE answers b"ok". +MC_MAGIC = b"MCK1" # cannot start a TCP-path header (rids are hex) +MC_RETRY_S = 0.005 +MC_WRITE_ATTEMPTS = 3 # then fall back to TCP for that transfer, so a failing peer never strands a request + + +def _mooncake_link(args: argparse.Namespace, device): # noqa: ANN001, ANN202 + from deepinfer.hicache.pd_mooncake import GpuArena, MooncakeLink + + arena = GpuArena(int(args.pd_arena_gb * 2**30), device) + host = args.mooncake_host or (args.kv_host if args.role == "ce" and args.kv_host else socket.gethostbyname(socket.gethostname())) + return MooncakeLink(host, args.mooncake_protocol, arena, args.mooncake_device) + + +def _recv_line(sock: socket.socket) -> bytes: + out = bytearray() + while not out.endswith(b"\n"): + chunk = sock.recv(1) + if not chunk: + raise ConnectionError("connection closed") + out += chunk + return bytes(out) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument("--role", required=True, choices=["pe", "ce"]) + p.add_argument("--model", required=True) + p.add_argument("--served-model-name", default=None) + p.add_argument("--dp", type=int, default=4) + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=8888) + p.add_argument("--kv-port", type=int, default=8900, help="ce: TCP receiver base port (+ rank)") + p.add_argument("--kv-host", default=None, help="ce: address the PE should push KV to (default: this host's name)") + p.add_argument("--pe-url", default=None, help="ce: prefill tray URL") + p.add_argument("--max-model-len", type=int, default=1 << 20) + p.add_argument("--max-num-seqs", type=int, default=128) + p.add_argument("--max-num-batched-tokens", type=int, default=32768) + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--yoco-keep-tokens", type=int, default=140) + p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"]) + p.add_argument("--log-interval-steps", type=int, default=200) + p.add_argument("--hicache-size-gb", type=float, default=32, help="pe: host prefix cache per rank (0 disables)") + p.add_argument("--hicache-state-slots", type=int, default=4096) + p.add_argument("--routing", default="affinity", choices=["affinity", "least-loaded"]) + p.add_argument("--prefill-any-interval", type=int, default=2) + p.add_argument("--eplb-interval", type=int, default=0, help="see openai_server.py; 0 (default) disables EPLB") + p.add_argument("--eplb-after-prefill", type=int, default=0, choices=[0, 1], help="see openai_server.py") + p.add_argument("--tokenizer-backend", default="auto", choices=["auto", "fastokens", "segment", "hf"], help="see openai_server.py") + p.add_argument("--tokenizer-verify-rate", type=float, default=0.005, help="see openai_server.py") + p.add_argument( + "--kv-transport", + default="tcp", + choices=["tcp", "mooncake"], + help="KV/state data plane PE -> CE. tcp: pinned host buffers over the KV socket. mooncake: GPU arena to GPU " + "arena through a Mooncake TransferEngine (nvlink_intra within a node, rdma/efa across hosts); the KV socket " + "carries only control messages", + ) + p.add_argument("--mooncake-protocol", default="nvlink_intra", help="TransferEngine protocol (nvlink_intra, rdma, efa, tcp)") + p.add_argument("--mooncake-device", default="", help="TransferEngine device name(s) for rdma/efa (empty: auto)") + p.add_argument("--mooncake-host", default=None, help="address peers use for this rank's TransferEngine handshake") + p.add_argument("--pd-arena-gb", type=float, default=4, help="mooncake: GPU staging arena per rank (PE send / CE receive)") + return p.parse_args() + + +def _recv_exact(sock: socket.socket, buf: memoryview) -> None: + view = buf + while len(view): + n = sock.recv_into(view, min(len(view), 1 << 24)) + if n == 0: + raise ConnectionError("peer closed during transfer") + view = view[n:] + + +# ============================================================================ engine processes + + +def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None: + import ctypes + import signal + + ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL) # PR_SET_PDEATHSIG + + import numpy as np + from _config import create_llm, setup_common_env + + from deepinfer.engine.engine_type import EngineType + + if args.yoco_keep_tokens > 0: + os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact" + os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens) + os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1" + os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") # identical state layout on PE and CE + os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval)) + if args.eplb_interval == 0: + # Same as openai_server.py: without EPLB the redundant expert slots only add expert weight per rank. + os.environ.setdefault("DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS", "0") + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False) + if args.eplb_interval > 0: + os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval) + os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill) + + engine_type = EngineType.PROMPT if args.role == "pe" else EngineType.COMPLETION + llm = create_llm( + args.model, + dp_rank, + ep_size=args.dp, + env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)}, + data_parallel_size=args.dp, + max_num_seqs=args.max_num_seqs, + max_num_batched_tokens=args.max_num_batched_tokens, + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))), + engine_type=engine_type, + stop_preempted_request=True, + ) + + from deepinfer._C.engine import RequestOutput # pyright: ignore[reportMissingModuleSource] + from deepinfer.hicache.pd_transfer import PDTransfer + from deepinfer.sampling_params import SamplingParams + + engine = llm.llm_engine + transfer = PDTransfer(engine) + mm = lambda: dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None # noqa: E731 + + def enqueue(rid: str, seq: Sequence, num_cached: int) -> None: + engine.scheduler.add_sequence(seq) + if engine.multi_modal_processor is not None: + engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached) + + if args.role == "pe": + _pe_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm, enqueue) + else: + _ce_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm) + + +def _pe_loop( + args: argparse.Namespace, + dp_rank: int, + in_q: mp.Queue, + out_q: mp.Queue, + engine: LLMEngine, + transfer: PDTransfer, + SamplingParams: type[SamplingParams], + RequestOutput: type[RequestOutput], + np: ModuleType, + mm: Callable[[], dict | None], + enqueue: Callable[[str, Sequence, int], None], +) -> None: + cache = None + if args.hicache_size_gb > 0: + from deepinfer.hicache import HostPrefixCache + + cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots) + + link = _mooncake_link(args, transfer.device) if args.kv_transport == "mooncake" else None + targets: dict[str, tuple[str, int]] = {} # rid -> (host, port) of the CE rank + awaiting_token: dict[str, dict] = {} # rid -> transfer waiting for the step's sampled first token + send_q: queue.Queue = queue.Queue() + + def on_prefill_step(scheduled: list[Sequence], finished: list[Sequence]) -> None: + # Engine hook at the start of post-processing for a prompt step: the prefill's KV is complete and the + # sequence still owns its blocks (with max_tokens=1 it finishes, and is freed, once outputs are processed). + if cache is not None: + cache.on_prefill_step(scheduled, finished) + for seq in scheduled: + rid = seq.request_id + if seq.chunked_rest_len != 0 or rid not in targets: + continue + prompt_len = seq.get_len() # no output appended yet + need = (prompt_len + transfer.block_size - 1) // transfer.block_size + if len(seq.block_table) < need: + out_q.put(("error", rid, f"prefilled sequence has {len(seq.block_table)} blocks, needs {need}")) + targets.pop(rid, None) + continue + region = None + if link is not None: + from deepinfer.hicache.pd_mooncake import region_layout + + nb = transfer.num_blocks(prompt_len) + state_off, total = region_layout(nb, transfer.page_bytes, transfer.state_bytes) + off = link.arena.alloc(total) + if off is not None: # else: arena full, this one goes over TCP + region = (off, total) + pages = link.arena.view(off, nb * transfer.page_bytes).view(nb, transfer.page_bytes) + state = link.arena.view(off + state_off, transfer.state_bytes) + ev = transfer.export_into(seq, prompt_len, pages, state) + if region is None: + pages, state, ev = transfer.export(seq, prompt_len) + item = { + "pages": pages, + "state": state, + "region": region, + "event": ev, + "prompt_len": prompt_len, + "target": targets.pop(rid), + "first_token": None, + "ready": threading.Event(), + } + awaiting_token[rid] = item + send_q.put((rid, item)) + + engine.pe_do_save_kvcache = on_prefill_step + + def send_mooncake(rid: str, item: dict, host: str, port: int) -> None: + pages, state = item["pages"], item["state"] + hdr = MC_MAGIC + HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel()) + off, total = item["region"] + deadline = time.monotonic() + 600 + while True: + with socket.create_connection((host, port), timeout=60) as s: + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + s.sendall(hdr) + resv = json.loads(_recv_line(s)) + if resv["status"] == "ok": + for attempt in range(MC_WRITE_ATTEMPTS): + try: + link.write(resv["segment"], off, resv["addr"], total) + break + except Exception as e: # noqa: BLE001 + print( + f"[pe rank {dp_rank}] mooncake write to {resv['segment']} failed (attempt {attempt + 1}): {e!r}", flush=True + ) + time.sleep(0.01 * (attempt + 1)) + else: + s.sendall(b"xx") + return False # caller falls back to TCP + s.sendall(b"go") + ack = s.recv(2) + if ack != b"ok": + raise RuntimeError(f"CE rejected transfer: {ack!r}") + return True + if resv["status"] != "busy" or time.monotonic() > deadline: + raise RuntimeError(f"CE reservation failed: {resv}") + time.sleep(MC_RETRY_S) # CE arena full: retry once it has admitted something + + def sender(): + if link is not None: + import torch + + torch.cuda.set_device(transfer.device) # TransferEngine copies use the calling thread's device + while True: + rid, item = send_q.get() + host, port = item["target"] + try: + item["ready"].wait(timeout=600) + if item["first_token"] is None: + raise RuntimeError("no first token produced for the prefilled sequence") + item["event"].synchronize() + if item["region"] is not None: + try: + sent = send_mooncake(rid, item, host, port) + if not sent: # copy the region to pinned host buffers and use the TCP path below + item["pages"] = item["pages"].to("cpu").pin_memory() + item["state"] = item["state"].to("cpu").pin_memory() + finally: + link.arena.release(*item["region"]) + if sent: + continue + pages, state = item["pages"], item["state"] + with socket.create_connection((host, port), timeout=60) as s: + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + s.sendall( + HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel()) + ) + s.sendall(memoryview(pages.numpy()).cast("B")) + s.sendall(memoryview(state.numpy()).cast("B")) + ack = s.recv(2) + if ack != b"ok": + out_q.put(("error", rid, f"CE rejected transfer: {ack!r}")) + except Exception as e: # noqa: BLE001 + out_q.put(("error", rid, f"KV transfer to {host}:{port} failed: {e!r}")) + + threading.Thread(target=sender, daemon=True).start() + out_q.put(("ready", dp_rank, None)) + + pending = collections.deque() + inflight = False + last_stats = time.monotonic() + while True: + while True: + try: + kind, rid, payload = in_q.get_nowait() + except queue.Empty: + break + if kind == "add": + try: + prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32) + sp = SamplingParams(**payload["sampling"]) + sp.max_tokens = 1 # prefill + first token only; the engine then finishes and frees the sequence + seq = engine.add_request( + user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=prompt_ids, multi_modal_data=mm(), enqueue=False + ) + if seq is None: + out_q.put(("error", rid, "request dropped: invalid prompt tokens")) + else: + targets[rid] = tuple(payload["kv_target"]) + pending.append((rid, seq, prompt_ids)) + except Exception as e: # noqa: BLE001 + out_q.put(("error", rid, repr(e))) + elif kind == "abort": + targets.pop(rid, None) + if any(p[0] == rid for p in pending): + pending = collections.deque(p for p in pending if p[0] != rid) + else: + engine.abort_request(rid) + if cache is not None: + cache.forget(rid) + while pending: + rid, seq, prompt_ids = pending[0] + hit = cache.try_resume(seq, rid, prompt_ids) if cache is not None else 0 + if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight): + break + pending.popleft() + enqueue(rid, seq, max(hit, 0)) + if cache is not None: + cache.poll() + if time.monotonic() - last_stats > 5: + last_stats = time.monotonic() + out_q.put(("stats", dp_rank, cache.stats_dict())) + + if engine.has_unfinished_requests(): + outputs = engine.step() + inflight = True + elif inflight: + outputs = engine.flush_inflight_reqs() + inflight = False + else: + time.sleep(0.001) + continue + if len(outputs): + batch = [ + ( + RequestOutput.get_request_id(p), + RequestOutput.get_token_id(p), + RequestOutput.get_finished(p), + RequestOutput.get_finish_reason(p), + RequestOutput.get_prompt_len(p), + ) + for p in outputs + ] + for rid, token_id, finished, _, _ in batch: + item = awaiting_token.pop(rid, None) + if item is not None: + item["first_token"] = int(token_id) + item["ready"].set() + if finished and cache is not None: + cache.forget(rid) + out_q.put(("tokens", None, batch)) + # A prompt-only engine never reclaims finished sequences by itself (that happens in the decode + # scheduler): drop them from `running` and free their blocks/state. The export and cache copies were + # enqueued on the stream before any kernel that could reuse those blocks. + for seq in engine.scheduler.free_finished_seqs(): + engine.scheduler.free_seq(seq) + + +def _ce_loop( + args: argparse.Namespace, + dp_rank: int, + in_q: mp.Queue, + out_q: mp.Queue, + engine: LLMEngine, + transfer: PDTransfer, + SamplingParams: type[SamplingParams], + RequestOutput: type[RequestOutput], + np: ModuleType, + mm: Callable[[], dict | None], +) -> None: + import torch + + recv_q: queue.Queue = queue.Queue() # (rid, first_token, prompt_len, pages, state, arena region or None) + link = _mooncake_link(args, transfer.device) if args.kv_transport == "mooncake" else None + + def receiver(): + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind((args.host, args.kv_port + dp_rank)) + srv.listen(64) + while True: + conn, _ = srv.accept() + threading.Thread(target=receive_one, args=(conn,), daemon=True).start() + + def receive_mooncake(conn: socket.socket, rid: str, first_token: int, prompt_len: int, num_pages: int) -> None: + from deepinfer.hicache.pd_mooncake import region_layout + + state_off, total = region_layout(num_pages, transfer.page_bytes, transfer.state_bytes) + off = link.arena.alloc(total) if link is not None else None + if off is None: + status = "busy" if link is not None and total <= link.arena.nbytes else "error" + conn.sendall(json.dumps({"status": status}).encode() + b"\n") + return + try: + conn.sendall(json.dumps({"status": "ok", "addr": link.arena.base + off, "segment": link.segment}).encode() + b"\n") + go = bytearray(2) + _recv_exact(conn, memoryview(go)) + if go != b"go": + raise RuntimeError(f"PE aborted the transfer: {bytes(go)!r}") + except BaseException: + link.arena.release(off, total) + raise + pages = link.arena.view(off, num_pages * transfer.page_bytes).view(num_pages, transfer.page_bytes) + state = link.arena.view(off + state_off, transfer.state_bytes) + recv_q.put((rid, first_token, prompt_len, pages, state, (off, total))) + conn.sendall(b"ok") + + def receive_one(conn: socket.socket): + try: + with conn: + hdr = bytearray(HDR.size) + _recv_exact(conn, memoryview(hdr)[: len(MC_MAGIC)]) + mooncake = bytes(hdr[: len(MC_MAGIC)]) == MC_MAGIC + _recv_exact(conn, memoryview(hdr)[0 if mooncake else len(MC_MAGIC) :]) + rid_b, first_token, prompt_len, num_pages, page_bytes, state_bytes = HDR.unpack(hdr) + rid = rid_b.decode() + assert page_bytes == transfer.page_bytes and state_bytes == transfer.state_bytes, "layout mismatch between PE and CE" + if mooncake: + receive_mooncake(conn, rid, first_token, prompt_len, num_pages) + return + pages = torch.empty((num_pages, page_bytes), dtype=torch.uint8, pin_memory=True) + state = torch.empty(state_bytes, dtype=torch.uint8, pin_memory=True) + _recv_exact(conn, memoryview(pages.numpy()).cast("B")) + _recv_exact(conn, memoryview(state.numpy()).cast("B")) + recv_q.put((rid, first_token, prompt_len, pages, state, None)) + conn.sendall(b"ok") + except Exception as e: # noqa: BLE001 + print(f"[ce rank {dp_rank}] transfer failed: {e!r}", flush=True) + + threading.Thread(target=receiver, daemon=True).start() + out_q.put(("ready", dp_rank, None)) + + expected: dict[str, dict] = {} # rid -> {"prompt_token_ids", "sampling"} from the frontend + arrived: dict[str, tuple] = {} # rid -> transfer + pending: collections.deque = collections.deque() # rids with both parts, waiting for GPU allocation + keep: dict[str, tuple] = {} # rid -> (admit event, pages, state, arena region or None) + + def release(region: tuple[int, int] | None) -> None: + if region is not None: + link.arena.release(*region) + + inflight = False + while True: + while True: + try: + kind, rid, payload = in_q.get_nowait() + except queue.Empty: + break + if kind == "expect": + expected[rid] = payload + if rid in arrived: + pending.append(rid) + elif kind == "abort": + expected.pop(rid, None) + gone = arrived.pop(rid, None) + if gone is not None: + release(gone[4]) + if rid in pending: + pending.remove(rid) + else: + engine.abort_request(rid) + while True: + try: + rid, first_token, prompt_len, pages, state, region = recv_q.get_nowait() + except queue.Empty: + break + arrived[rid] = (first_token, prompt_len, pages, state, region) + if rid in expected: + pending.append(rid) + while pending: + rid = pending[0] + spec, (first_token, prompt_len, pages, state, region) = expected[rid], arrived[rid] + tokens = np.asarray(list(spec["prompt_token_ids"]) + [first_token], dtype=np.int32) + assert len(tokens) - 1 == prompt_len, (len(tokens), prompt_len) + sp = SamplingParams(**spec["sampling"]) + sp.max_tokens = max(1, sp.max_tokens - 1) # the first token was produced on the PE + seq = engine.add_request( + user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=tokens, multi_modal_data=mm(), enqueue=False + ) + if seq is None: + out_q.put(("error", rid, "request dropped: invalid prompt tokens")) + pending.popleft() + expected.pop(rid, None) + arrived.pop(rid, None) + release(region) + continue + try: + allocated = engine.scheduler.try_allocate_seq(seq) + except RuntimeError as e: + out_q.put(("error", rid, f"cannot allocate: {e!r}")) + allocated = None + if allocated is False and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight): + break # retry once something finishes + pending.popleft() + expected.pop(rid, None) + arrived.pop(rid, None) + if not allocated: + if allocated is False: + out_q.put(("error", rid, "cannot allocate KV blocks for the transferred sequence")) + release(region) + continue + # Keep the pinned receive buffers alive until the H2D page-copy kernels have consumed them: the kernel + # reads them by raw pointer, so the caching host allocator would otherwise hand the memory to the next + # transfer (or unmap it) while the copy is still queued -> corrupted KV / cudaErrorLaunchFailure. + keep[rid] = (transfer.admit(seq, pages, state, prompt_len), pages, state, region) + if engine.multi_modal_processor is not None: + engine.multi_modal_processor.add(rid, seq.multi_modal_data, prompt_len) + for rid in [r for r, k in keep.items() if k[0].query()]: + release(keep.pop(rid)[3]) + + if engine.has_unfinished_requests(): + outputs = engine.step() + inflight = True + elif inflight: + outputs = engine.flush_inflight_reqs() + inflight = False + else: + time.sleep(0.001) + continue + if len(outputs): + batch = [ + ( + RequestOutput.get_request_id(p), + RequestOutput.get_token_id(p), + RequestOutput.get_finished(p), + RequestOutput.get_finish_reason(p), + RequestOutput.get_prompt_len(p), + ) + for p in outputs + ] + out_q.put(("tokens", None, batch)) + + +# ============================================================================ frontends + + +class PDFrontend(oas.Frontend): + def __init__(self, args: argparse.Namespace) -> None: + super().__init__(args) + ctx = mp.get_context("spawn") + self.out_q = ctx.Queue() + self.in_qs = [ctx.Queue() for _ in range(args.dp)] + self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)] + + +def build_pe_app(fe: PDFrontend) -> FastAPI: + """Prefill tray: POST /prefill {rid, prompt_token_ids, sampling, kv_target:[host,port]} -> {first_token}.""" + from fastapi import FastAPI + from fastapi.responses import JSONResponse + + app = FastAPI() + + @app.on_event("startup") + async def _startup() -> None: + fe.loop = asyncio.get_running_loop() + threading.Thread(target=fe.dispatch_forever, daemon=True).start() + + @app.get("/health") + async def health(): + ok = fe.ready == fe.args.dp + return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready, "role": "pe"}, status_code=200 if ok else 503) + + @app.get("/deepinfer/metrics") + async def metrics(): + lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()] + for r, st in sorted(fe.cache_stats.items()): + lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()] + return "\n".join(lines) + "\n" + + @app.post("/prefill") + async def prefill(request: Request): + if fe.ready != fe.args.dp: + return JSONResponse({"error": "engine loading"}, status_code=503) + body = await request.json() + rid = body["rid"] + rank = min(range(fe.args.dp), key=lambda r: fe.load[r]) + key = bytes.fromhex(body["affinity"]) if body.get("affinity") else None + if fe.args.routing == "affinity" and key is not None: + rank = int.from_bytes(key[:4], "little") % fe.args.dp + fe.load[rank] += 1 + fe.owner[rid] = rank + q: asyncio.Queue = asyncio.Queue() + fe.streams[rid] = q + fe.stats["requests"] += 1 + fe.in_qs[rank].put( + ("add", rid, {"prompt_token_ids": body["prompt_token_ids"], "sampling": body["sampling"], "kv_target": body["kv_target"]}) + ) + try: + item = await asyncio.wait_for(q.get(), timeout=1800) + finally: + fe.release(rid, aborted=False) + if item[0] == "error": + return JSONResponse({"error": item[1]}, status_code=500) + return {"first_token": item[1], "prompt_len": item[4]} + + return app + + +def build_ce_app(fe: PDFrontend, pe_url: str, kv_host: str) -> FastAPI: + """Decode tray: the user-facing OpenAI API.""" + import httpx + from fastapi import FastAPI + from fastapi.responses import JSONResponse, StreamingResponse + + app = FastAPI() + client = httpx.AsyncClient(timeout=1800) + + @app.on_event("startup") + async def _startup() -> None: + fe.loop = asyncio.get_running_loop() + threading.Thread(target=fe.dispatch_forever, daemon=True).start() + + @app.get("/health") + async def health(): + ok = fe.ready == fe.args.dp + return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready, "role": "ce"}, status_code=200 if ok else 503) + + @app.get("/v1/models") + async def models(): + return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]} + + @app.get("/deepinfer/metrics") + async def metrics(): + lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()] + lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)] + return "\n".join(lines) + "\n" + + @app.post("/v1/chat/completions") + async def chat(request: Request): + if fe.ready != fe.args.dp: + return JSONResponse({"error": {"message": "engine loading"}}, status_code=503) + body = await request.json() + prompt_ids = await asyncio.to_thread(fe.render_prompt, body) + max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096 + if len(prompt_ids) + 1 > fe.args.max_model_len: + return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400) + max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids))) + sampling = dict( + temperature=float(body.get("temperature", 1.0)), + top_p=float(body.get("top_p", 1.0)), + ignore_eos=bool(body.get("ignore_eos", False)), + max_tokens=max_tokens, + ) + if body.get("seed") is not None: + sampling["seq_seed"] = int(body["seed"]) + fe.stats["prompt_tokens"] += len(prompt_ids) + include_usage = bool((body.get("stream_options") or {}).get("include_usage")) + key = fe.affinity_key(body) + + # Register with the CE rank first so the transfer can be admitted as soon as it lands. + rid = uuid.uuid4().hex[:16] + rank = ( + int.from_bytes(key[:4], "little") % fe.args.dp + if (fe.args.routing == "affinity" and key) + else min(range(fe.args.dp), key=lambda r: fe.load[r]) + ) + fe.load[rank] += 1 + fe.owner[rid] = rank + q: asyncio.Queue = asyncio.Queue() + fe.streams[rid] = q + fe.stats["requests"] += 1 + fe.in_qs[rank].put(("expect", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling})) + + created = int(time.time()) + cid = f"chatcmpl-{rid}" + detok = oas.IncrementalDetokenizer(fe.tokenizer) + + def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str: + obj = { + "id": cid, + "object": "chat.completion.chunk", + "created": created, + "model": fe.model_name, + "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}], + } + if usage: + obj["usage"] = usage + return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n" + + async def prefill_first(): + r = await client.post( + f"{pe_url}/prefill", + json={ + "rid": rid, + "prompt_token_ids": prompt_ids, + "sampling": sampling, + "kv_target": [kv_host, fe.args.kv_port + rank], + "affinity": key.hex() if key else None, + }, + ) + if r.status_code != 200: + raise RuntimeError(f"PE: {r.text[:200]}") + return int(r.json()["first_token"]) + + async def tokens(): + first = await prefill_first() + text = detok.add(first) + fe.stats["completion_tokens"] += 1 + if max_tokens == 1 or first == fe.tokenizer.eos_token_id and not sampling["ignore_eos"]: + yield text, ("length" if max_tokens == 1 else "stop") + return + yield text, None + while True: + item = await q.get() + if item[0] == "error": + raise RuntimeError(item[1]) + _, token_id, finished, reason, _ = item + text = detok.add(token_id) + fe.stats["completion_tokens"] += 1 + if finished: + yield text, ("length" if detok.num_tokens >= max_tokens else "stop") + return + yield text, None + + usage = lambda: { + "prompt_tokens": len(prompt_ids), + "completion_tokens": detok.num_tokens, # noqa: E731 + "total_tokens": len(prompt_ids) + detok.num_tokens, + } + + if body.get("stream"): + + async def sse(): + done = False + try: + yield chunk({"role": "assistant", "content": ""}) + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) + if text or finish: + yield chunk({"content": text} if text else {}, finish) + done = True + if include_usage: + yield chunk({}, usage=usage()) + yield "data: [DONE]\n\n" + finally: + fe.release(rid, aborted=not done) + + return StreamingResponse(sse(), media_type="text/event-stream") + + done = False + try: + parts, finish = [], None + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) + parts.append(text) + done = True + finally: + fe.release(rid, aborted=not done) + return { + "id": cid, + "object": "chat.completion", + "created": created, + "model": fe.model_name, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}], + "usage": usage(), + } + + return app + + +def main() -> None: + args = parse_args() + args.model = str(Path(args.model).expanduser().resolve()) + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("RANK", "0") + fe = PDFrontend(args) + fe.start() + import uvicorn + + if args.role == "pe": + app = build_pe_app(fe) + else: + assert args.pe_url, "--pe-url is required for the decode tray" + app = build_ce_app(fe, args.pe_url.rstrip("/"), args.kv_host or socket.gethostname()) + uvicorn.run(app, host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900) + + +if __name__ == "__main__": + main() 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/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 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 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"