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

deepinfer/hicache/pd_transfer.py

@@ -0,0 +1,125 @@
+"""Prefill/decode split: move a prefilled sequence's KV blocks and SWA state between engines.
+
+The prefill engine (EngineType.PROMPT) exports, after the last prefill chunk, the sequence's full KV blocks
+(all tokens of the prompt) and the complete SWA/compression state in the "to CE" layout
+(GPUCacheInterface.select_state_slice_to_store(..., TO_CE=True)), both into pinned host buffers through the
+page-copy kernel. The decode engine (EngineType.COMPLETION) receives those bytes into pinned buffers, allocates
+blocks and a state slot for a new sequence (prompt + first sampled token), fills them with the kernel, and admits
+the sequence into its `prompted` queue (Scheduler.add_prompted_sequence). Transport between the two is left to
+the caller (plain TCP in serve/pd_server.py).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from deepinfer import envs
+from deepinfer.hicache.kernels import copy_pages_module
+from deepinfer.worker.cache_engine import GPUCacheInterface
+
+if TYPE_CHECKING:
+    from deepinfer._C.engine import Sequence
+    from deepinfer.engine.llm_engine import LLMEngine
+
+
+def load_state_from_ce_transfer(state_id: int, last_position: int, state_tensor: torch.Tensor, state_cache: torch.Tensor) -> None:
+    """Inverse of select_state_slice_to_store(..., TO_CE=True): write a full inference-layout state (stored SWA,
+    dropped/consumer SWA and compression state) into state_cache[state_id] with the ring ending at last_position."""
+    from deepinfer.model_executor.layers.tilelang_ops.state_cache import load_state_to_cache
+
+    cls = GPUCacheInterface
+    cache_views, tensor_views = cls._carve_state_views(state_id, state_cache, state_tensor, TO_CE=True)
+    for tensor_view, cache_view in zip(tensor_views, cache_views, strict=True):
+        if tensor_view.numel():
+            load_state_to_cache(tensor_view, cache_view, last_position)
+    dropped_tensor_offset = cls.swa_total_bytes_inference_pad
+    dropped_cache_offset = cls.swa_total_bytes_kvstore_pad
+    compress_bytes = cls.dropped_bytes_inference - cls.swa_drop_total_bytes_inference_pad
+    if compress_bytes:
+        state_cache[state_id, dropped_cache_offset : dropped_cache_offset + compress_bytes].copy_(
+            state_tensor[dropped_tensor_offset : dropped_tensor_offset + compress_bytes]
+        )
+
+
+class PDTransfer:
+    """Per-engine helper owning the GPU pools' pointers and small pinned staging for transfers."""
+
+    def __init__(self, engine: LLMEngine) -> None:
+        cache_engine = engine.driver_worker.cache_engine
+        self.engine = engine
+        self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks
+        self.gpu_states: torch.Tensor = cache_engine.gpu_states
+        self.device = self.gpu_blocks.device
+        self.block_size: int = engine.cache_config.block_size
+        self.page_bytes: int = self.gpu_blocks.size(1)
+        self.state_bytes: int = GPUCacheInterface.total_bytes_inference  # full "to CE" state
+        assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0
+        self.kernel = copy_pages_module()
+
+    def _ids(self, ids: list[int]) -> torch.Tensor:
+        return torch.tensor(ids, dtype=torch.int32).pin_memory().to(self.device, non_blocking=True)
+
+    # ---------------------------------------------------------------- prefill side
+
+    def export(self, seq: Sequence, prompt_len: int | None = None) -> tuple[torch.Tensor, torch.Tensor, torch.cuda.Event]:
+        """Copy seq's KV blocks for its prompt and the full state to pinned host buffers.
+
+        prompt_len: tokens with KV (default: seq.get_len() - 1, i.e. after the first output token was appended;
+        pass seq.get_len() when called before outputs are processed). KV exists for the prompt only, so a prompt
+        ending exactly on a block boundary must not count an extra block.
+        Returns (pages[num_blocks, page_bytes] pinned, state[state_bytes] pinned, event); the buffers are valid once
+        the event has completed. Enqueued on the current stream, so it is ordered after the prefill's kernels.
+        """
+        if prompt_len is None:
+            prompt_len = seq.get_len() - 1
+        num_blocks = (prompt_len + self.block_size - 1) // self.block_size
+        block_table = list(seq.block_table)[:num_blocks]
+        pages = torch.empty((num_blocks, self.page_bytes), dtype=torch.uint8, pin_memory=True)
+        self.kernel.copy_pages(
+            self.gpu_blocks.data_ptr(), pages.data_ptr(), self._ids(block_table), self._ids(list(range(num_blocks))), self.page_bytes
+        )
+        state = torch.empty(self.state_bytes, dtype=torch.uint8, pin_memory=True)
+        staging = torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device)
+        GPUCacheInterface.select_state_slice_to_store(seq.state_id, prompt_len, self.gpu_states, state_tensor=staging, TO_CE=True)
+        self.kernel.copy_pages(staging.data_ptr(), state.data_ptr(), self._ids([0]), self._ids([0]), self.state_bytes)
+        ev = torch.cuda.Event()
+        ev.record()
+        # keep the staging alive until the copy ran (same-stream ordering makes later reuse safe)
+        ev.staging = staging  # type: ignore[attr-defined]
+        return pages, state, ev
+
+    # ---------------------------------------------------------------- decode side
+
+    def admit(self, seq: Sequence, pages: torch.Tensor, state: torch.Tensor, prompt_len: int) -> torch.cuda.Event:
+        """Fill a freshly created sequence (prompt + first token, created with add_request(enqueue=False)) from
+        pinned buffers and put it in the scheduler's prompted queue. Caller must have called try_allocate_seq.
+
+        prompt_len: tokens with KV in `pages` (= seq.get_len() - 1); the state ring ends there.
+        """
+        num_blocks = pages.size(0)
+        block_table = list(seq.block_table)
+        assert len(block_table) >= num_blocks, (len(block_table), num_blocks)
+        self.kernel.copy_pages(
+            pages.data_ptr(),
+            self.gpu_blocks.data_ptr(),
+            self._ids(list(range(num_blocks))),
+            self._ids(block_table[:num_blocks]),
+            self.page_bytes,
+        )
+        staging = torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device)
+        self.kernel.copy_pages(state.data_ptr(), staging.data_ptr(), self._ids([0]), self._ids([0]), self.state_bytes)
+        load_state_from_ce_transfer(seq.state_id, prompt_len, staging, self.gpu_states)
+        seq.num_cached_tokens = prompt_len
+        seq.cached_hit_tokens = prompt_len
+        # A finished prefill leaves `nextn` draft tokens on the sequence (pad tokens when the prompt step produced
+        # none, as a prompt-only engine does); the decode sampler requires exactly that many.
+        nextn = int(envs.DEEPINFER_USE_NEXTN_DECODE)
+        if nextn > 0:
+            seq.last_next_n_token = self.engine.get_nextn_pad_token_ids().tolist()[:nextn]
+        self.engine.scheduler.add_prompted_sequence(seq)
+        ev = torch.cuda.Event()
+        ev.record()
+        ev.staging = staging  # type: ignore[attr-defined]
+        return ev

serve/README.md

@@ -43,3 +43,21 @@ stored when they come back as part of the next turn's prompt).
 
 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`).
+
+```bash
+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`.

serve/check_pd.py

@@ -0,0 +1,81 @@
+"""Compare a PD (prefill/decode split) deployment against a single-engine server on the same prompts.
+
+    python serve/check_pd.py --ref http://localhost:8888 --pd http://tray10:8888 --turns 3 --prefix-tokens 3000
+
+Greedy, multi-turn with a long shared prefix (so the PE's prefix cache gets exercised); prints per-turn TTFT and
+whether the outputs are identical. Exits non-zero on any mismatch.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+
+import requests
+
+p = argparse.ArgumentParser()
+p.add_argument("--ref", default="http://localhost:8888")
+p.add_argument("--pd", required=True)
+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=3000)
+p.add_argument("--max-tokens", type=int, default=64)
+args = p.parse_args()
+
+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.",
+    "Quote the first five words of the document.",
+    "What is the last word of the document?",
+    "How many distinct word stems are used? Answer briefly.",
+]
+
+
+def ask(url: str, messages: list[dict]) -> tuple[str, float, float, int]:
+    body = {
+        "model": args.model,
+        "messages": messages,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        "temperature": 0,
+        "max_tokens": args.max_tokens,
+        "seed": 1,
+    }
+    t0 = time.perf_counter()
+    ttft, text, n = None, [], 0
+    with requests.post(f"{url}/v1/chat/completions", 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"):
+                n = obj["usage"]["completion_tokens"]
+            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, n
+
+
+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)]})
+    ref, ttft_r, tot_r, n_r = ask(args.ref, messages)
+    pd, ttft_p, tot_p, n_p = ask(args.pd, messages)
+    same = ref == pd
+    mismatches += not same
+    print(
+        f"turn {turn + 1}: ref TTFT {ttft_r * 1e3:6.0f} ms total {tot_r * 1e3:6.0f} ms ({n_r} tok) | PD TTFT {ttft_p * 1e3:6.0f} ms total {tot_p * 1e3:6.0f} ms ({n_p} tok) | "
+        f"outputs {'identical' if same else 'DIFFER'}"
+    )
+    if not same:
+        print("   ref:", repr(ref[:160]))
+        print("   pd :", repr(pd[:160]))
+    messages.append({"role": "assistant", "content": ref})
+raise SystemExit(1 if mismatches else 0)

serve/nccl_xtray_test.py

@@ -0,0 +1,30 @@
+"""Cross-tray NCCL check: run one rank on each of two trays and report the transport NCCL picks + bandwidth.
+
+# on tray A (rank 0) and tray B (rank 1), same MASTER_ADDR:
+MASTER_ADDR=<trayA> MASTER_PORT=29999 RANK=0 WORLD_SIZE=2 CUDA_VISIBLE_DEVICES=0 NCCL_DEBUG=INFO python serve/nccl_xtray_test.py
+"""
+
+import os
+import time
+
+import torch
+import torch.distributed as dist
+
+dist.init_process_group("nccl")
+rank = dist.get_rank()
+torch.cuda.set_device(0)
+x = torch.ones(256 << 20, dtype=torch.uint8, device="cuda")  # 256 MiB
+for _ in range(3):
+    dist.all_reduce(x)
+torch.cuda.synchronize()
+t0 = time.perf_counter()
+n = 10
+for _ in range(n):
+    dist.all_reduce(x)
+torch.cuda.synchronize()
+dt = (time.perf_counter() - t0) / n
+print(
+    f"rank {rank} on {os.uname().nodename}: all_reduce 256 MiB avg {dt * 1e3:.2f} ms -> {x.numel() / dt / 2**30:.1f} GiB/s algbw",
+    flush=True,
+)
+dist.destroy_process_group()

serve/pd_run.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+# Start/stop the prefill (PE) and decode (CE) trays of serve/pd_server.py over ssh.
+#
+#   serve/pd_run.sh start <pe-host> <ce-host> [extra pd_server args...]
+#   serve/pd_run.sh stop  <pe-host> <ce-host>
+#   serve/pd_run.sh log   <host>              # tail the tray's log
+#
+# Ports: PE HTTP 8889, CE HTTP (user-facing) 8888, CE KV receivers 8900+rank. Logs: ~/pd_<role>.log on each tray.
+set -euo pipefail
+CMD="${1:?start|stop|log}"; shift
+MODEL="${DEEPINFER_HF_MODEL_DIR:-$HOME/models/DeepSeek-V4.1-Flash-1017}"
+SERVED=deepseek-ai/DeepSeek-V4.1-Flash-1017
+ENV='cd ~/DeepInfer && source env.sh && export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}'
+
+launch() {  # host role args...
+    local host=$1 role=$2; shift 2
+    # ssh -n + nohup with every fd redirected, otherwise the remote session stays attached to the server.
+    ssh -n "$host" "$ENV && rm -f ~/pd_$role.pgid && nohup setsid python -u serve/pd_server.py --role $role --model $MODEL --served-model-name $SERVED $* > ~/pd_$role.log 2>&1 < /dev/null & echo \$! > ~/pd_$role.pgid; echo started $role on \$(hostname) pgid \$(cat ~/pd_$role.pgid)"
+}
+wait_ready() {  # host port
+    local i
+    for i in $(seq 200); do
+        if curl -fs "http://$1:$2/health" >/dev/null 2>&1; then echo "$1:$2 ready"; return 0; fi
+        sleep 3
+    done
+    echo "$1:$2 did not become ready" >&2; return 1
+}
+
+case "$CMD" in
+    start)
+        PE=${1:?pe-host}; CE=${2:?ce-host}; shift 2
+        launch "$PE" pe --port 8889 "$@"
+        launch "$CE" ce --port 8888 --kv-port 8900 --kv-host "$CE" --pe-url "http://$PE:8889" "$@"
+        wait_ready "$PE" 8889; wait_ready "$CE" 8888
+        ;;
+    stop)
+        for spec in "${1:?pe-host}:pe" "${2:?ce-host}:ce"; do
+            h=${spec%%:*}; r=${spec##*:}
+            # Kill by pid list, excluding the remote shell itself: a pkill -f pattern would match this very
+            # command line (it names pd_server.py) and kill the shell before anything else runs.
+            # Frontend + its spawned engine workers, SIGKILL (uvicorn's graceful shutdown otherwise keeps the
+            # children, and the GPUs, alive for minutes); then wait until the GPUs are actually free.
+            ssh -n "$h" "for p in \$(pgrep -f '[p]d_server.py --role $r|multiprocessing.spawn|multiprocessing.resource_tracker'); do [ \"\$p\" != \"\$\$\" ] && kill -9 \"\$p\" 2>/dev/null; done; rm -f ~/pd_$r.pgid; for i in \$(seq 60); do [ \"\$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | sort -n | tail -1)\" -lt 1000 ] && break; sleep 3; done; echo stopped $r on \$(hostname), gpu mem \$(nvidia-smi --query-gpu=memory.used --format=csv,noheader | tr '\\n' ' ')" || true
+        done
+        ;;
+    log)
+        ssh "${1:?host}" 'tail -n 40 ~/pd_*.log'
+        ;;
+esac

serve/pd_server.py

@@ -0,0 +1,682 @@
+"""Prefill/decode-disaggregated serving: a prompt-only engine tray (PE) and a decode-only engine tray (CE).
+
+    # decode tray (user-facing OpenAI API on --port, KV receivers on --kv-port + rank):
+    python serve/pd_server.py --role ce --model ... --dp 4 --port 8888 --kv-port 8900 --pe-url http://<pe-host>:8889
+    # prefill tray:
+    python serve/pd_server.py --role pe --model ... --dp 4 --port 8889
+
+Flow per request: CE frontend renders the prompt, picks a CE rank, POSTs /prefill to the PE with the KV target
+(ce-host, kv-port + rank). The PE prefills (host prefix cache on), exports the sequence's KV blocks + full SWA
+state (deepinfer.hicache.pd_transfer) and the first sampled token, streams the first token back in the HTTP
+response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence straight into decode
+(Scheduler.add_prompted_sequence); its tokens stream to the client through the CE frontend.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import collections
+import json
+import multiprocessing as mp
+import os
+import queue
+import socket
+import struct
+import sys
+import threading
+import time
+import uuid
+from collections.abc import Callable
+from pathlib import Path
+from types import ModuleType
+from typing import TYPE_CHECKING
+
+from fastapi import FastAPI, Request
+
+if TYPE_CHECKING:
+    from deepinfer._C.engine import RequestOutput, Sequence
+    from deepinfer.engine.llm_engine import LLMEngine
+    from deepinfer.hicache.pd_transfer import PDTransfer
+    from deepinfer.sampling_params import SamplingParams
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+sys.path.insert(0, str(HERE.parent / "examples"))
+
+import openai_server as oas  # noqa: E402  (the single-engine server: prompt rendering, detokenizer, frontend plumbing)
+
+HDR = struct.Struct("<16sqqqqq")  # rid(16 bytes hex) | first_token | prompt_len | num_pages | page_bytes | state_bytes
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--role", required=True, choices=["pe", "ce"])
+    p.add_argument("--model", required=True)
+    p.add_argument("--served-model-name", default=None)
+    p.add_argument("--dp", type=int, default=4)
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8888)
+    p.add_argument("--kv-port", type=int, default=8900, help="ce: TCP receiver base port (+ rank)")
+    p.add_argument("--kv-host", default=None, help="ce: address the PE should push KV to (default: this host's name)")
+    p.add_argument("--pe-url", default=None, help="ce: prefill tray URL")
+    p.add_argument("--max-model-len", type=int, default=1 << 20)
+    p.add_argument("--max-num-seqs", type=int, default=128)
+    p.add_argument("--max-num-batched-tokens", type=int, default=32768)
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    p.add_argument("--yoco-keep-tokens", type=int, default=140)
+    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="pe: host prefix cache per rank (0 disables)")
+    p.add_argument("--hicache-state-slots", type=int, default=4096)
+    p.add_argument("--routing", default="affinity", choices=["affinity", "least-loaded"])
+    p.add_argument("--prefill-any-interval", type=int, default=2)
+    return p.parse_args()
+
+
+def _recv_exact(sock: socket.socket, buf: memoryview) -> None:
+    view = buf
+    while len(view):
+        n = sock.recv_into(view, min(len(view), 1 << 24))
+        if n == 0:
+            raise ConnectionError("peer closed during transfer")
+        view = view[n:]
+
+
+# ============================================================================ engine processes
+
+
+def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None:
+    import ctypes
+    import signal
+
+    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL)  # PR_SET_PDEATHSIG
+
+    import numpy as np
+    from _config import create_llm, setup_common_env
+
+    from deepinfer.engine.engine_type import EngineType
+
+    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"
+    os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512")  # identical state layout on PE and CE
+    os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval))
+    setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True)
+
+    engine_type = EngineType.PROMPT if args.role == "pe" else EngineType.COMPLETION
+    llm = create_llm(
+        args.model,
+        dp_rank,
+        ep_size=args.dp,
+        env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)},
+        data_parallel_size=args.dp,
+        max_num_seqs=args.max_num_seqs,
+        max_num_batched_tokens=args.max_num_batched_tokens,
+        max_model_len=args.max_model_len,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
+        engine_type=engine_type,
+        stop_preempted_request=True,
+    )
+
+    from deepinfer._C.engine import RequestOutput  # pyright: ignore[reportMissingModuleSource]
+    from deepinfer.hicache.pd_transfer import PDTransfer
+    from deepinfer.sampling_params import SamplingParams
+
+    engine = llm.llm_engine
+    transfer = PDTransfer(engine)
+    mm = lambda: dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None  # noqa: E731
+
+    def enqueue(rid: str, seq: Sequence, num_cached: int) -> None:
+        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)
+
+    if args.role == "pe":
+        _pe_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm, enqueue)
+    else:
+        _ce_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm)
+
+
+def _pe_loop(
+    args: argparse.Namespace,
+    dp_rank: int,
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    engine: LLMEngine,
+    transfer: PDTransfer,
+    SamplingParams: type[SamplingParams],
+    RequestOutput: type[RequestOutput],
+    np: ModuleType,
+    mm: Callable[[], dict | None],
+    enqueue: Callable[[str, Sequence, int], None],
+) -> None:
+    cache = None
+    if args.hicache_size_gb > 0:
+        from deepinfer.hicache import HostPrefixCache
+
+        cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots)
+
+    targets: dict[str, tuple[str, int]] = {}  # rid -> (host, port) of the CE rank
+    awaiting_token: dict[str, dict] = {}  # rid -> transfer waiting for the step's sampled first token
+    send_q: queue.Queue = queue.Queue()
+
+    def on_prefill_step(scheduled: list[Sequence], finished: list[Sequence]) -> None:
+        # Engine hook at the start of post-processing for a prompt step: the prefill's KV is complete and the
+        # sequence still owns its blocks (with max_tokens=1 it finishes, and is freed, once outputs are processed).
+        if cache is not None:
+            cache.on_prefill_step(scheduled, finished)
+        for seq in scheduled:
+            rid = seq.request_id
+            if seq.chunked_rest_len != 0 or rid not in targets:
+                continue
+            prompt_len = seq.get_len()  # no output appended yet
+            need = (prompt_len + transfer.block_size - 1) // transfer.block_size
+            if len(seq.block_table) < need:
+                out_q.put(("error", rid, f"prefilled sequence has {len(seq.block_table)} blocks, needs {need}"))
+                targets.pop(rid, None)
+                continue
+            pages, state, ev = transfer.export(seq, prompt_len)
+            item = {
+                "pages": pages,
+                "state": state,
+                "event": ev,
+                "prompt_len": prompt_len,
+                "target": targets.pop(rid),
+                "first_token": None,
+                "ready": threading.Event(),
+            }
+            awaiting_token[rid] = item
+            send_q.put((rid, item))
+
+    engine.pe_do_save_kvcache = on_prefill_step
+
+    def sender():
+        while True:
+            rid, item = send_q.get()
+            host, port = item["target"]
+            try:
+                item["ready"].wait(timeout=600)
+                if item["first_token"] is None:
+                    raise RuntimeError("no first token produced for the prefilled sequence")
+                item["event"].synchronize()
+                pages, state = item["pages"], item["state"]
+                with socket.create_connection((host, port), timeout=60) as s:
+                    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+                    s.sendall(
+                        HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel())
+                    )
+                    s.sendall(memoryview(pages.numpy()).cast("B"))
+                    s.sendall(memoryview(state.numpy()).cast("B"))
+                    ack = s.recv(2)
+                if ack != b"ok":
+                    out_q.put(("error", rid, f"CE rejected transfer: {ack!r}"))
+            except Exception as e:  # noqa: BLE001
+                out_q.put(("error", rid, f"KV transfer to {host}:{port} failed: {e!r}"))
+
+    threading.Thread(target=sender, daemon=True).start()
+    out_q.put(("ready", dp_rank, None))
+
+    pending = collections.deque()
+    inflight = False
+    last_stats = time.monotonic()
+    while True:
+        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)
+                    sp = SamplingParams(**payload["sampling"])
+                    sp.max_tokens = 1  # prefill + first token only; the engine then finishes and frees the sequence
+                    seq = engine.add_request(
+                        user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=prompt_ids, multi_modal_data=mm(), enqueue=False
+                    )
+                    if seq is None:
+                        out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
+                    else:
+                        targets[rid] = tuple(payload["kv_target"])
+                        pending.append((rid, seq, prompt_ids))
+                except Exception as e:  # noqa: BLE001
+                    out_q.put(("error", rid, repr(e)))
+            elif kind == "abort":
+                targets.pop(rid, None)
+                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)
+        while pending:
+            rid, seq, prompt_ids = pending[0]
+            hit = cache.try_resume(seq, rid, prompt_ids) if cache is not None 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()
+                out_q.put(("stats", dp_rank, cache.stats_dict()))
+
+        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):
+            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
+            ]
+            for rid, token_id, finished, _, _ in batch:
+                item = awaiting_token.pop(rid, None)
+                if item is not None:
+                    item["first_token"] = int(token_id)
+                    item["ready"].set()
+                if finished and cache is not None:
+                    cache.forget(rid)
+            out_q.put(("tokens", None, batch))
+        # A prompt-only engine never reclaims finished sequences by itself (that happens in the decode
+        # scheduler): drop them from `running` and free their blocks/state. The export and cache copies were
+        # enqueued on the stream before any kernel that could reuse those blocks.
+        for seq in engine.scheduler.free_finished_seqs():
+            engine.scheduler.free_seq(seq)
+
+
+def _ce_loop(
+    args: argparse.Namespace,
+    dp_rank: int,
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    engine: LLMEngine,
+    transfer: PDTransfer,
+    SamplingParams: type[SamplingParams],
+    RequestOutput: type[RequestOutput],
+    np: ModuleType,
+    mm: Callable[[], dict | None],
+) -> None:
+    import torch
+
+    recv_q: queue.Queue = queue.Queue()  # (rid, first_token, prompt_len, pages, state)
+
+    def receiver():
+        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        srv.bind((args.host, args.kv_port + dp_rank))
+        srv.listen(64)
+        while True:
+            conn, _ = srv.accept()
+            threading.Thread(target=receive_one, args=(conn,), daemon=True).start()
+
+    def receive_one(conn: socket.socket):
+        try:
+            with conn:
+                hdr = bytearray(HDR.size)
+                _recv_exact(conn, memoryview(hdr))
+                rid_b, first_token, prompt_len, num_pages, page_bytes, state_bytes = HDR.unpack(hdr)
+                rid = rid_b.decode()
+                assert page_bytes == transfer.page_bytes and state_bytes == transfer.state_bytes, "layout mismatch between PE and CE"
+                pages = torch.empty((num_pages, page_bytes), dtype=torch.uint8, pin_memory=True)
+                state = torch.empty(state_bytes, dtype=torch.uint8, pin_memory=True)
+                _recv_exact(conn, memoryview(pages.numpy()).cast("B"))
+                _recv_exact(conn, memoryview(state.numpy()).cast("B"))
+                recv_q.put((rid, first_token, prompt_len, pages, state))
+                conn.sendall(b"ok")
+        except Exception as e:  # noqa: BLE001
+            print(f"[ce rank {dp_rank}] transfer failed: {e!r}", flush=True)
+
+    threading.Thread(target=receiver, daemon=True).start()
+    out_q.put(("ready", dp_rank, None))
+
+    expected: dict[str, dict] = {}  # rid -> {"prompt_token_ids", "sampling"} from the frontend
+    arrived: dict[str, tuple] = {}  # rid -> transfer
+    pending: collections.deque = collections.deque()  # rids with both parts, waiting for GPU allocation
+    keep: dict[str, torch.cuda.Event] = {}
+    inflight = False
+    while True:
+        while True:
+            try:
+                kind, rid, payload = in_q.get_nowait()
+            except queue.Empty:
+                break
+            if kind == "expect":
+                expected[rid] = payload
+                if rid in arrived:
+                    pending.append(rid)
+            elif kind == "abort":
+                expected.pop(rid, None)
+                arrived.pop(rid, None)
+                if rid in pending:
+                    pending.remove(rid)
+                else:
+                    engine.abort_request(rid)
+        while True:
+            try:
+                rid, first_token, prompt_len, pages, state = recv_q.get_nowait()
+            except queue.Empty:
+                break
+            arrived[rid] = (first_token, prompt_len, pages, state)
+            if rid in expected:
+                pending.append(rid)
+        while pending:
+            rid = pending[0]
+            spec, (first_token, prompt_len, pages, state) = expected[rid], arrived[rid]
+            tokens = np.asarray(list(spec["prompt_token_ids"]) + [first_token], dtype=np.int32)
+            assert len(tokens) - 1 == prompt_len, (len(tokens), prompt_len)
+            sp = SamplingParams(**spec["sampling"])
+            sp.max_tokens = max(1, sp.max_tokens - 1)  # the first token was produced on the PE
+            seq = engine.add_request(
+                user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=tokens, multi_modal_data=mm(), enqueue=False
+            )
+            if seq is None:
+                out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
+                pending.popleft()
+                expected.pop(rid, None)
+                arrived.pop(rid, None)
+                continue
+            try:
+                allocated = engine.scheduler.try_allocate_seq(seq)
+            except RuntimeError as e:
+                out_q.put(("error", rid, f"cannot allocate: {e!r}"))
+                allocated = None
+            if allocated is False and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
+                break  # retry once something finishes
+            pending.popleft()
+            expected.pop(rid, None)
+            arrived.pop(rid, None)
+            if not allocated:
+                if allocated is False:
+                    out_q.put(("error", rid, "cannot allocate KV blocks for the transferred sequence"))
+                continue
+            # Keep the pinned receive buffers alive until the H2D page-copy kernels have consumed them: the kernel
+            # reads them by raw pointer, so the caching host allocator would otherwise hand the memory to the next
+            # transfer (or unmap it) while the copy is still queued -> corrupted KV / cudaErrorLaunchFailure.
+            keep[rid] = (transfer.admit(seq, pages, state, prompt_len), pages, state)
+            if engine.multi_modal_processor is not None:
+                engine.multi_modal_processor.add(rid, seq.multi_modal_data, prompt_len)
+        for rid in [r for r, (ev, _, _) in keep.items() if ev.query()]:
+            del keep[rid]
+
+        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):
+            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
+            ]
+            out_q.put(("tokens", None, batch))
+
+
+# ============================================================================ frontends
+
+
+class PDFrontend(oas.Frontend):
+    def __init__(self, args: argparse.Namespace) -> None:
+        super().__init__(args)
+        ctx = mp.get_context("spawn")
+        self.out_q = ctx.Queue()
+        self.in_qs = [ctx.Queue() for _ in range(args.dp)]
+        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)]
+
+
+def build_pe_app(fe: PDFrontend) -> FastAPI:
+    """Prefill tray: POST /prefill {rid, prompt_token_ids, sampling, kv_target:[host,port]} -> {first_token}."""
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse
+
+    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, "role": "pe"}, status_code=200 if ok else 503)
+
+    @app.get("/deepinfer/metrics")
+    async def metrics():
+        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
+        for r, st in sorted(fe.cache_stats.items()):
+            lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()]
+        return "\n".join(lines) + "\n"
+
+    @app.post("/prefill")
+    async def prefill(request: Request):
+        if fe.ready != fe.args.dp:
+            return JSONResponse({"error": "engine loading"}, status_code=503)
+        body = await request.json()
+        rid = body["rid"]
+        rank = min(range(fe.args.dp), key=lambda r: fe.load[r])
+        key = bytes.fromhex(body["affinity"]) if body.get("affinity") else None
+        if fe.args.routing == "affinity" and key is not None:
+            rank = int.from_bytes(key[:4], "little") % fe.args.dp
+        fe.load[rank] += 1
+        fe.owner[rid] = rank
+        q: asyncio.Queue = asyncio.Queue()
+        fe.streams[rid] = q
+        fe.stats["requests"] += 1
+        fe.in_qs[rank].put(
+            ("add", rid, {"prompt_token_ids": body["prompt_token_ids"], "sampling": body["sampling"], "kv_target": body["kv_target"]})
+        )
+        try:
+            item = await asyncio.wait_for(q.get(), timeout=1800)
+        finally:
+            fe.release(rid, aborted=False)
+        if item[0] == "error":
+            return JSONResponse({"error": item[1]}, status_code=500)
+        return {"first_token": item[1], "prompt_len": item[4]}
+
+    return app
+
+
+def build_ce_app(fe: PDFrontend, pe_url: str, kv_host: str) -> FastAPI:
+    """Decode tray: the user-facing OpenAI API."""
+    import httpx
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse, StreamingResponse
+
+    app = FastAPI()
+    client = httpx.AsyncClient(timeout=1800)
+
+    @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, "role": "ce"}, 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"}]}
+
+    @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)]
+        return "\n".join(lines) + "\n"
+
+    @app.post("/v1/chat/completions")
+    async def chat(request: Request):
+        if fe.ready != fe.args.dp:
+            return JSONResponse({"error": {"message": "engine loading"}}, status_code=503)
+        body = await request.json()
+        prompt_ids = await asyncio.to_thread(fe.render_prompt, body)
+        max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096
+        if len(prompt_ids) + 1 > fe.args.max_model_len:
+            return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400)
+        max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids)))
+        sampling = dict(
+            temperature=float(body.get("temperature", 1.0)),
+            top_p=float(body.get("top_p", 1.0)),
+            ignore_eos=bool(body.get("ignore_eos", False)),
+            max_tokens=max_tokens,
+        )
+        if body.get("seed") is not None:
+            sampling["seq_seed"] = int(body["seed"])
+        fe.stats["prompt_tokens"] += len(prompt_ids)
+        include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
+        key = fe.affinity_key(body)
+
+        # Register with the CE rank first so the transfer can be admitted as soon as it lands.
+        rid = uuid.uuid4().hex[:16]
+        rank = (
+            int.from_bytes(key[:4], "little") % fe.args.dp
+            if (fe.args.routing == "affinity" and key)
+            else min(range(fe.args.dp), key=lambda r: fe.load[r])
+        )
+        fe.load[rank] += 1
+        fe.owner[rid] = rank
+        q: asyncio.Queue = asyncio.Queue()
+        fe.streams[rid] = q
+        fe.stats["requests"] += 1
+        fe.in_qs[rank].put(("expect", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling}))
+
+        created = int(time.time())
+        cid = f"chatcmpl-{rid}"
+        detok = oas.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 prefill_first():
+            r = await client.post(
+                f"{pe_url}/prefill",
+                json={
+                    "rid": rid,
+                    "prompt_token_ids": prompt_ids,
+                    "sampling": sampling,
+                    "kv_target": [kv_host, fe.args.kv_port + rank],
+                    "affinity": key.hex() if key else None,
+                },
+            )
+            if r.status_code != 200:
+                raise RuntimeError(f"PE: {r.text[:200]}")
+            return int(r.json()["first_token"])
+
+        async def tokens():
+            first = await prefill_first()
+            text = detok.add(first)
+            fe.stats["completion_tokens"] += 1
+            if max_tokens == 1 or first == fe.tokenizer.eos_token_id and not sampling["ignore_eos"]:
+                yield text, ("length" if max_tokens == 1 else "stop")
+                return
+            yield text, None
+            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 = PDFrontend(args)
+    fe.start()
+    import uvicorn
+
+    if args.role == "pe":
+        app = build_pe_app(fe)
+    else:
+        assert args.pe_url, "--pe-url is required for the decode tray"
+        app = build_ce_app(fe, args.pe_url.rstrip("/"), args.kv_host or socket.gethostname())
+    uvicorn.run(app, host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900)
+
+
+if __name__ == "__main__":
+    main()