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

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

deepinfer/environ.py

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

deepinfer/model_executor/models/decoder_layer.py

@@ -140,11 +140,19 @@ 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
-            hidden_states = mlp(
-                hidden_states,
-                hidden_states_bf16,
-                input_metadata,
-            )
+            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,
+                    input_metadata,
+                )
 
         return hidden_states, residual
 

deepinfer/model_executor/models/v4_modules/mlp.py

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

deepinfer/neuron_ops/regions.py

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

deepinfer/neuron_ops/static_moe.py

@@ -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,21 +198,37 @@ 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
-    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)
-        act = swiglu_forward_and_per_token_cast(
-            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()
+    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)
+            act = swiglu_forward_and_per_token_cast(
+                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)

examples/neuron/bench_neuron.py

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

serve/README.md

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

serve/check_prefix_cache.py

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

serve/openai_server.py

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

serve/prompt_tokenizer.py

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

serve/server.sh

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

tests/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"