From eb1c7a51ab1cafecaf4e74006649924be7eb129e Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:11:38 +0000 Subject: [PATCH 1/7] Add a minimal OpenAI-compatible chat server over the offline engine serve/openai_server.py exposes /v1/chat/completions (streaming, tools, usage) on top of LLMEngine: one engine process per GPU (attention DP + MoE EP, the layout of examples/benchmark.py), prompts rendered with the checkpoint's own encoding.py, conversation-affinity routing, a watchdog that fails open requests if an engine rank dies, and /deepinfer/metrics. serve/server.sh runs it as a process group so a stop also reaps the engine's helper processes, which otherwise keep the GPUs. It integrates the host prefix cache from deepinfer/hicache when present (--hicache-size-gb, --hicache-state-slots, per-request "prefix_cache": false bypass; serve/check_prefix_cache.py is the A/B check) and serves without it otherwise. It also sets DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL to 2 by default (--prefill-any-interval): with the engine default of 100, a rank whose peers are idle admits a waiting prompt only every 100 decode steps, a ~1.4 s TTFT floor on every request in a serving setting. Used to drive the InferenceX AgentX replay (aiperf) on 4x GB300. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/README.md | 45 ++++ serve/check_prefix_cache.py | 87 ++++++ serve/openai_server.py | 518 ++++++++++++++++++++++++++++++++++++ serve/server.sh | 41 +++ 4 files changed, 691 insertions(+) create mode 100644 serve/README.md create mode 100644 serve/check_prefix_cache.py create mode 100644 serve/openai_server.py create mode 100755 serve/server.sh diff --git a/serve/README.md b/serve/README.md new file mode 100644 index 0000000..9f204eb --- /dev/null +++ b/serve/README.md @@ -0,0 +1,45 @@ +# Serving DeepInfer + host prefix cache (local additions, not part of upstream) + +## Pieces + +- `serve/openai_server.py` — OpenAI-compatible `/v1/chat/completions` (streaming, `tools`, usage) on top of the + offline engine: one engine process per GPU (attention DP + MoE EP), prompts rendered with the checkpoint's own + `encoding.py`. `serve/server.sh start|stop` runs it as a process group (stop reaps the engine helpers, which + otherwise keep the GPUs). +- `deepinfer/hicache/` — host-memory prefix cache (SGLang HiCache style): pinned host pools organized in pages + (one page = one 128-token KV block) plus SWA state slots, hash-chain index, LRU eviction, write-through at the + end of every prefill, "kernel" transfers (a CUDA kernel copying pages through UVA pointers on the engine + stream, `page_copy.cu`, JIT-built on first use). Needs the read-only `Sequence.block_table` binding added in + `csrc/engine/pybind_sequence.cpp` (rebuild `deepinfer._C` with `bash scripts/build.sh build`). +- `tests/hicache/test_host_cache.py` — cache logic against a fake engine (needs one GPU). +- `serve/check_prefix_cache.py` — A/B on a live server: each turn sent with `"prefix_cache": false` and with + the cache, greedy; compares outputs and TTFT. + +## Run + +```bash +export DEEPINFER_ACCELERATOR=nvidia DEEPINFER_HF_MODEL_DIR=/path/to/DeepSeek-V4.1-Flash-1017 +serve/server.sh start --dp 4 --port 8888 # cache on: 32 GiB pages + 4096 states per rank +serve/server.sh start --dp 4 --port 8888 --hicache-size-gb 0 # cache off +python serve/check_prefix_cache.py --port 8888 --turns 3 +curl -s localhost:8888/deepinfer/metrics | grep hicache # per-rank hit/store counters +serve/server.sh stop +``` + +Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-loaded` (affinity keeps a +conversation's turns on one rank, keyed by the first user message; required for hits since the cache is +per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, the default 140). + +## How a hit works + +A new prompt's page hashes are matched from the start; the match is capped so that at least `keep_tokens` +tokens remain to be recomputed (the decoder SWA state is rebuilt from that suffix by bounded replay), then +snapped down to the nearest 128-aligned position that has a stored SWA state. The sequence is allocated up +front, pages and state are copied in on the engine stream, `num_cached_tokens` is set, and the sequence is +enqueued; the scheduler prefills only the suffix. Generated tokens are not stored at decode end (they are +stored when they come back as part of the next turn's prompt). + +## AgentX + +The server speaks what InferenceX's AgentX replay (aiperf, `--endpoint /v1/chat/completions --streaming +--use-server-token-count`) needs; `prefill-any-interval` matters there, since a single lane keeps most ranks idle. diff --git a/serve/check_prefix_cache.py b/serve/check_prefix_cache.py new file mode 100644 index 0000000..bb8fae1 --- /dev/null +++ b/serve/check_prefix_cache.py @@ -0,0 +1,87 @@ +"""A/B check of the host prefix cache on a running serve/openai_server.py. + +Builds a multi-turn conversation with a long shared prefix, sends every turn twice (cache bypassed via +"prefix_cache": false, then with the cache) with greedy sampling, and compares outputs and TTFT. + + python serve/check_prefix_cache.py --port 8888 --turns 3 --prefix-tokens 6000 +""" + +from __future__ import annotations + +import argparse +import json +import time + +import requests + +p = argparse.ArgumentParser() +p.add_argument("--port", type=int, default=8888) +p.add_argument("--model", default="deepseek-ai/DeepSeek-V4.1-Flash-1017") +p.add_argument("--turns", type=int, default=3) +p.add_argument("--prefix-tokens", type=int, default=6000, help="approximate size of the shared document") +p.add_argument("--max-tokens", type=int, default=64) +args = p.parse_args() +URL = f"http://localhost:{args.port}/v1/chat/completions" + +# A deterministic pseudo-document; ~1 token per word for this tokenizer is close enough. +words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa", "lambda", "mu"] +doc = " ".join(f"{words[i % len(words)]}{i % 97}" for i in range(args.prefix_tokens)) +questions = [ + "Which word appears right after 'gamma2'? Answer briefly.", + "How many distinct word stems are used in the document? Answer briefly.", + "Quote the first five words of the document.", + "What is the last word of the document?", +] + + +def ask(messages: list[dict], use_cache: bool) -> tuple[str, float, float, dict]: + body = { + "model": args.model, + "messages": messages, + "stream": True, + "stream_options": {"include_usage": True}, + "temperature": 0, + "max_tokens": args.max_tokens, + "seed": 1, + "prefix_cache": use_cache, + } + t0 = time.perf_counter() + ttft = None + text, usage = [], {} + with requests.post(URL, json=body, stream=True, timeout=1800) as r: + r.raise_for_status() + for line in r.iter_lines(): + if not line or not line.startswith(b"data: ") or line == b"data: [DONE]": + continue + obj = json.loads(line[6:]) + if obj.get("usage"): + usage = obj["usage"] + for ch in obj.get("choices", []): + delta = ch["delta"].get("content") + if delta: + if ttft is None: + ttft = time.perf_counter() - t0 + text.append(delta) + return "".join(text), ttft or 0.0, time.perf_counter() - t0, usage + + +messages = [{"role": "system", "content": "You answer questions about the document below.\n\n" + doc}] +mismatches = 0 +for turn in range(args.turns): + messages.append({"role": "user", "content": questions[turn % len(questions)]}) + base, ttft_b, tot_b, usage_b = ask(messages, use_cache=False) + cached, ttft_c, tot_c, usage_c = ask(messages, use_cache=True) + same = base == cached + mismatches += not same + print( + f"turn {turn + 1}: prompt {usage_c.get('prompt_tokens')} tok | TTFT no-cache {ttft_b * 1e3:7.0f} ms, cache {ttft_c * 1e3:7.0f} ms " + f"({ttft_b / max(ttft_c, 1e-6):.1f}x) | outputs {'identical' if same else 'DIFFER'}" + ) + if not same: + print(" no-cache:", repr(base[:200])) + print(" cache :", repr(cached[:200])) + messages.append({"role": "assistant", "content": base}) + +stats = requests.get(f"http://localhost:{args.port}/deepinfer/metrics", timeout=30).text +print("\n".join(line for line in stats.splitlines() if "hicache:hits" in line or "hicache:lookups" in line or "hit_rate" in line)) +raise SystemExit(1 if mismatches else 0) diff --git a/serve/openai_server.py b/serve/openai_server.py new file mode 100644 index 0000000..d34acdc --- /dev/null +++ b/serve/openai_server.py @@ -0,0 +1,518 @@ +"""Minimal OpenAI-compatible chat server on top of the DeepInfer offline engine. + +DeepInfer ships offline inference only. This wraps one `LLM` per data-parallel rank +(attention DP + MoE EP, same layout as examples/benchmark.py) behind a FastAPI frontend +so request-replay benchmarks such as InferenceX AgentX (aiperf) can drive it. + + python serve/openai_server.py --model ~/models/DeepSeek-V4.1-Flash-1017 --dp 4 --port 8888 + +Limitations, by design of the open-source engine: +- No cross-request prefix cache: every request prefills its full prompt. +- Responses are streamed as plain `content`; DSML tool calls are not parsed into `tool_calls`. +""" + +from __future__ import annotations + +import argparse +import asyncio +import collections +import hashlib +import json +import multiprocessing as mp +import os +import queue +import sys +import threading +import time +import uuid +from pathlib import Path + +# Module level so FastAPI can resolve the string annotation `request: Request`. +from fastapi import Request + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "examples")) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument("--model", required=True, help="local HF checkpoint directory") + p.add_argument("--served-model-name", default=None) + p.add_argument("--dp", type=int, default=4, help="data-parallel ranks on this node (= EP size)") + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=8888) + p.add_argument("--max-model-len", type=int, default=1 << 20) + p.add_argument("--max-num-seqs", type=int, default=128, help="per DP rank") + p.add_argument("--max-num-batched-tokens", type=int, default=32768, help="prefill chunk per DP rank") + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--yoco-keep-tokens", type=int, default=140, help="0 disables decoder SWA bounded replay") + p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"]) + p.add_argument("--log-interval-steps", type=int, default=200) + p.add_argument("--hicache-size-gb", type=float, default=32, help="pinned host KV page pool per DP rank; 0 disables the prefix cache") + p.add_argument("--hicache-state-slots", type=int, default=4096, help="host SWA state slots per DP rank") + p.add_argument( + "--routing", + default="affinity", + choices=["affinity", "least-loaded"], + help="affinity: a conversation's turns go to one rank (needed for cache hits)", + ) + p.add_argument( + "--prefill-any-interval", + type=int, + default=2, + help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this " + "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor)", + ) + return p.parse_args() + + +# --------------------------------------------------------------------------- engine side + + +def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None: + """One DP rank: owns an engine and steps it in lockstep with the other ranks.""" + import ctypes + import signal + + # Die with the frontend instead of lingering as an orphan that holds the GPU. + ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL) # PR_SET_PDEATHSIG + + from _config import create_llm, setup_common_env + + if args.yoco_keep_tokens > 0: + os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact" + os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens) + os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1" + if args.hicache_size_gb > 0: + # Keep more SWA history in the state ring so the prefix cache can store states at several + # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks). + os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") + os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval)) + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + + llm = create_llm( + args.model, + dp_rank, + ep_size=args.dp, + env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)}, + data_parallel_size=args.dp, + max_num_seqs=args.max_num_seqs, + max_num_batched_tokens=args.max_num_batched_tokens, + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))), + ) + + import numpy as np + + from deepinfer._C.engine import RequestOutput # pyright: ignore[reportMissingModuleSource] + from deepinfer.sampling_params import SamplingParams + + engine = llm.llm_engine + cache = None + if args.hicache_size_gb > 0: + try: + from deepinfer.hicache import HostPrefixCache + except ImportError: # tree without the host prefix cache: serve without it + print(f"[server] deepinfer.hicache not available, prefix cache disabled (rank {dp_rank})", flush=True) + else: + cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots) + engine.pe_do_save_kvcache = cache.on_prefill_step # write-through at the end of every prefill + out_q.put(("ready", dp_rank, None)) + + def mm_data(): + return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None + + def enqueue(rid, seq, num_cached_tokens): + # Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict + # (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length. + engine.scheduler.add_sequence(seq) + if engine.multi_modal_processor is not None: + engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached_tokens) + + pending = collections.deque() # (rid, seq, prompt_ids, use_cache): created, waiting for a cache decision + inflight = False + last_stats = time.monotonic() + last_logged_lookups = 0 + while True: + # Requests must be enqueued before the DP-wide has_unfinished_requests all-reduce, + # so every rank takes the same step / flush / idle branch in the same iteration. + while True: + try: + kind, rid, payload = in_q.get_nowait() + except queue.Empty: + break + if kind == "add": + try: + prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32) + seq = engine.add_request( + user_id="0", + request_id=rid, + sampling_params=SamplingParams(**payload["sampling"]), + prompt_token_ids=prompt_ids, + multi_modal_data=mm_data(), + enqueue=cache is None, + ) + if cache is not None: + if seq is None: + out_q.put(("error", rid, "request dropped: invalid prompt tokens")) + else: + pending.append((rid, seq, prompt_ids, payload.get("use_cache", True))) + except Exception as e: # noqa: BLE001 - surface to the client, keep the rank alive + out_q.put(("error", rid, repr(e))) + elif kind == "abort": + if any(p[0] == rid for p in pending): + pending = collections.deque(p for p in pending if p[0] != rid) + else: + engine.abort_request(rid) + if cache is not None: + cache.forget(rid) + + # Cache decisions, in arrival order. A hit pre-allocates the sequence's blocks and fills them from the + # host; if the GPU cannot allocate yet, wait (while something is running that will free memory). + while pending: + rid, seq, prompt_ids, use_cache = pending[0] + hit = cache.try_resume(seq, rid, prompt_ids) if use_cache else 0 + if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight): + break + pending.popleft() + enqueue(rid, seq, max(hit, 0)) + if cache is not None: + cache.poll() + if time.monotonic() - last_stats > 5: + last_stats = time.monotonic() + st = cache.stats_dict() + out_q.put(("stats", dp_rank, st)) + if st["lookups"] != last_logged_lookups: + last_logged_lookups = st["lookups"] + print( + f"[hicache rank {dp_rank}] lookups {st['lookups']} hits {st['hits']} " + f"token hit rate {st['hit_rate_tokens']:.1%} deferred {st['deferred']} stores {st['stores']} " + f"pages {st['pages_used']}/{st['pages_total']} states {st['states_used']}/{st['states_total']} " + f"evicted {st['evicted_pages']}", + flush=True, + ) + + if engine.has_unfinished_requests(): + outputs = engine.step() + inflight = True + elif inflight: + outputs = engine.flush_inflight_reqs() + inflight = False + else: + time.sleep(0.001) + continue + + if len(outputs): + # Pointers are recycled on the next step, so read the fields now. + batch = [ + ( + RequestOutput.get_request_id(p), + RequestOutput.get_token_id(p), + RequestOutput.get_finished(p), + RequestOutput.get_finish_reason(p), + RequestOutput.get_prompt_len(p), + ) + for p in outputs + ] + if cache is not None: + for rid, _, finished, _, _ in batch: + if finished: + cache.forget(rid) + out_q.put(("tokens", None, batch)) + + +# --------------------------------------------------------------------------- HTTP side + + +class IncrementalDetokenizer: + """vLLM-style prefix/read offsets so each delta decodes only a short tail.""" + + def __init__(self, tokenizer) -> None: + self.tok = tokenizer + self.ids: list[int] = [] + self.num_tokens = 0 + self.prefix_offset = 0 + self.read_offset = 0 + + def add(self, token_id: int) -> str: + self.num_tokens += 1 + if token_id == self.tok.eos_token_id: + return "" + self.ids.append(token_id) + prefix = self.tok.decode(self.ids[self.prefix_offset : self.read_offset], skip_special_tokens=False) + full = self.tok.decode(self.ids[self.prefix_offset :], skip_special_tokens=False) + if len(full) > len(prefix) and not full.endswith("�"): + self.prefix_offset, self.read_offset = self.read_offset, len(self.ids) + return full[len(prefix) :] + return "" + + +class Frontend: + def __init__(self, args: argparse.Namespace) -> None: + from transformers import AutoTokenizer + + sys.path.insert(0, str(Path(args.model) / "encoding")) + from encoding import encode_messages # the checkpoint's own prompt-format reference + + self.args = args + self.model_name = args.served_model_name or args.model + self.encode_messages = encode_messages + self.tokenizer = AutoTokenizer.from_pretrained(args.model) + ctx = mp.get_context("spawn") + self.out_q: mp.Queue = ctx.Queue() + self.in_qs: list[mp.Queue] = [ctx.Queue() for _ in range(args.dp)] + # Not daemonic: the engine forks its own helper processes (as examples/_launch.py allows). + self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)] + self.load = [0] * args.dp + self.owner: dict[str, int] = {} + self.streams: dict[str, asyncio.Queue] = {} + self.ready = 0 + self.stats = dict(requests=0, finished=0, aborted=0, errors=0, prompt_tokens=0, completion_tokens=0, encode_fallbacks=0) + self.cache_stats: dict[int, dict] = {} + self.loop: asyncio.AbstractEventLoop | None = None + + def start(self) -> None: + for p in self.procs: + p.start() + + def dispatch_forever(self) -> None: + """Thread: fan engine outputs out to the per-request asyncio queues.""" + while True: + try: + kind, rid, payload = self.out_q.get(timeout=1.0) + except queue.Empty: + dead = [r for r, p in enumerate(self.procs) if not p.is_alive()] + if dead: + print(f"[server] engine rank(s) {dead} died; failing open requests and exiting", flush=True) + for rid in list(self.streams): + self.loop.call_soon_threadsafe(self._deliver, rid, ("error", f"engine rank died: {dead}")) + time.sleep(2) + os._exit(1) + continue + if kind == "ready": + self.ready += 1 + print(f"[server] DP rank {rid} ready ({self.ready}/{self.args.dp})", flush=True) + elif kind == "error": + self.loop.call_soon_threadsafe(self._deliver, rid, ("error", payload)) + elif kind == "stats": + self.cache_stats[rid] = payload + else: + self.loop.call_soon_threadsafe(self._deliver_batch, payload) + + def _deliver(self, rid: str, item) -> None: + q = self.streams.get(rid) + if q is not None: + q.put_nowait(item) + + def _deliver_batch(self, batch) -> None: + for rid, token_id, finished, reason, prompt_len in batch: + q = self.streams.get(rid) + if q is not None: + q.put_nowait(("token", token_id, finished, reason, prompt_len)) + + def render_prompt(self, body: dict) -> list[int]: + messages = [dict(m) for m in body["messages"]] + if body.get("tools"): + if not messages or messages[0].get("role") != "system": + messages.insert(0, {"role": "system", "content": ""}) + messages[0]["tools"] = body["tools"] + try: + prompt = self.encode_messages(messages, thinking_mode=self.args.thinking_mode) + except Exception: # noqa: BLE001 - trace messages the reference encoder rejects + # Keep the request (and its prompt length) instead of failing the benchmark. + self.stats["encode_fallbacks"] += 1 + parts = ["<｜begin▁of▁sentence｜>"] + for m in messages: + content = m.get("content") or "" + if isinstance(content, list): + content = "".join(b.get("text", "") for b in content if isinstance(b, dict)) + if m.get("tool_calls"): + content += json.dumps(m["tool_calls"], ensure_ascii=False) + tag = {"user": "<｜User｜>", "assistant": "<｜Assistant｜>", "tool": "<｜User｜>"}.get(m.get("role"), "") + parts.append(tag + content) + if body.get("tools"): + parts.insert(1, json.dumps(body["tools"], ensure_ascii=False)) + prompt = "".join(parts) + "<｜Assistant｜>" + return self.tokenizer.encode(prompt, add_special_tokens=False) + + @staticmethod + def affinity_key(body: dict) -> bytes | None: + """Stable per-conversation key: the first user message (AgentX's cache-bust marker lives there and is + constant across a trajectory's turns, and forked subagents inherit it).""" + for m in body.get("messages", []): + if m.get("role") == "user": + content = m.get("content") or "" + if isinstance(content, list): + content = "".join(b.get("text", "") for b in content if isinstance(b, dict)) + return hashlib.sha1(content[:4096].encode("utf-8", "ignore")).digest() + return None + + def submit(self, prompt_ids: list[int], sampling: dict, key: bytes | None = None, use_cache: bool = True) -> tuple[str, asyncio.Queue]: + rid = uuid.uuid4().hex + if self.args.routing == "affinity" and key is not None: + rank = int.from_bytes(key[:4], "little") % self.args.dp + else: + rank = min(range(self.args.dp), key=lambda r: self.load[r]) + self.load[rank] += 1 + self.owner[rid] = rank + q: asyncio.Queue = asyncio.Queue() + self.streams[rid] = q + self.stats["requests"] += 1 + self.in_qs[rank].put(("add", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling, "use_cache": use_cache})) + return rid, q + + def release(self, rid: str, aborted: bool) -> None: + rank = self.owner.pop(rid, None) + self.streams.pop(rid, None) + if rank is None: + return + self.load[rank] -= 1 + if aborted: + self.stats["aborted"] += 1 + self.in_qs[rank].put(("abort", rid, None)) + else: + self.stats["finished"] += 1 + + +def build_app(fe: Frontend): + from fastapi import FastAPI + from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse + + app = FastAPI() + + @app.on_event("startup") + async def _startup() -> None: + fe.loop = asyncio.get_running_loop() + threading.Thread(target=fe.dispatch_forever, daemon=True).start() + + @app.get("/health") + async def health(): + ok = fe.ready == fe.args.dp + return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready}, status_code=200 if ok else 503) + + @app.get("/v1/models") + async def models(): + return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]} + + # Not at /metrics: aiperf scrapes that path automatically and InferenceX's aggregator + # rejects metric families it has no backend adapter for. + @app.get("/deepinfer/metrics") + async def metrics(): + lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()] + lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)] + for r, st in sorted(fe.cache_stats.items()): + lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()] + return PlainTextResponse("\n".join(lines) + "\n") + + @app.post("/v1/chat/completions") + async def chat(request: Request): + if fe.ready != fe.args.dp: + return JSONResponse({"error": {"message": "engine loading"}}, status_code=503) + body = await request.json() + prompt_ids = await asyncio.to_thread(fe.render_prompt, body) + max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096 + if len(prompt_ids) + 1 > fe.args.max_model_len: + return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400) + max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids))) + sampling = dict( + temperature=float(body.get("temperature", 1.0)), + top_p=float(body.get("top_p", 1.0)), + ignore_eos=bool(body.get("ignore_eos", False)), + max_tokens=max_tokens, + ) + if body.get("seed") is not None: + sampling["seq_seed"] = int(body["seed"]) + fe.stats["prompt_tokens"] += len(prompt_ids) + include_usage = bool((body.get("stream_options") or {}).get("include_usage")) + # Non-standard request field: "prefix_cache": false bypasses the host prefix cache lookup (A/B checks). + rid, q = fe.submit(prompt_ids, sampling, fe.affinity_key(body), use_cache=bool(body.get("prefix_cache", True))) + created = int(time.time()) + cid = f"chatcmpl-{rid}" + detok = IncrementalDetokenizer(fe.tokenizer) + + def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str: + obj = { + "id": cid, + "object": "chat.completion.chunk", + "created": created, + "model": fe.model_name, + "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}], + } + if usage: + obj["usage"] = usage + return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n" + + async def tokens(): + """Yield (text_delta, finish_reason) until the request finishes.""" + while True: + item = await q.get() + if item[0] == "error": + raise RuntimeError(item[1]) + _, token_id, finished, reason, _ = item + text = detok.add(token_id) + fe.stats["completion_tokens"] += 1 + if finished: + yield text, ("length" if detok.num_tokens >= max_tokens else "stop") + return + yield text, None + + usage = lambda: { + "prompt_tokens": len(prompt_ids), + "completion_tokens": detok.num_tokens, # noqa: E731 + "total_tokens": len(prompt_ids) + detok.num_tokens, + } + + if body.get("stream"): + + async def sse(): + done = False + try: + yield chunk({"role": "assistant", "content": ""}) + async for text, finish in tokens(): + if text or finish: + yield chunk({"content": text} if text else {}, finish) + done = True + if include_usage: + yield chunk({}, usage=usage()) + yield "data: [DONE]\n\n" + finally: + fe.release(rid, aborted=not done) + + return StreamingResponse(sse(), media_type="text/event-stream") + + done = False + try: + parts, finish = [], None + async for text, finish in tokens(): + parts.append(text) + done = True + finally: + fe.release(rid, aborted=not done) + return { + "id": cid, + "object": "chat.completion", + "created": created, + "model": fe.model_name, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}], + "usage": usage(), + } + + return app + + +def main() -> None: + args = parse_args() + args.model = str(Path(args.model).expanduser().resolve()) + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("RANK", "0") + fe = Frontend(args) + fe.start() + + import uvicorn + + uvicorn.run(build_app(fe), host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900) + + +if __name__ == "__main__": + main() diff --git a/serve/server.sh b/serve/server.sh new file mode 100755 index 0000000..b612fd2 --- /dev/null +++ b/serve/server.sh @@ -0,0 +1,41 @@ +#!/bin/bash +# Start/stop serve/openai_server.py as one process group, so stopping it also reaps the +# engine's helper processes (otherwise they outlive the frontend and keep the GPUs). +# +# serve/server.sh start [openai_server.py args...] +# serve/server.sh stop +set -euo pipefail +ROOT=$(cd "$(dirname "$0")/.." && pwd) +PIDFILE="$ROOT/serve/server.pgid" +LOG="${LOG:-$ROOT/serve.log}" + +case "${1:-}" in + start) + shift + [ -s "$PIDFILE" ] && kill -0 -- "-$(cat "$PIDFILE")" 2>/dev/null && { echo "already running (pgid $(cat "$PIDFILE"))" >&2; exit 1; } + [ -f "$ROOT/env.sh" ] && source "$ROOT/env.sh" # optional local environment (venv, DEEPINFER_* vars) + : "${DEEPINFER_HF_MODEL_DIR:?set DEEPINFER_HF_MODEL_DIR to the local checkpoint directory}" + export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" + setsid python -u "$ROOT/serve/openai_server.py" --model "$DEEPINFER_HF_MODEL_DIR" \ + --served-model-name deepseek-ai/DeepSeek-V4.1-Flash-1017 "$@" > "$LOG" 2>&1 < /dev/null & + echo $! > "$PIDFILE" + port=8888 prev= + for a in "$@"; do [[ "$prev" == --port ]] && port=$a; prev=$a; done + until curl -fs "localhost:$port/health" >/dev/null 2>&1; do + kill -0 "$(cat "$PIDFILE")" 2>/dev/null || { echo "server died; see $LOG" >&2; exit 1; } + sleep 3 + done + echo "ready (pgid $(cat "$PIDFILE"), log $LOG)" + ;; + stop) + [ -s "$PIDFILE" ] || { echo "not running"; exit 0; } + pgid=$(cat "$PIDFILE") + kill -TERM -- "-$pgid" 2>/dev/null || true + for _ in $(seq 30); do kill -0 -- "-$pgid" 2>/dev/null || break; sleep 1; done + kill -KILL -- "-$pgid" 2>/dev/null || true + rm -f "$PIDFILE" + echo "stopped" + ;; + *) + echo "usage: $0 start|stop" >&2; exit 2 ;; +esac From aaf1e62c9eab1ce5bd260a594dfaea613760e89f Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:26:07 +0000 Subject: [PATCH 2/7] serve: type annotations and B007 for ruff Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr --- serve/openai_server.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/serve/openai_server.py b/serve/openai_server.py index d34acdc..1e41bf2 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -26,9 +26,13 @@ import threading import time import uuid from pathlib import Path +from typing import TYPE_CHECKING, Any # Module level so FastAPI can resolve the string annotation `request: Request`. -from fastapi import Request +from fastapi import FastAPI, Request + +if TYPE_CHECKING: + from deepinfer._C.engine import Sequence REPO_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(REPO_ROOT / "examples")) @@ -123,7 +127,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m def mm_data(): return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None - def enqueue(rid, seq, num_cached_tokens): + def enqueue(rid: str, seq: Sequence, num_cached_tokens: int) -> None: # Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict # (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length. engine.scheduler.add_sequence(seq) @@ -228,7 +232,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m class IncrementalDetokenizer: """vLLM-style prefix/read offsets so each delta decodes only a short tail.""" - def __init__(self, tokenizer) -> None: + def __init__(self, tokenizer: Any) -> None: self.tok = tokenizer self.ids: list[int] = [] self.num_tokens = 0 @@ -300,12 +304,12 @@ class Frontend: else: self.loop.call_soon_threadsafe(self._deliver_batch, payload) - def _deliver(self, rid: str, item) -> None: + def _deliver(self, rid: str, item: tuple) -> None: q = self.streams.get(rid) if q is not None: q.put_nowait(item) - def _deliver_batch(self, batch) -> None: + def _deliver_batch(self, batch: list[tuple]) -> None: for rid, token_id, finished, reason, prompt_len in batch: q = self.streams.get(rid) if q is not None: @@ -375,7 +379,7 @@ class Frontend: self.stats["finished"] += 1 -def build_app(fe: Frontend): +def build_app(fe: Frontend) -> FastAPI: from fastapi import FastAPI from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse @@ -469,7 +473,7 @@ def build_app(fe: Frontend): done = False try: yield chunk({"role": "assistant", "content": ""}) - async for text, finish in tokens(): + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) if text or finish: yield chunk({"content": text} if text else {}, finish) done = True @@ -484,7 +488,7 @@ def build_app(fe: Frontend): done = False try: parts, finish = [], None - async for text, finish in tokens(): + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) parts.append(text) done = True finally: From d7844a221d7c7e8ec0694795a68dbfcad3bc9068 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:17:22 +0000 Subject: [PATCH 3/7] serve: EPLB rearrangement off by default, --eplb-interval / --eplb-after-prefill flags Hard-wired enable_eplb=True rearranged experts every 30 steps and after every prefill step; under multi-turn traffic that fires almost continuously and each rearrangement stalls all ranks for ~0.3 s. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/README.md | 4 ++++ serve/openai_server.py | 21 ++++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index 9f204eb..1f30b79 100644 --- a/serve/README.md +++ b/serve/README.md @@ -30,6 +30,10 @@ Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-l conversation's turns on one rank, keyed by the first user message; required for hits since the cache is per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, the default 140). +EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearrangement stalls all ranks for about +0.3 s, and with `--eplb-after-prefill 1` it fires after nearly every step under multi-turn traffic; enable it with +`--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. + ## How a hit works A new prompt's page hashes are matched from the start; the match is capped so that at least `keep_tokens` diff --git a/serve/openai_server.py b/serve/openai_server.py index 1e41bf2..e117e1b 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -67,6 +67,22 @@ def parse_args() -> argparse.Namespace: help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this " "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor)", ) + p.add_argument( + "--eplb-interval", + type=int, + default=0, + help="DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL: rearrange experts every N steps; 0 (default) disables EPLB. " + "Each rearrangement stalls every rank (~0.3 s on B200/GB300), which under online traffic costs far more than " + "the balance it buys", + ) + p.add_argument( + "--eplb-after-prefill", + type=int, + default=0, + choices=[0, 1], + help="DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE: also rearrange after every prefill step (only with " + "--eplb-interval > 0); with agentic traffic nearly every step batch contains a prefill", + ) return p.parse_args() @@ -92,7 +108,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks). os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval)) - setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False) + if args.eplb_interval > 0: + os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval) + os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill) llm = create_llm( args.model, From 269a61f7c54a2d56145bee46c566808541cd3467 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:49:52 +0000 Subject: [PATCH 4/7] serve: exact fast prompt tokenization (fastokens or added-token segment cache), checked against tokenizer.encode Agentic prompts are re-tokenized in full every turn; tokenizer.encode takes ~280 ms p50 on AgentX prompts (p50 163K tokens) and dominates TTFT once the prefix cache hits. serve/prompt_tokenizer.py picks fastokens if installed, else a dependency-free segment cache that splits at added tokens (where HF already encodes pieces independently) and reuses earlier turns' pieces; both are enabled only after matching tokenizer.encode on startup probes, and a sampled background re-check switches to hf on any mismatch. Suggested by whn09 in review. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 18 ++ serve/openai_server.py | 33 +++- serve/prompt_tokenizer.py | 264 +++++++++++++++++++++++++++ tests/serve/__init__.py | 0 tests/serve/test_prompt_tokenizer.py | 136 ++++++++++++++ 5 files changed, 450 insertions(+), 1 deletion(-) create mode 100644 serve/prompt_tokenizer.py create mode 100644 tests/serve/__init__.py create mode 100644 tests/serve/test_prompt_tokenizer.py diff --git a/serve/README.md b/serve/README.md index 1f30b79..20b2f1f 100644 --- a/serve/README.md +++ b/serve/README.md @@ -34,6 +34,24 @@ EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearran 0.3 s, and with `--eplb-after-prefill 1` it fires after nearly every step under multi-turn traffic; enable it with `--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. +## Prompt tokenization + +Agentic clients resend the whole conversation every turn, and `tokenizer.encode` needs ~1.7 ms per 1K tokens on a +GB300 Grace core (p50 278 ms per request on replayed AgentX prompts of p50 163K tokens), which dominates TTFT once the +prefix cache hits. `serve/prompt_tokenizer.py` gives the same ids faster; `--tokenizer-backend` picks one: + +- `fastokens` (`pip install fastokens`; x86-64 and aarch64 wheels): Rust reimplementation of the HF pipeline, built + from the loaded tokenizer. 3.5 ms p50 on the same prompts. +- `segment` (no dependency): splits the prompt at added tokens, where HF tokenizers already encode the pieces + independently, and caches each piece's ids (LRU, 512 MiB), so a turn only encodes what it adds. 18 ms p50. + Enabled only when the tokenizer config makes the split exact (no `single_word` tokens, no `prepend_scheme="first"`). +- `hf`: `tokenizer.encode`. + +`auto` (default) takes the first of these that matches `tokenizer.encode` on startup probes built from the +tokenizer's added tokens. `--tokenizer-verify-rate` (default 0.005) re-encodes a sample of prompts with +`tokenizer.encode` in the background; a mismatch switches to `hf`. The backend, encode latency and check counters are +in `/deepinfer/metrics` (`deepinfer_tokenizer:*`). Tests: `tests/serve/test_prompt_tokenizer.py`. + ## How a hit works A new prompt's page hashes are matched from the start; the match is capped so that at least `keep_tokens` diff --git a/serve/openai_server.py b/serve/openai_server.py index e117e1b..772390d 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -36,6 +36,8 @@ if TYPE_CHECKING: REPO_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(REPO_ROOT / "examples")) +sys.path.insert(0, str(REPO_ROOT / "serve")) +from prompt_tokenizer import PromptTokenizer # noqa: E402 def parse_args() -> argparse.Namespace: @@ -83,6 +85,19 @@ def parse_args() -> argparse.Namespace: help="DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE: also rearrange after every prefill step (only with " "--eplb-interval > 0); with agentic traffic nearly every step batch contains a prefill", ) + p.add_argument( + "--tokenizer-backend", + default="auto", + choices=["auto", "fastokens", "segment", "hf"], + help="prompt tokenization (serve/prompt_tokenizer.py); auto = the first of fastokens (if installed), segment, hf " + "that matches tokenizer.encode on startup probes. All produce tokenizer.encode's ids", + ) + p.add_argument( + "--tokenizer-verify-rate", + type=float, + default=0.005, + help="fraction of prompts re-encoded with tokenizer.encode in the background; a mismatch switches to hf", + ) return p.parse_args() @@ -282,6 +297,20 @@ class Frontend: self.model_name = args.served_model_name or args.model self.encode_messages = encode_messages self.tokenizer = AutoTokenizer.from_pretrained(args.model) + probe = [ + {"role": "system", "content": "sys", "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}]}, + {"role": "user", "content": "hi 你好 123"}, + {"role": "assistant", "content": "ok", "tool_calls": [{"type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "content": " out\n"}, + ] + try: + probes = [encode_messages(probe, thinking_mode=args.thinking_mode)] + except Exception: # noqa: BLE001 - the probe only adds coverage + probes = [] + # getattr: frontends built on this one with their own argument parser (serve/pd_server.py) get the defaults + backend, verify_rate = getattr(args, "tokenizer_backend", "auto"), getattr(args, "tokenizer_verify_rate", 0.005) + self.prompt_tokenizer = PromptTokenizer(self.tokenizer, backend, verify_rate, extra_probes=probes) + print(f"[server] {self.prompt_tokenizer.describe()}", flush=True) ctx = mp.get_context("spawn") self.out_q: mp.Queue = ctx.Queue() self.in_qs: list[mp.Queue] = [ctx.Queue() for _ in range(args.dp)] @@ -357,7 +386,7 @@ class Frontend: if body.get("tools"): parts.insert(1, json.dumps(body["tools"], ensure_ascii=False)) prompt = "".join(parts) + "<｜Assistant｜>" - return self.tokenizer.encode(prompt, add_special_tokens=False) + return self.prompt_tokenizer.encode(prompt) @staticmethod def affinity_key(body: dict) -> bytes | None: @@ -426,6 +455,8 @@ def build_app(fe: Frontend) -> FastAPI: lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)] for r, st in sorted(fe.cache_stats.items()): lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()] + pt = fe.prompt_tokenizer + lines += [f'deepinfer_tokenizer:{k}{{backend="{pt.name}"}} {v}' for k, v in pt.metrics().items()] return PlainTextResponse("\n".join(lines) + "\n") @app.post("/v1/chat/completions") diff --git a/serve/prompt_tokenizer.py b/serve/prompt_tokenizer.py new file mode 100644 index 0000000..feda0f0 --- /dev/null +++ b/serve/prompt_tokenizer.py @@ -0,0 +1,264 @@ +"""Prompt tokenization for the serving frontends: exact token ids, much faster than `tokenizer.encode` on long prompts. + +Agentic traffic re-sends the whole conversation every turn, so the frontend re-tokenizes 100K+ token prompts on +nearly every request. HF `tokenizer.encode` takes ~1.3 ms per 1K tokens (GB300 Grace core), which becomes the largest +part of TTFT once the prefix cache hits. Backends, all required to produce the ids of +`tokenizer.encode(text, add_special_tokens=False)`: + +- `fastokens` (optional `pip install fastokens`): a faster Rust implementation of the HF pipeline, built from the + tokenizer exactly as transformers configured it (`backend_tokenizer.to_str()`, not the file on disk). +- `segment`: no extra dependency. HF tokenizers first split the input at non-normalized added tokens and encode the + text between them independently. Splitting there ourselves and caching each piece's ids by its text is exact, and + a new turn only pays for the pieces it adds (the new tool result, the new user message). It is enabled only when + the tokenizer's config shows the split is exact (see `_segment_split_tokens`). +- `hf`: `tokenizer.encode`. + +`auto` takes the first of fastokens, segment, hf that is available and matches `tokenizer.encode` on a set of probes +built from the tokenizer's own added tokens. At runtime a sample of prompts is re-encoded with `tokenizer.encode` +on a background thread; the first mismatch switches the frontend to `hf`. +""" + +from __future__ import annotations + +import array +import collections +import contextlib +import json +import queue +import random +import re +import threading +import time +from typing import Any + +BACKENDS = ("auto", "fastokens", "segment", "hf") + + +class PromptTokenizer: + def __init__( + self, + tokenizer: Any, + backend: str = "auto", + verify_rate: float = 0.005, + cache_bytes: int = 512 << 20, + extra_probes: list[str] | None = None, + ) -> None: + assert backend in BACKENDS, backend + self.tokenizer = tokenizer + self.verify_rate = verify_rate + self.stats = dict(calls=0, errors=0, verified=0, verify_mismatches=0) + self.notes: list[str] = [] + self._times: collections.deque[float] = collections.deque(maxlen=4096) + self._probes = _probes(tokenizer) + list(extra_probes or []) + + candidates = ("fastokens", "segment") if backend == "auto" else (() if backend == "hf" else (backend,)) + self.name, self._encode = "hf", self._hf_encode + for name in candidates: + try: + encode = _FastokensBackend(tokenizer).encode if name == "fastokens" else _SegmentBackend(tokenizer, cache_bytes).encode + if (bad := self._first_mismatch(encode)) is not None: + raise ValueError(f"ids differ from tokenizer.encode on probe {bad!r}") + except Exception as e: # noqa: BLE001 - a backend that cannot be built or fails parity is skipped + self.notes.append(f"{name}: {type(e).__name__}: {e}"[:300]) + if backend != "auto": + raise RuntimeError(f"tokenizer backend {name} unavailable: {self.notes[-1]}") from e + continue + self.name, self._encode = name, encode + break + + self._verify_q: queue.Queue[tuple[str, list[int]]] = queue.Queue(maxsize=8) + if self.name != "hf" and verify_rate > 0: + threading.Thread(target=self._verify_forever, name="prompt-tokenizer-verify", daemon=True).start() + + def encode(self, text: str) -> list[int]: + t0 = time.perf_counter() + name = self.name + try: + ids = self._encode(text) + except Exception: # noqa: BLE001 - e.g. a regex engine limit on pathological input + self.stats["errors"] += 1 + name, ids = "hf", self._hf_encode(text) + self._times.append(time.perf_counter() - t0) + self.stats["calls"] += 1 + if name != "hf" and self.verify_rate > 0 and random.random() < self.verify_rate: + with contextlib.suppress(queue.Full): # the checker is behind: skip this sample + self._verify_q.put_nowait((text, ids)) + return ids + + def metrics(self) -> dict[str, float]: + """Counters plus encode latency percentiles (ms) over the last 4096 calls.""" + out: dict[str, float] = dict(self.stats) + times = sorted(self._times) + for q in (50, 90, 99): + out[f"encode_ms_p{q}"] = round(times[min(len(times) - 1, len(times) * q // 100)] * 1e3, 3) if times else 0.0 + return out + + def describe(self) -> str: + skipped = "; ".join(self.notes) + return f"prompt tokenizer backend: {self.name}" + (f" (skipped {skipped})" if skipped else "") + + def _hf_encode(self, text: str) -> list[int]: + return self.tokenizer.encode(text, add_special_tokens=False) + + def _first_mismatch(self, encode: Any) -> str | None: + for p in self._probes: + if list(encode(p)) != self._hf_encode(p): + return p[:80] + return None + + def _verify_forever(self) -> None: + while True: + text, ids = self._verify_q.get() + ref = self._hf_encode(text) + self.stats["verified"] += 1 + if ids != ref: + self.stats["verify_mismatches"] += 1 + at = next((i for i, (a, b) in enumerate(zip(ids, ref, strict=False)) if a != b), min(len(ids), len(ref))) + print( + f"[server] prompt tokenizer {self.name} differs from tokenizer.encode at token {at} of {len(ref)} " + f"(text {text[:60]!r}...); switching to hf", + flush=True, + ) + self.name, self._encode = "hf", self._hf_encode + + +class _FastokensBackend: + def __init__(self, tokenizer: Any) -> None: + from fastokens._native import Tokenizer # optional dependency + + if not getattr(tokenizer, "is_fast", False): + raise ValueError("needs a fast (tokenizer.json) tokenizer") + self._tok = Tokenizer.from_json_str(tokenizer.backend_tokenizer.to_str()) + + def encode(self, text: str) -> list[int]: + return self._tok.encode(text, add_special_tokens=False).ids + + +class _SegmentBackend: + """Split at added tokens, encode only the pieces not seen before (one batched call), LRU-cache pieces by text.""" + + def __init__(self, tokenizer: Any, cache_bytes: int) -> None: + if not getattr(tokenizer, "is_fast", False): + raise ValueError("needs a fast (tokenizer.json) tokenizer") + self._backend = tokenizer.backend_tokenizer + split = _segment_split_tokens(self._backend) + # Longest first: Python's alternation is leftmost-first, HF's matcher is leftmost-longest. lstrip/rstrip + # tokens swallow the whitespace next to them, as in HF (Unicode White_Space, Rust's char::is_whitespace). + alts = [ + f"{_WS}*" * t.lstrip + re.escape(t.content) + f"{_WS}*" * t.rstrip + for t in sorted(split.values(), key=lambda t: -len(t.content)) + ] + self._split = re.compile("(" + "|".join(alts) + ")") + self._ids = {c: t.id for c, t in split.items()} + self._stripped = any(t.lstrip or t.rstrip for t in split.values()) + # ids as uint32 arrays: 4 B/token instead of ~40 B for a list of Python ints + self._cache: collections.OrderedDict[str, array.array] = collections.OrderedDict() + self._cache_bytes = 0 + self._max_bytes = cache_bytes + self._lock = threading.Lock() + + def encode(self, text: str) -> list[int]: + parts = self._split.split(text) # odd indices are the split tokens + pieces: list[Any] = [None] * len(parts) + missing: list[int] = [] + with self._lock: + for i, p in enumerate(parts): + if i % 2: + pieces[i] = [self._ids[_strip_ws(p) if self._stripped else p]] + elif p: + ids = self._cache.get(p) + if ids is None: + missing.append(i) + else: + self._cache.move_to_end(p) + pieces[i] = ids + if missing: + encs = self._backend.encode_batch([parts[i] for i in missing], add_special_tokens=False) + with self._lock: + for i, e in zip(missing, encs, strict=True): + pieces[i] = e.ids + if parts[i] not in self._cache: + self._cache[parts[i]] = array.array("I", e.ids) + self._cache_bytes += _cost(parts[i], e.ids) + while self._cache_bytes > self._max_bytes and self._cache: + k, v = self._cache.popitem(last=False) + self._cache_bytes -= _cost(k, v) + out: list[int] = [] + for ids in pieces: + if ids: + out.extend(ids) + return out + + +_WS = "[\\t\\n\\x0b\\x0c\\r \\x85\\xa0\\u1680\\u2000-\\u200a\\u2028\\u2029\\u202f\\u205f\\u3000]" +_WS_CHARS = "\t\n\x0b\x0c\r \x85\xa0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000" + + +def _cost(text: str, ids: Any) -> int: + return len(text) + 4 * len(ids) + + +def _strip_ws(s: str) -> str: + return s.strip(_WS_CHARS) + + +_SplitToken = collections.namedtuple("_SplitToken", "content id lstrip rstrip") + + +def _segment_split_tokens(backend: Any) -> dict[str, _SplitToken]: + """Added tokens at which splitting the input and encoding the pieces separately gives tokenizer.encode's ids. + + HF's AddedVocabulary first matches the non-normalized added tokens on the raw input (leftmost-longest, an + lstrip/rstrip token also taking the whitespace next to it), and the rest of the pipeline (normalizer, + pre-tokenizer, model) then runs on each piece in between on its own. Raises when the config can make a piece + depend on its neighbours. + """ + cfg = json.loads(backend.to_str()) + added = cfg.get("added_tokens") or [] + if any(t["single_word"] for t in added): + raise ValueError("single_word added tokens depend on the characters around them") + if _has_first_only_prefix(cfg.get("pre_tokenizer")) or _has_first_only_prefix(cfg.get("normalizer")): + raise ValueError("prepend_scheme='first' prefixes only the first piece of the whole input") + split = {t["content"]: _SplitToken(t["content"], t["id"], t["lstrip"], t["rstrip"]) for t in added if not t["normalized"]} + if not split: + raise ValueError("no non-normalized added tokens to split at") + if any(_strip_ws(c) != c for c in split): + raise ValueError("an added token starts or ends with whitespace") + return split + + +def _has_first_only_prefix(node: Any) -> bool: + if isinstance(node, dict): + if node.get("prepend_scheme") == "first": + return True + return any(_has_first_only_prefix(v) for v in node.values()) + if isinstance(node, list): + return any(_has_first_only_prefix(v) for v in node) + return False + + +def _probes(tokenizer: Any) -> list[str]: + """Inputs that exercise added-token boundaries, scripts, digits and whitespace for this tokenizer.""" + added = [t.content for t in tokenizer.added_tokens_decoder.values()] + rng = random.Random(0) + sample = added if len(added) <= 64 else added[:16] + rng.sample(added, 48) + texts = [ + "", + "hello world", + " leading space", + "trailing space ", + "tabs\tand\r\nwindows\nlines\n\n\n", + " indented code:\n return x # comment\n", + "数字 12345678 和中文、日本語のテキスト、カタカナ", + "emoji 😀👍🏽 and combining é é ñ", + '{"name": "read_file", "arguments": {"path": "/tmp/a b.py", "n": 3.14}}', + "a" * 3000, + " " * 500 + "x", + "-" * 2000, + ] + probes = list(texts) + for i, tok in enumerate(sample): + a, b = texts[(i + 1) % len(texts)][:120], texts[(i + 5) % len(texts)][:120] + probes += [tok, f"{a}{tok}{b}", f"{a} {tok} {b}", f"{tok}{tok}", f"\n{tok}\n"] + probes.append("".join(f"{t}{texts[i % len(texts)][:40]}" for i, t in enumerate(sample))) + return probes diff --git a/tests/serve/__init__.py b/tests/serve/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/serve/test_prompt_tokenizer.py b/tests/serve/test_prompt_tokenizer.py new file mode 100644 index 0000000..cf18997 --- /dev/null +++ b/tests/serve/test_prompt_tokenizer.py @@ -0,0 +1,136 @@ +"""serve/prompt_tokenizer.py: every backend must give tokenizer.encode's ids (CPU only, needs the checkpoint's tokenizer).""" + +from __future__ import annotations + +import copy +import json +import os +import sys +import time +from pathlib import Path +from typing import Any + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT / "serve")) +from prompt_tokenizer import PromptTokenizer # noqa: E402 + +MODEL = os.environ.get("DEEPINFER_HF_MODEL_DIR") +pytestmark = pytest.mark.skipif(not MODEL or not Path(MODEL, "tokenizer.json").exists(), reason="DEEPINFER_HF_MODEL_DIR not set") + + +@pytest.fixture(scope="module") +def tokenizer() -> Any: + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(MODEL) + + +@pytest.fixture(scope="module") +def turns() -> list[str]: + """A growing tool-use conversation rendered with the checkpoint's encode_messages, one prompt per turn.""" + sys.path.insert(0, str(Path(MODEL, "encoding"))) + from encoding import encode_messages + + files = sorted((REPO_ROOT / "serve").glob("*.py")) + sorted((REPO_ROOT / "deepinfer" / "engine").glob("*.py"))[:8] + tools = [ + { + "type": "function", + "function": {"name": "read_file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}, + } + ] + msgs = [ + {"role": "system", "content": "You are a coding agent. 你是一个编程助手。", "tools": tools}, + {"role": "user", "content": "Summarise serve/. 12345"}, + ] + out = [] + for i, f in enumerate(files): + call = {"type": "function", "function": {"name": "read_file", "arguments": json.dumps({"path": str(f.name)})}} + msgs.append({"role": "assistant", "content": f"Step {i}: reading.", "tool_calls": [call]}) + msgs.append({"role": "tool", "content": f.read_text()}) + out.append(encode_messages(copy.deepcopy(msgs), thinking_mode="chat")) + return out + + +def available(tokenizer: Any) -> list[str]: + names = ["segment", "hf"] + try: + import fastokens # noqa: F401 + + names.insert(0, "fastokens") + except ImportError: + pass + return names + + +def test_backends_match_hf(tokenizer: Any, turns: list[str]) -> None: + for name in available(tokenizer): + pt = PromptTokenizer(tokenizer, backend=name, verify_rate=0) + assert pt.name == name + for prompt in turns: + assert pt.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False), name + + +def test_auto_picks_a_fast_backend(tokenizer: Any) -> None: + assert PromptTokenizer(tokenizer, verify_rate=0).name in ("fastokens", "segment") + + +def test_segment_cache_reuses_earlier_turns(tokenizer: Any, turns: list[str]) -> None: + pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0) + hf_s = seg_s = 0.0 + for prompt in turns: + t0 = time.perf_counter() + ref = tokenizer.encode(prompt, add_special_tokens=False) + t1 = time.perf_counter() + assert pt.encode(prompt) == ref + hf_s, seg_s = hf_s + t1 - t0, seg_s + time.perf_counter() - t1 + assert seg_s < hf_s + + +def test_segment_eviction_stays_exact(tokenizer: Any, turns: list[str]) -> None: + pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0, cache_bytes=4096) + for prompt in turns + turns[::-1]: + assert pt.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False) + + +def test_segment_with_strip_tokens(tokenizer: Any) -> None: + """Added tokens that swallow neighbouring whitespace (lstrip/rstrip, as Phi-3 and others use) stay exact.""" + from tokenizers import AddedToken + + tok = copy.deepcopy(tokenizer) + tok.add_tokens([AddedToken("<|l|>", lstrip=True, normalized=False), AddedToken("<|r|>", rstrip=True, normalized=False)]) + pt = PromptTokenizer(tok, backend="segment", verify_rate=0) + for text in ["a <|l|> b", "a<|r|> \n\t b", " <|l|><|r|> ", "x　<|l|> y<|r|> z", "<|r|> ", " <|l|>"]: + assert pt.encode(text) == tok.encode(text, add_special_tokens=False), repr(text) + + +def test_runtime_mismatch_switches_to_hf(tokenizer: Any) -> None: + pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=1.0) + pt._encode = lambda text: [0] # a backend gone wrong after startup + pt.encode("hello world") + for _ in range(200): + if pt.name == "hf": + break + time.sleep(0.01) + assert pt.name == "hf" and pt.stats["verify_mismatches"] == 1 + assert pt.encode("hello world") == tokenizer.encode("hello world", add_special_tokens=False) + + +def test_backend_error_falls_back_to_hf(tokenizer: Any) -> None: + pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0) + + def broken(text: str) -> list[int]: + raise RuntimeError("regex limit") + + pt._encode = broken + assert pt.encode("hello") == tokenizer.encode("hello", add_special_tokens=False) + assert pt.stats["errors"] == 1 + + +def test_forced_unavailable_backend_raises(tokenizer: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "fastokens", None) + monkeypatch.setitem(sys.modules, "fastokens._native", None) + with pytest.raises(RuntimeError, match="fastokens"): + PromptTokenizer(tokenizer, backend="fastokens") + assert PromptTokenizer(tokenizer, verify_rate=0).name == "segment" From 83be1fa44c34231911cfc223b9212d6364254462 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Sun, 25 Oct 2026 21:35:59 +0000 Subject: [PATCH 5/7] serve: drop the 32 redundant expert slots when EPLB is off Without rearrangement they are fixed copies of experts 0-31 and only add expert weight per rank. Decode SPS +4.8% / +3.9% at 4 / 16 requests per rank on GB300 (examples/benchmark.py, 128K context), matching whn09's B200 numbers; 6 GiB less model memory per GPU. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 4 +++- serve/openai_server.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index 20b2f1f..3a7a33a 100644 --- a/serve/README.md +++ b/serve/README.md @@ -32,7 +32,9 @@ per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearrangement stalls all ranks for about 0.3 s, and with `--eplb-after-prefill 1` it fires after nearly every step under multi-turn traffic; enable it with -`--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. +`--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. With EPLB off the server also drops +the 32 redundant expert slots (`DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=0`): without rearrangement they are fixed copies +of experts 0-31, and removing them speeds decode up by 4-5%. ## Prompt tokenization diff --git a/serve/openai_server.py b/serve/openai_server.py index 772390d..ca6c557 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -123,6 +123,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks). os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval)) + if args.eplb_interval == 0: + # The 32 redundant expert slots only pay off with EPLB rearrangement; without it they are fixed copies of + # experts 0-31 and only add expert weight to read per rank (decode +4-5% without them on B200 and GB300). + os.environ.setdefault("DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS", "0") setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False) if args.eplb_interval > 0: os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval) From 8b1ed3f26aefbd956178f5a3577931b34874d8fe Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Sun, 25 Oct 2026 21:37:45 +0000 Subject: [PATCH 6/7] serve: document the --prefill-any-interval TTFT / decode-speed trade-off Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 5 +++++ serve/openai_server.py | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/serve/README.md b/serve/README.md index 3a7a33a..44ce939 100644 --- a/serve/README.md +++ b/serve/README.md @@ -36,6 +36,11 @@ EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearran the 32 redundant expert slots (`DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=0`): without rearrangement they are fixed copies of experts 0-31, and removing them speeds decode up by 4-5%. +`--prefill-any-interval` (default 2) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and +with 2 almost every arriving request triggers one. That keeps TTFT low; the engine default of 100 adds a ~1.4 s +floor. At high concurrency, 30-100 gives faster decode and more throughput at the cost of TTFT: on B200 AgentX c256, +interval 100 gave 255K vs 211K tok/s/GPU and interactivity p50 119 vs 55, with TTFT p50 2.5 vs 0.7 s. + ## Prompt tokenization Agentic clients resend the whole conversation every turn, and `tokenizer.encode` needs ~1.7 ms per 1K tokens on a diff --git a/serve/openai_server.py b/serve/openai_server.py index ca6c557..b11f94e 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -67,7 +67,8 @@ def parse_args() -> argparse.Namespace: type=int, default=2, help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this " - "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor)", + "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor). A prompt step " + "pauses decode on every rank, so at high concurrency a larger value (30-100) trades TTFT for decode speed", ) p.add_argument( "--eplb-interval", From f6388940003e9c962337bcd1f8d0ba693d977f86 Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Mon, 26 Oct 2026 04:43:27 +0000 Subject: [PATCH 7/7] serve: default --prefill-any-interval to 30 (throughput over TTFT) At AgentX c256 on one GB300 tray, 30 vs 2 gives +9% tok/s/GPU and 2.2x P90 interactivity, for +0.4 s TTFT p50. Use 2 for low TTFT. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01YNgkNrJL5oR1pZqVYzZMHW --- serve/README.md | 19 ++++++++++++++----- serve/openai_server.py | 6 +++--- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/serve/README.md b/serve/README.md index 44ce939..e7b119f 100644 --- a/serve/README.md +++ b/serve/README.md @@ -36,10 +36,18 @@ EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearran the 32 redundant expert slots (`DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=0`): without rearrangement they are fixed copies of experts 0-31, and removing them speeds decode up by 4-5%. -`--prefill-any-interval` (default 2) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and -with 2 almost every arriving request triggers one. That keeps TTFT low; the engine default of 100 adds a ~1.4 s -floor. At high concurrency, 30-100 gives faster decode and more throughput at the cost of TTFT: on B200 AgentX c256, -interval 100 gave 255K vs 211K tok/s/GPU and interactivity p50 119 vs 55, with TTFT p50 2.5 vs 0.7 s. +`--prefill-any-interval` (default 30) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and +a rank with a waiting prompt forces one at most this many decode steps later. The default favours throughput. With 2, +almost every arriving request triggers a prompt step, which keeps TTFT low but slows decode at high concurrency. The +engine default of 100 adds a ~1.4 s TTFT floor. AgentX at c256 on one GB300 tray (20-minute runs): + +| interval | tok/s/GPU | TTFT p50 / p90 | ITL p50 / p90 | P90 interactivity | +|---|---|---|---|---| +| 2 | 258K | 1.26 / 4.53 s | 14.3 / 31.5 ms | 32 | +| 30 | 283K | 1.62 / 6.51 s | 10.0 / 14.1 ms | 71 | +| 100 | 276K | 3.00 / 14.95 s | 8.1 / 13.5 ms | 74 | + +B200 shows the same trend: at c256, interval 100 gave 255K vs 211K tok/s/GPU for interval 2. ## Prompt tokenization @@ -71,4 +79,5 @@ stored when they come back as part of the next turn's prompt). ## AgentX The server speaks what InferenceX's AgentX replay (aiperf, `--endpoint /v1/chat/completions --streaming ---use-server-token-count`) needs; `prefill-any-interval` matters there, since a single lane keeps most ranks idle. +--use-server-token-count`) needs. At low concurrency, use `--prefill-any-interval 2`: a single lane keeps most +ranks idle, and the default adds a few hundred ms of TTFT. diff --git a/serve/openai_server.py b/serve/openai_server.py index b11f94e..4bfd939 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -65,10 +65,10 @@ def parse_args() -> argparse.Namespace: p.add_argument( "--prefill-any-interval", type=int, - default=2, + default=30, help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this " - "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor). A prompt step " - "pauses decode on every rank, so at high concurrency a larger value (30-100) trades TTFT for decode speed", + "many decode steps later even if other ranks are idle. A prompt step pauses decode on every rank, so a larger " + "value gives faster decode and more throughput at the cost of TTFT. 30 favours throughput; use 2 for low TTFT", ) p.add_argument( "--eplb-interval",