Jiejing 8e488bbd55 serve: prefill/decode split prototype over two hosts (pd_server.py)
A prompt-only engine tray (PE) and a decode-only engine tray (CE)
behind the OpenAI-compatible frontend from serve/openai_server.py:

- serve/pd_server.py: --role pe / --role ce. The CE frontend renders
  the prompt, picks a CE rank (affinity or least-loaded routing) and
  POSTs /prefill to the PE. The PE prefills with the host prefix cache
  on, exports the prompt's KV blocks and the full SWA state in the
  pe_do_save_kvcache hook, streams the first sampled token back in the
  HTTP response and pushes the KV bytes to the CE rank over TCP. The CE
  rank admits the sequence into decode via Scheduler.add_prompted_sequence
  and reclaims PE sequences through Scheduler.free_finished_seqs.
- deepinfer/hicache/pd_transfer.py: export/admit of KV pages and state
  through the hicache page-copy kernels. The CE keeps the pinned receive
  buffers alive until the admit copy kernels finish (the kernels read
  them by raw pointer; the caching host allocator otherwise reused the
  memory while the copy was queued -> cudaErrorLaunchFailure).
- serve/pd_run.sh start/stop/logs over ssh; serve/check_pd.py compares
  outputs against a single-engine server token by token;
  serve/nccl_xtray_test.py reports the cross-host NCCL transport.

Measured on 4 GB300 (PE) + 4 GB300 (CE) against one 4-GPU tray with
the AgentX replay, outputs identical: ITL p90 stays within 8.6-15 ms
from conc 8 to 256 (single tray 12-98 ms), P90 interactivity 66-117
tok/s/user where the single tray falls to 10; per total chip PD is
behind up to conc 64 and ahead at 256 (doc/benchmark_gb300.md).

Depends on the host prefix cache (#144), the OpenAI server (#147) and
the scheduler bindings (add_prompted_sequence / free_finished_seqs),
which are included as the first three commits.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr
2026-10-21 22:26:29 +00:00
..
check_pd.py serve: prefill/decode split prototype over two hosts (pd_server.py) 2026-10-21 22:26:29 +00:00
check_prefix_cache.py Add a minimal OpenAI-compatible chat server over the offline engine 2026-10-21 22:26:29 +00:00
nccl_xtray_test.py serve: prefill/decode split prototype over two hosts (pd_server.py) 2026-10-21 22:26:29 +00:00
openai_server.py serve: type annotations and B007 for ruff 2026-10-21 22:26:29 +00:00
pd_run.sh serve: prefill/decode split prototype over two hosts (pd_server.py) 2026-10-21 22:26:29 +00:00
pd_server.py serve: prefill/decode split prototype over two hosts (pd_server.py) 2026-10-21 22:26:29 +00:00
README.md serve: prefill/decode split prototype over two hosts (pd_server.py) 2026-10-21 22:26:29 +00:00
server.sh Add a minimal OpenAI-compatible chat server over the offline engine 2026-10-21 22:26:29 +00:00

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

export DEEPINFER_ACCELERATOR=nvidia DEEPINFER_HF_MODEL_DIR=/path/to/DeepSeek-V4.1-Flash-1017
serve/server.sh start --dp 4 --port 8888                       # cache on: 32 GiB pages + 4096 states per rank
serve/server.sh start --dp 4 --port 8888 --hicache-size-gb 0   # cache off
python serve/check_prefix_cache.py --port 8888 --turns 3
curl -s localhost:8888/deepinfer/metrics | grep hicache         # per-rank hit/store counters
serve/server.sh stop

Knobs: --hicache-size-gb, --hicache-state-slots, --routing affinity|least-loaded (affinity keeps a conversation's turns on one rank, keyed by the first user message; required for hits since the cache is per rank). The cache requires decoder bounded replay (--yoco-keep-tokens > 0, the default 140).

How a hit works

A new prompt's page hashes are matched from the start; the match is capped so that at least keep_tokens tokens remain to be recomputed (the decoder SWA state is rebuilt from that suffix by bounded replay), then snapped down to the nearest 128-aligned position that has a stored SWA state. The sequence is allocated up front, pages and state are copied in on the engine stream, num_cached_tokens is set, and the sequence is enqueued; the scheduler prefills only the suffix. Generated tokens are not stored at decode end (they are stored when they come back as part of the next turn's prompt).

AgentX

The server speaks what InferenceX's AgentX replay (aiperf, --endpoint /v1/chat/completions --streaming --use-server-token-count) needs; prefill-any-interval matters there, since a single lane keeps most ranks idle.

Prefill/decode split (pd_server.py)

serve/pd_server.py runs the same OpenAI-compatible frontend over two engines on two hosts: a prompt-only engine (PE, --role pe, host prefix cache on) and a decode-only engine (CE, --role ce, user-facing). Per request the CE frontend renders the prompt, picks a CE rank and POSTs /prefill to the PE; the PE prefills, exports the sequence's KV blocks and SWA state (deepinfer/hicache/pd_transfer.py) plus the first sampled token, pushes the bytes to the CE rank over TCP and the CE admits the sequence straight into decode (Scheduler.add_prompted_sequence).

serve/pd_run.sh start <pe-host> <ce-host>          # PE HTTP 8889, CE HTTP 8888, CE KV receivers 8900+rank
python serve/check_pd.py --ref http://<single-engine-host>:8888 --pd http://<ce-host>:8888   # token-identical?
serve/pd_run.sh stop  <pe-host> <ce-host>

serve/nccl_xtray_test.py reports the transport and bandwidth NCCL picks between the two hosts. Results on 4+4 GB300 vs a single 4-GPU tray: doc/benchmark_gb300.md.