whn09 d54fe25993 serve: Mooncake Transfer Engine KV transport for the prefill/decode split (--kv-transport mooncake)
Stacked on #197 (pd_server.py / pd_transfer.py), which moves the prefilled KV blocks and SWA
state PE -> CE as pinned host buffers over TCP. --kv-transport mooncake keeps that flow and
control socket, and replaces the data plane with a Mooncake TransferEngine:

- deepinfer/hicache/pd_mooncake.py: GpuArena (one registered device buffer per rank, first-fit
  regions; cuMemAlloc like the KV pools, since expandable_segments memory has no cudaIpc
  handle) and MooncakeLink (TransferEngine with P2PHANDSHAKE metadata, register_memory,
  transfer_sync_write).
- PDTransfer.export_into / admit accept device buffers: the PE exports a sequence straight
  into an arena region (state included, no extra staging) and the CE admits from its arena.
- pd_server.py: MC_MAGIC header -> CE reserves an arena region and replies {addr, segment}
  (or busy -> retry) -> one contiguous write -> go/ok. Writes are retried, then the region is
  sent over the TCP path, so a failing peer cannot strand a request. TCP is unchanged and is
  also the fallback when the PE arena is full.
- --mooncake-protocol nvlink_intra (default, one host) or rdma/efa across hosts.
- serve/README.md: usage; containers must share IPC and PID namespaces (--pid=host).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 15:09:01 +08: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: default --prefill-any-interval to 30 (throughput over TTFT) 2026-10-26 04:43:27 +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: Mooncake Transfer Engine KV transport for the prefill/decode split (--kv-transport mooncake) 2026-10-26 15:09:01 +08:00
prompt_tokenizer.py serve: exact fast prompt tokenization (fastokens or added-token segment cache), checked against tokenizer.encode 2026-10-23 03:50:06 +00:00
README.md serve: Mooncake Transfer Engine KV transport for the prefill/decode split (--kv-transport mooncake) 2026-10-26 15:09:01 +08: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).

EPLB expert rearrangement is off by default (--eplb-interval 0). Every rearrangement stalls all ranks for about 0.3 s, and with --eplb-after-prefill 1 it fires after nearly every step under multi-turn traffic; enable it with --eplb-interval 30 only for decode-heavy workloads whose expert load is skewed. With EPLB off the server also drops the 32 redundant expert slots (DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=0): without rearrangement they are fixed copies of experts 0-31, and removing them speeds decode up by 4-5%.

--prefill-any-interval (default 30) trades TTFT for decode speed. A prompt step pauses decode on every DP rank, and a rank with a waiting prompt forces one at most this many decode steps later. The default favours throughput. With 2, almost every arriving request triggers a prompt step, which keeps TTFT low but slows decode at high concurrency. The engine default of 100 adds a ~1.4 s TTFT floor. AgentX at c256 on one GB300 tray (20-minute runs):

interval tok/s/GPU TTFT p50 / p90 ITL p50 / p90 P90 interactivity
2 258K 1.26 / 4.53 s 14.3 / 31.5 ms 32
30 283K 1.62 / 6.51 s 10.0 / 14.1 ms 71
100 276K 3.00 / 14.95 s 8.1 / 13.5 ms 74

B200 shows the same trend: at c256, interval 100 gave 255K vs 211K tok/s/GPU for interval 2.

Prompt tokenization

Agentic clients resend the whole conversation every turn, and tokenizer.encode needs ~1.7 ms per 1K tokens on a GB300 Grace core (p50 278 ms per request on replayed AgentX prompts of p50 163K tokens), which dominates TTFT once the prefix cache hits. serve/prompt_tokenizer.py gives the same ids faster; --tokenizer-backend picks one:

  • fastokens (pip install fastokens; x86-64 and aarch64 wheels): Rust reimplementation of the HF pipeline, built from the loaded tokenizer. 3.5 ms p50 on the same prompts.
  • segment (no dependency): splits the prompt at added tokens, where HF tokenizers already encode the pieces independently, and caches each piece's ids (LRU, 512 MiB), so a turn only encodes what it adds. 18 ms p50. Enabled only when the tokenizer config makes the split exact (no single_word tokens, no prepend_scheme="first").
  • hf: tokenizer.encode.

auto (default) takes the first of these that matches tokenizer.encode on startup probes built from the tokenizer's added tokens. --tokenizer-verify-rate (default 0.005) re-encodes a sample of prompts with tokenizer.encode in the background; a mismatch switches to hf. The backend, encode latency and check counters are in /deepinfer/metrics (deepinfer_tokenizer:*). Tests: tests/serve/test_prompt_tokenizer.py.

How a hit works

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

AgentX

The server speaks what InferenceX's AgentX replay (aiperf, --endpoint /v1/chat/completions --streaming --use-server-token-count) needs. At low concurrency, use --prefill-any-interval 2: a single lane keeps most ranks idle, and the default adds a few hundred ms of TTFT.

Prefill/decode split (pd_server.py)

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

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.

KV transport: Mooncake Transfer Engine (--kv-transport mooncake)

--kv-transport mooncake (both roles) moves the KV blocks and SWA state with a Mooncake TransferEngine instead of pinned host buffers over TCP (deepinfer/hicache/pd_mooncake.py).

  • Each rank registers one GPU arena (--pd-arena-gb, default 4).
  • The PE exports a sequence into a region of its arena (PDTransfer.export_into) and writes it with one contiguous transfer_sync_write into a region the CE rank reserved in its own arena.
  • The CE admits from there with the same page-copy kernels.
  • Control (header, reservation, completion) stays on the KV socket. Each write is retried, then falls back to the TCP path, so a failing peer cannot strand a request.

--mooncake-protocol picks the transport: nvlink_intra (default; PE and CE on one host, CUDA IPC + copy engines), or rdma / efa across hosts. Use --mooncake-device for the NICs and --mooncake-host for the address peers use for the handshake.

pip install mooncake-transfer-engine-cuda13   # or mooncake-transfer-engine for CUDA 12
# one 8-GPU host: PE on GPUs 0-3, CE on GPUs 4-7
python serve/pd_server.py --role pe --port 8889 --kv-transport mooncake --mooncake-host 127.0.0.1 ...
python serve/pd_server.py --role ce --port 8888 --kv-port 8900 --kv-host 127.0.0.1 --pe-url http://127.0.0.1:8889 \
    --kv-transport mooncake --mooncake-host 127.0.0.1 ...

Notes for nvlink_intra:

  • Allocation. The arena is allocated with cuMemAlloc (cuda_utils.alloc_tensor, like the KV pools), not torch.empty. The worker runs PyTorch with expandable_segments (deepinfer/worker/base.py), and that VMM memory has no cudaIpc handle.
  • Containers. If PE and CE run in separate containers, share the IPC and PID namespaces (docker run --ipc=host --pid=host). A CUDA IPC handle carries the exporter's PID. With per-container PID namespaces, the PE's and CE's rank-r workers get the same PID, and cudaIpcOpenMemHandle rejects the peer handle as its own ("invalid device context"). Only the same-index rank pairs fail.