coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron: PR #147 OpenAI server (+ /v1/completions, Neuron worker setup), offline bench, MoE fixes (#26, #27, #29)

Browse Source 
- serve/: the minimal OpenAI-compatible server of deepseek-ai/DeepInfer#147,
  with a Neuron worker setup (basic mode env, per-rank NeuronCore, fixed KV
  pool via --num-kv-blocks, no MTP lookahead) and a /v1/completions endpoint
  (prompt text or token ids, no chat template): what vllm bench serve
  --backend openai drives, as vllm-neuron#71's run_bench.py does.
- examples/neuron/bench_neuron.py: the same TTFT/TPOT matrix offline.
- static MoE: one packed fp32 all-gather instead of three (each 64-rank
  collective is ~2 ms of fixed latency a layer); a capacity-grouped expert
  path for prefill chunks (DEEPINFER_NEURON_MOE_CAPACITY, dense rerun on
  overflow; not yet validated on device).
- Sub-region timings at split() boundaries under DEEPINFER_NEURON_REGION_TIMING.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-23 15:07:36 +08:00
parent 7d897aedb2
commit d8f7a2e898
13 changed files with 1558 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
deepinfer/environ.py
View File

@ -122,6 +122,9 @@ class envs:
# (issue #26, the vllm-neuron#71 layout). 0 = eager prefill. Needs max_num_batched_tokens <= bucket and
# DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket (the SWA / compressor rings must hold a whole chunk).
DEEPINFER_NEURON_PREFILL_BUCKET = EnvVar.int(default=0)
# Neuron only: routed rows per local expert in a prefill chunk's MoE (0 = dense over all gathered rows). An expert
# with more routed rows reruns the layer's MoE densely.
DEEPINFER_NEURON_MOE_CAPACITY = EnvVar.int(default=1024)
# Neuron only: run the sampler's process_sampling as one compiled graph instead of ~1900 eager ops a step.
DEEPINFER_NEURON_COMPILE_SAMPLER = EnvVar.bool(default=True)
# Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).

8
deepinfer/model_executor/models/decoder_layer.py
View File

@ -140,6 +140,14 @@ class DeepseekDecoderLayer(HCLayerMixin):
cuda_graph_timer.record("mlp", enabled=(self.layer_idx == 1 and not input_metadata.is_prompt), scale=mlp_timer_scale),
):
mlp = neuron_regions.entry(self.mlp, "forward") if use_regions else self.mlp
capacity = envs.DEEPINFER_NEURON_MOE_CAPACITY
if use_regions and capacity and hidden_states_bf16.shape[0] == envs.DEEPINFER_NEURON_PREFILL_BUCKET:
# A prefill chunk: the experts over their routed rows only, the dense path if one overflowed.
out, overflow = neuron_regions.entry(self.mlp, "forward_capacity")(
hidden_states, hidden_states_bf16, input_metadata
)
hidden_states = mlp(hidden_states, hidden_states_bf16, input_metadata) if bool(overflow) else out
else:
hidden_states = mlp(
hidden_states,
hidden_states_bf16,

20
deepinfer/model_executor/models/v4_modules/mlp.py
View File

@ -578,6 +578,26 @@ class DeepseekMoE(nn.Module):
self.shared_forward(hidden_states, out_acc=out)
return out
def forward_capacity(
self,
hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
hidden_states_bf16: torch.Tensor,
input_metadata: InputMetadata,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Neuron prefill chunk (issue #26): the static MoE with each local expert over at most
DEEPINFER_NEURON_MOE_CAPACITY routed rows. Returns ``(output, overflow)``; on overflow the caller reruns
``forward``."""
topk_idx, topk_weights = self.gate_forward(hidden_states_bf16, input_metadata)
return static_moe_forward(
self,
hidden_states,
topk_idx,
topk_weights,
hidden_states_bf16,
getattr(input_metadata, "valid_token_mask", None),
capacity=envs.DEEPINFER_NEURON_MOE_CAPACITY,
)
def forward(
self,
hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

19
deepinfer/neuron_ops/regions.py
View File

@ -83,18 +83,33 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
return out
torch.neuron.synchronize()
started = time.perf_counter()
_MARK[0], _MARK[1], _MARK[2] = started, 0, label
out = compiled(mod, *args, **kwargs)
torch.neuron.synchronize()
_TIMES[label].append(time.perf_counter() - started)
done = time.perf_counter()
if _MARK[1]:
_TIMES[f"{label}.split{_MARK[1]}"].append(done - _MARK[0])
_MARK[0] = 0.0
_TIMES[label].append(done - started)
return out
_ENTRIES[key] = fn
return functools.partial(fn, module)
_MARK: list = [0.0, 0, ""]
@torch._dynamo.disable
def _boundary() -> None:
"""A graph boundary and nothing else (vllm-neuron#71's ``_graph_break``)."""
"""A graph boundary and nothing else (vllm-neuron#71's ``_graph_break``). Under DEEPINFER_NEURON_REGION_TIMING
also records the time since the previous boundary (or the region's start) as ``split<i>``."""
if envs.DEEPINFER_NEURON_REGION_TIMING:
torch.neuron.synchronize()
now = time.perf_counter()
if _MARK[0]:
_TIMES[f"{_MARK[2]}.split{_MARK[1]}"].append(now - _MARK[0])
_MARK[0], _MARK[1] = now, _MARK[1] + 1
def split() -> None:

65
deepinfer/neuron_ops/static_moe.py
View File

@ -23,7 +23,7 @@ import torch.distributed._functional_collectives as funcol
from tile_kernels.quant import swiglu_forward_and_per_token_cast
from deepinfer.model_executor.layers.linear import get_fp8_cast_params, get_fp8_casting_granularity
from deepinfer.neuron_ops import nki_gemm
from deepinfer.neuron_ops import nki_gemm, regions
if TYPE_CHECKING:
from deepinfer.model_executor.models.v4_modules.mlp import DeepseekMoE
@ -82,14 +82,16 @@ def static_moe_forward(
topk_weights: torch.Tensor,
hidden_states_bf16: torch.Tensor | None = None,
valid: torch.Tensor | None = None,
) -> torch.Tensor:
capacity: int = 0,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""``valid`` ([tokens] bool, the decode metadata's valid_token_mask): padding rows of a graph batch select no
expert, so they cost nothing past the dense compute and never reach a real row's output."""
if valid is not None:
topk_idx = torch.where(valid.view(-1, 1), topk_idx, torch.full_like(topk_idx, -1))
topk_weights = torch.where(valid.view(-1, 1), topk_weights, torch.zeros_like(topk_weights))
if hidden_states_bf16 is not None and nki_gemm.has_kernel(moe.gate_up_proj_weight):
return _nki_moe_forward(moe, hidden_states_bf16, topk_idx, topk_weights)
return _nki_moe_forward(moe, hidden_states_bf16, topk_idx, topk_weights, capacity)
assert not capacity, "the grouped expert path needs the NKI kernels"
group, size, rank = moe.ep_group, moe.ep_size, moe.ep_rank
num_local = moe.num_local_physical_experts
x_parts = list(hidden_states) if moe.use_fp8_act else [hidden_states]
@ -142,9 +144,42 @@ def _expert_gemm(a, b, d, e: int, fp8: bool, recipe_a, recipe_b) -> None: # noq
d.copy_((a.float() @ b[e].float().transpose(0, 1)).to(d.dtype))
def _grouped_experts(moe: DeepseekMoE, x: torch.Tensor, local_idx: torch.Tensor, w: torch.Tensor, capacity: int):
"""The local experts over at most ``capacity`` routed rows each instead of all ``rows`` (issue #26: a prefill
chunk gathers 64 ranks x 128 rows, and each row selects only 6 of the 384 experts). Returns the fp32
``[rows, hidden]`` sum and a bool scalar that is True when some expert had more routed rows than ``capacity``
(the result is then incomplete; the caller reruns the dense path)."""
rows = x.shape[0]
x_pad = torch.cat([x, x.new_zeros((1, x.shape[1]))], dim=0) # row `rows` is the empty slot
acc = torch.zeros((rows + 1, moe.hidden_size), dtype=torch.float32, device=x.device)
row_ids = torch.arange(rows, device=x.device, dtype=torch.int32)
overflow = torch.zeros((), dtype=torch.bool, device=x.device)
for e in range(moe.num_local_physical_experts):
hit = local_idx == e
w_e = torch.where(hit, w, torch.zeros_like(w)).sum(-1)
row_hit = hit.any(dim=-1)
# Slot of each routed row (in row order); rows past the capacity and unrouted rows go to the spare slot.
slot = torch.cumsum(row_hit.to(torch.int32), dim=0) - 1
dest = torch.where(row_hit & (slot < capacity), slot, torch.full_like(slot, capacity))
table = torch.full((capacity + 1,), rows, dtype=torch.int32, device=x.device).scatter(0, dest, row_ids)[:capacity]
overflow = overflow | (row_hit.sum() > capacity)
xs = x_pad.index_select(0, table)
ws = torch.cat([w_e, w_e.new_zeros(1)]).index_select(0, table)
gate_up = nki_gemm.linear_qdq(xs, moe.gate_up_proj_weight, group=e)
act = swiglu_forward_and_per_token_cast(
gate_up, fmt="bf16", num_per_channels=None, topk_weights=ws, clamp_value=moe.swiglu_clamp_value
)
acc = acc.index_add(0, table, nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float())
return acc[:rows], overflow
def _nki_moe_forward(
moe: DeepseekMoE, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor
) -> torch.Tensor:
moe: DeepseekMoE,
x: torch.Tensor,
topk_idx: torch.Tensor,
topk_weights: torch.Tensor,
capacity: int = 0,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""The same dense static MoE on bf16 activations through the NKI block-quant kernels (vllm-neuron#71 style).
The kernels quantize activations to fp8 themselves (1x32, power-of-two scales), exactly as the fp8 path does, so
@ -163,9 +198,21 @@ def _nki_moe_forward(
x = torch.nn.functional.pad(x, (0, 0, 0, pad))
idx = torch.nn.functional.pad(idx, (0, 0, 0, pad), value=-1)
w = torch.nn.functional.pad(w, (0, 0, 0, pad))
x, idx, w = (_all_gather(t, group, size) for t in (x, idx, w))
acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device)
# One collective instead of three (each 64-rank collective costs ~2 ms of fixed latency a layer, issue #27):
# fp32 holds the bf16 activations and the expert ids (< 2**24) exactly, and needs no dtype bitcast.
topk = idx.shape[1]
packed = torch.cat([x.float(), idx.float(), w], dim=1)
packed = _all_gather(packed, group, size)
x = packed[:, : x.shape[1]].to(torch.bfloat16)
idx = packed[:, x.shape[1] : x.shape[1] + topk].to(torch.int32)
w = packed[:, x.shape[1] + topk :]
regions.split()
local_idx = idx - moe.neuron_rank_tensor() * num_local
overflow = None
if capacity and capacity < x.shape[0]:
acc, overflow = _grouped_experts(moe, x, local_idx, w, capacity)
else:
acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device)
for e in range(num_local):
w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1)
gate_up = nki_gemm.linear_qdq(x, moe.gate_up_proj_weight, group=e)
@ -173,11 +220,15 @@ def _nki_moe_forward(
gate_up, fmt="bf16", num_per_channels=None, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value
)
acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float()
regions.split()
out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
regions.split()
if moe.has_shared:
shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
shared = swiglu_forward_and_per_token_cast(
shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value
)
out = out + nki_gemm.linear_qdq(shared, moe.shared_down_proj.weight).float()
if capacity:
return out.to(torch.bfloat16), overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device)
return out.to(torch.bfloat16)

150
examples/neuron/bench_neuron.py Normal file
View File

@ -0,0 +1,150 @@
"""TTFT / TPOT benchmark of the full DeepSeek-V4.1-Flash on one trn2.48xlarge, the vllm-neuron#71 method offline.
vllm-neuron#71 (``contrib/models/deepseek_v41/test/performance/run_bench.py``) serves the model and drives it with
``vllm bench serve --dataset-name random``: random token ids fix the input length exactly and share no prefix,
output 128 tokens, and per cell it reports the median TTFT, the median TPOT (decode only:
``(latency - ttft) / (output_tokens - 1)``) and the decode tokens per second. DeepInfer has no server, so this drives
the engine directly with the same requests: random prompt token ids of the cell's input length, ``ignore_eos`` and
``max_tokens`` 128, timed per request from submission.
Topology: one engine process per DP rank (EP=64, DP-attention), and every rank runs its own request of the cell at
the same time: batch 1 per rank, 64 requests in flight (the engine steps all DP ranks together). #71's bs1 cells are
one request for the whole machine; per request the two loads cost the same here, since the static MoE computes
every rank's row either way.
export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0 DEEPINFER_NEURON_PREFILL_BUCKET=128 \\
DEEPINFER_MAX_BACKTRACKING_TOKENS=128
python examples/neuron/bench_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64
"""
from __future__ import annotations
import argparse
import json
import os
import statistics
import sys
import time
from pathlib import Path
EXAMPLES = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(EXAMPLES))
from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices # noqa: E402
from _launch import launch_dp # noqa: E402
parser = argparse.ArgumentParser(description="V4.1 TTFT/TPOT benchmark on Neuron (vllm-neuron#71 method)")
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
parser.add_argument("--ep", default=64, type=int)
parser.add_argument("--input_lens", default="512,1024,2048,8192", help="comma-separated prompt lengths (tokens)")
parser.add_argument("--output_len", default=128, type=int)
parser.add_argument("--repeats", default=3, type=int, help="requests per cell (medians are reported)")
parser.add_argument("--warmup", default=1, type=int, help="untimed requests before the first cell")
parser.add_argument("--num_kv_blocks", default=400, type=int)
parser.add_argument("--max_model_len", default=16384, type=int)
parser.add_argument("--max_num_batched_tokens", default=128, type=int)
parser.add_argument("--max_num_seqs", default=1, type=int)
parser.add_argument("--out", default="", help="write the results as JSON here (rank 0)")
args = parser.parse_args()
dp_size = args.ep
def _fixed_kv_pool(num_blocks: int) -> None:
from deepinfer.worker.base import WorkerBase
def profile_num_available_blocks(self): # noqa: ANN001, ANN202
self.model_runner.set_block_size(self.cache_config.block_size)
self.cache_config.num_gpu_blocks = num_blocks
return num_blocks
WorkerBase.profile_num_available_blocks = profile_num_available_blocks
def _random_prompt(length: int, seed: int, vocab: int) -> list[int]:
import numpy as np
rng = np.random.default_rng(seed)
return rng.integers(1000, vocab - 1000, size=length, dtype=np.int64).astype(np.int32).tolist()
def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int] | None: # noqa: ANN001
"""Run the engine until it drains. With a prompt, returns (ttft_s, latency_s, output_tokens) for it."""
engine = llm.llm_engine
start = time.perf_counter()
if prompt is not None:
llm._add_request(None, params, prompt)
first = last = None
tokens = 0
while True:
unfinished = engine.has_unfinished_requests()
outputs = engine.step() if unfinished else engine.flush_inflight_reqs()
now = time.perf_counter()
for out in outputs:
tokens += 1
first = first if first is not None else now
last = now
if not unfinished:
break
if prompt is None:
return None
return first - start, last - start, tokens
def main(dp_size: int, dp_rank: int) -> None:
os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
from deep_cluster.hardware import get_actual_num_gpus
setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus())
_fixed_kv_pool(args.num_kv_blocks)
llm = create_llm(
args.model_name,
dp_rank,
ep_size=args.ep,
cp_size=1,
env={"DEEPINFER_LOGGING_INTERVAL_STEP": "100"},
data_parallel_size=dp_size,
max_num_batched_tokens=args.max_num_batched_tokens,
max_model_len=args.max_model_len,
max_num_seqs=args.max_num_seqs,
gpu_memory_utilization=0.85,
lookahead_tokens=0,
)
from deepinfer.sampling_params import SamplingParams
params = SamplingParams(temperature=0.0, max_tokens=args.output_len, ignore_eos=True)
vocab = llm.llm_engine.tokenizer.vocab_size if hasattr(llm.llm_engine.tokenizer, "vocab_size") else 128000
busy = True
lens = [int(x) for x in args.input_lens.split(",")]
for w in range(args.warmup):
_run_one(llm, _random_prompt(lens[0], 10_000 + w, vocab) if busy else None, params)
results = []
for length in lens:
cell = []
for r in range(args.repeats):
got = _run_one(llm, _random_prompt(length, 1000 * length + r, vocab) if busy else None, params)
if got is not None:
ttft, latency, n = got
tpot = (latency - ttft) / max(n - 1, 1)
cell.append((ttft, tpot, n))
if dp_rank == 0:
print(f"BENCH input={length} rep={r} ttft_ms={ttft * 1e3:.1f} tpot_ms={tpot * 1e3:.1f} tokens={n}", flush=True)
if dp_rank == 0 and cell:
ttft = statistics.median(c[0] for c in cell)
tpot = statistics.median(c[1] for c in cell)
results.append({"input": length, "output": args.output_len, "ttft_p50_ms": ttft * 1e3, "tpot_p50_ms": tpot * 1e3,
"decode_tok_s_per_request": 1.0 / tpot})
if dp_rank == 0:
print("BENCH_RESULTS " + json.dumps(results), flush=True)
print(f"| input | output | TTFT p50 (ms) | TPOT p50 (ms) | decode tok/s per request | (batch 1 per DP rank)")
for r in results:
print(f"| {r['input']} | {r['output']} | {r['ttft_p50_ms']:.0f} | {r['tpot_p50_ms']:.1f} | {r['decode_tok_s_per_request']:.2f} |")
if args.out:
Path(args.out).write_text(json.dumps(results, indent=2))
llm.close()
if __name__ == "__main__":
launch_dp(main, dp_size)

67
serve/README.md Normal file
View File

@ -0,0 +1,67 @@
# 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.
## 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; `prefill-any-interval` matters there, since a single lane keeps most ranks idle.

87
serve/check_prefix_cache.py Normal file
View File

@ -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)

695
serve/openai_server.py Normal file
View File

@ -0,0 +1,695 @@
"""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=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)",
)
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(
"--num-kv-blocks",
type=int,
default=0,
help="fixed KV block pool per rank instead of the memory-profiling prefill (Neuron: the profile prefill does not "
"fit next to the graphs; examples/neuron/offline_inference_neuron.py)",
)
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))
neuron = os.environ.get("DEEPINFER_ACCELERATOR") == "neuron"
if neuron:
# Neuron basic mode (examples/neuron/offline_inference_neuron.py): no MTP drafting, no forced top-k sampling,
# no expert override; this rank's NeuronCore is picked before deepinfer.PLATFORM initializes.
os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
setup_common_env(model_name=args.model, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
from deep_cluster.hardware import get_actual_num_gpus
from _config import setup_cuda_visible_devices
setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus())
else:
setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False)
if args.num_kv_blocks > 0:
from deepinfer.worker.base import WorkerBase
def profile_num_available_blocks(self, _n: int = args.num_kv_blocks) -> int: # noqa: ANN001
self.model_runner.set_block_size(self.cache_config.block_size)
self.cache_config.num_gpu_blocks = _n
return _n
WorkerBase.profile_num_available_blocks = profile_num_available_blocks
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,
**({"cp_size": 1} if neuron else {}),
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=0 if neuron else min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
)
import numpy as np
from deepinfer._C.engine import RequestOutput # pyright: ignore[reportMissingModuleSource]
from deepinfer.sampling_params import SamplingParams
engine = llm.llm_engine
cache = None
if args.hicache_size_gb > 0:
try:
from deepinfer.hicache import HostPrefixCache
except ImportError: # tree without the host prefix cache: serve without it
print(f"[server] deepinfer.hicache not available, prefix cache disabled (rank {dp_rank})", flush=True)
else:
cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots)
engine.pe_do_save_kvcache = cache.on_prefill_step # write-through at the end of every prefill
out_q.put(("ready", dp_rank, None))
def mm_data():
return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None
def enqueue(rid: str, seq: Sequence, num_cached_tokens: int) -> None:
# Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict
# (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length.
engine.scheduler.add_sequence(seq)
if engine.multi_modal_processor is not None:
engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached_tokens)
pending = collections.deque() # (rid, seq, prompt_ids, use_cache): created, waiting for a cache decision
inflight = False
last_stats = time.monotonic()
last_logged_lookups = 0
while True:
# Requests must be enqueued before the DP-wide has_unfinished_requests all-reduce,
# so every rank takes the same step / flush / idle branch in the same iteration.
while True:
try:
kind, rid, payload = in_q.get_nowait()
except queue.Empty:
break
if kind == "add":
try:
prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32)
seq = engine.add_request(
user_id="0",
request_id=rid,
sampling_params=SamplingParams(**payload["sampling"]),
prompt_token_ids=prompt_ids,
multi_modal_data=mm_data(),
enqueue=cache is None,
)
if cache is not None:
if seq is None:
out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
else:
pending.append((rid, seq, prompt_ids, payload.get("use_cache", True)))
except Exception as e: # noqa: BLE001 - surface to the client, keep the rank alive
out_q.put(("error", rid, repr(e)))
elif kind == "abort":
if any(p[0] == rid for p in pending):
pending = collections.deque(p for p in pending if p[0] != rid)
else:
engine.abort_request(rid)
if cache is not None:
cache.forget(rid)
# Cache decisions, in arrival order. A hit pre-allocates the sequence's blocks and fills them from the
# host; if the GPU cannot allocate yet, wait (while something is running that will free memory).
while pending:
rid, seq, prompt_ids, use_cache = pending[0]
hit = cache.try_resume(seq, rid, prompt_ids) if use_cache else 0
if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
break
pending.popleft()
enqueue(rid, seq, max(hit, 0))
if cache is not None:
cache.poll()
if time.monotonic() - last_stats > 5:
last_stats = time.monotonic()
st = cache.stats_dict()
out_q.put(("stats", dp_rank, st))
if st["lookups"] != last_logged_lookups:
last_logged_lookups = st["lookups"]
print(
f"[hicache rank {dp_rank}] lookups {st['lookups']} hits {st['hits']} "
f"token hit rate {st['hit_rate_tokens']:.1%} deferred {st['deferred']} stores {st['stores']} "
f"pages {st['pages_used']}/{st['pages_total']} states {st['states_used']}/{st['states_total']} "
f"evicted {st['evicted_pages']}",
flush=True,
)
if engine.has_unfinished_requests():
outputs = engine.step()
inflight = True
elif inflight:
outputs = engine.flush_inflight_reqs()
inflight = False
else:
time.sleep(0.001)
continue
if len(outputs):
# Pointers are recycled on the next step, so read the fields now.
batch = [
(
RequestOutput.get_request_id(p),
RequestOutput.get_token_id(p),
RequestOutput.get_finished(p),
RequestOutput.get_finish_reason(p),
RequestOutput.get_prompt_len(p),
)
for p in outputs
]
if cache is not None:
for rid, _, finished, _, _ in batch:
if finished:
cache.forget(rid)
out_q.put(("tokens", None, batch))
# --------------------------------------------------------------------------- HTTP side
class IncrementalDetokenizer:
"""vLLM-style prefix/read offsets so each delta decodes only a short tail."""
def __init__(self, tokenizer: Any) -> None:
self.tok = tokenizer
self.ids: list[int] = []
self.num_tokens = 0
self.prefix_offset = 0
self.read_offset = 0
def add(self, token_id: int) -> str:
self.num_tokens += 1
if token_id == self.tok.eos_token_id:
return ""
self.ids.append(token_id)
prefix = self.tok.decode(self.ids[self.prefix_offset : self.read_offset], skip_special_tokens=False)
full = self.tok.decode(self.ids[self.prefix_offset :], skip_special_tokens=False)
if len(full) > len(prefix) and not full.endswith("<EFBFBD>"):
self.prefix_offset, self.read_offset = self.read_offset, len(self.ids)
return full[len(prefix) :]
return ""
class Frontend:
def __init__(self, args: argparse.Namespace) -> None:
from transformers import AutoTokenizer
sys.path.insert(0, str(Path(args.model) / "encoding"))
from encoding import encode_messages # the checkpoint's own prompt-format reference
self.args = args
self.model_name = args.served_model_name or args.model
self.encode_messages = encode_messages
self.tokenizer = AutoTokenizer.from_pretrained(args.model)
probe = [
{"role": "system", "content": "sys", "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}]},
{"role": "user", "content": "hi 你好 123"},
{"role": "assistant", "content": "ok", "tool_calls": [{"type": "function", "function": {"name": "f", "arguments": "{}"}}]},
{"role": "tool", "content": " out\n"},
]
try:
probes = [encode_messages(probe, thinking_mode=args.thinking_mode)]
except Exception: # noqa: BLE001 - the probe only adds coverage
probes = []
# getattr: frontends built on this one with their own argument parser (serve/pd_server.py) get the defaults
backend, verify_rate = getattr(args, "tokenizer_backend", "auto"), getattr(args, "tokenizer_verify_rate", 0.005)
self.prompt_tokenizer = PromptTokenizer(self.tokenizer, backend, verify_rate, extra_probes=probes)
print(f"[server] {self.prompt_tokenizer.describe()}", flush=True)
ctx = mp.get_context("spawn")
self.out_q: mp.Queue = ctx.Queue()
self.in_qs: list[mp.Queue] = [ctx.Queue() for _ in range(args.dp)]
# Not daemonic: the engine forks its own helper processes (as examples/_launch.py allows).
self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)]
self.load = [0] * args.dp
self.owner: dict[str, int] = {}
self.streams: dict[str, asyncio.Queue] = {}
self.ready = 0
self.stats = dict(requests=0, finished=0, aborted=0, errors=0, prompt_tokens=0, completion_tokens=0, encode_fallbacks=0)
self.cache_stats: dict[int, dict] = {}
self.loop: asyncio.AbstractEventLoop | None = None
def start(self) -> None:
for p in self.procs:
p.start()
def dispatch_forever(self) -> None:
"""Thread: fan engine outputs out to the per-request asyncio queues."""
while True:
try:
kind, rid, payload = self.out_q.get(timeout=1.0)
except queue.Empty:
dead = [r for r, p in enumerate(self.procs) if not p.is_alive()]
if dead:
print(f"[server] engine rank(s) {dead} died; failing open requests and exiting", flush=True)
for rid in list(self.streams):
self.loop.call_soon_threadsafe(self._deliver, rid, ("error", f"engine rank died: {dead}"))
time.sleep(2)
os._exit(1)
continue
if kind == "ready":
self.ready += 1
print(f"[server] DP rank {rid} ready ({self.ready}/{self.args.dp})", flush=True)
elif kind == "error":
self.loop.call_soon_threadsafe(self._deliver, rid, ("error", payload))
elif kind == "stats":
self.cache_stats[rid] = payload
else:
self.loop.call_soon_threadsafe(self._deliver_batch, payload)
def _deliver(self, rid: str, item: tuple) -> None:
q = self.streams.get(rid)
if q is not None:
q.put_nowait(item)
def _deliver_batch(self, batch: list[tuple]) -> None:
for rid, token_id, finished, reason, prompt_len in batch:
q = self.streams.get(rid)
if q is not None:
q.put_nowait(("token", token_id, finished, reason, prompt_len))
def render_prompt(self, body: dict) -> list[int]:
messages = [dict(m) for m in body["messages"]]
if body.get("tools"):
if not messages or messages[0].get("role") != "system":
messages.insert(0, {"role": "system", "content": ""})
messages[0]["tools"] = body["tools"]
try:
prompt = self.encode_messages(messages, thinking_mode=self.args.thinking_mode)
except Exception: # noqa: BLE001 - trace messages the reference encoder rejects
# Keep the request (and its prompt length) instead of failing the benchmark.
self.stats["encode_fallbacks"] += 1
parts = ["<｜begin▁of▁sentence｜>"]
for m in messages:
content = m.get("content") or ""
if isinstance(content, list):
content = "".join(b.get("text", "") for b in content if isinstance(b, dict))
if m.get("tool_calls"):
content += json.dumps(m["tool_calls"], ensure_ascii=False)
tag = {"user": "<｜User｜>", "assistant": "<｜Assistant｜>", "tool": "<｜User｜>"}.get(m.get("role"), "")
parts.append(tag + content)
if body.get("tools"):
parts.insert(1, json.dumps(body["tools"], ensure_ascii=False))
prompt = "".join(parts) + "<｜Assistant｜></think>"
return self.prompt_tokenizer.encode(prompt)
@staticmethod
def affinity_key(body: dict) -> bytes | None:
"""Stable per-conversation key: the first user message (AgentX's cache-bust marker lives there and is
constant across a trajectory's turns, and forked subagents inherit it)."""
for m in body.get("messages", []):
if m.get("role") == "user":
content = m.get("content") or ""
if isinstance(content, list):
content = "".join(b.get("text", "") for b in content if isinstance(b, dict))
return hashlib.sha1(content[:4096].encode("utf-8", "ignore")).digest()
return None
def submit(self, prompt_ids: list[int], sampling: dict, key: bytes | None = None, use_cache: bool = True) -> tuple[str, asyncio.Queue]:
rid = uuid.uuid4().hex
if self.args.routing == "affinity" and key is not None:
rank = int.from_bytes(key[:4], "little") % self.args.dp
else:
rank = min(range(self.args.dp), key=lambda r: self.load[r])
self.load[rank] += 1
self.owner[rid] = rank
q: asyncio.Queue = asyncio.Queue()
self.streams[rid] = q
self.stats["requests"] += 1
self.in_qs[rank].put(("add", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling, "use_cache": use_cache}))
return rid, q
def release(self, rid: str, aborted: bool) -> None:
rank = self.owner.pop(rid, None)
self.streams.pop(rid, None)
if rank is None:
return
self.load[rank] -= 1
if aborted:
self.stats["aborted"] += 1
self.in_qs[rank].put(("abort", rid, None))
else:
self.stats["finished"] += 1
def build_app(fe: Frontend) -> FastAPI:
from fastapi import FastAPI
from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
app = FastAPI()
@app.on_event("startup")
async def _startup() -> None:
fe.loop = asyncio.get_running_loop()
threading.Thread(target=fe.dispatch_forever, daemon=True).start()
@app.get("/health")
async def health():
ok = fe.ready == fe.args.dp
return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready}, status_code=200 if ok else 503)
@app.get("/v1/models")
async def models():
return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]}
# Not at /metrics: aiperf scrapes that path automatically and InferenceX's aggregator
# rejects metric families it has no backend adapter for.
@app.get("/deepinfer/metrics")
async def metrics():
lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)]
for r, st in sorted(fe.cache_stats.items()):
lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()]
pt = fe.prompt_tokenizer
lines += [f'deepinfer_tokenizer:{k}{{backend="{pt.name}"}} {v}' for k, v in pt.metrics().items()]
return PlainTextResponse("\n".join(lines) + "\n")
@app.post("/v1/completions")
async def completions(request: Request):
"""Plain completions (prompt text or token ids, no chat template): what ``vllm bench serve --backend openai``
drives, so a benchmark's input length is the one it names (vllm-neuron#71's method)."""
if fe.ready != fe.args.dp:
return JSONResponse({"error": {"message": "engine loading"}}, status_code=503)
body = await request.json()
prompt = body.get("prompt")
if isinstance(prompt, list) and len(prompt) == 1 and not isinstance(prompt[0], int):
prompt = prompt[0]
if isinstance(prompt, str):
prompt_ids = await asyncio.to_thread(lambda: fe.tokenizer.encode(prompt, add_special_tokens=False))
else:
prompt_ids = [int(t) for t in prompt]
max_tokens = int(body.get("max_tokens") or 16)
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(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,
)
fe.stats["prompt_tokens"] += len(prompt_ids)
include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
rid, q = fe.submit(prompt_ids, sampling, None, use_cache=bool(body.get("prefix_cache", True)))
created = int(time.time())
cid = f"cmpl-{rid}"
detok = IncrementalDetokenizer(fe.tokenizer)
def chunk(text: str | None, finish: str | None = None, usage: dict | None = None) -> str:
obj = {
"id": cid,
"object": "text_completion",
"created": created,
"model": fe.model_name,
"choices": [] if usage else [{"index": 0, "text": text or "", "finish_reason": finish}],
}
if usage:
obj["usage"] = usage
return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"
async def tokens():
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
def usage() -> dict:
n = detok.num_tokens
return {"prompt_tokens": len(prompt_ids), "completion_tokens": n, "total_tokens": len(prompt_ids) + n}
if body.get("stream"):
async def sse():
done = False
try:
async for text, finish in tokens():
# One chunk a token, text or not: clients time every token (TPOT / ITL).
yield chunk(text, finish)
done = True
if include_usage:
yield chunk(None, 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
parts.append(text)
done = True
finally:
fe.release(rid, aborted=not done)
return {
"id": cid,
"object": "text_completion",
"created": created,
"model": fe.model_name,
"choices": [{"index": 0, "text": "".join(parts), "finish_reason": finish}],
"usage": usage(),
}
@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()

264
serve/prompt_tokenizer.py Normal file
View File

@ -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

41
serve/server.sh Executable file
View File

@ -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

0
tests/serve/__init__.py Normal file
View File

136
tests/serve/test_prompt_tokenizer.py Normal file
View File

@ -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"