diff --git a/deepinfer/hicache/pd_transfer.py b/deepinfer/hicache/pd_transfer.py new file mode 100644 index 0000000..abc69e1 --- /dev/null +++ b/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 diff --git a/serve/README.md b/serve/README.md index 9f204eb..a362759 100644 --- a/serve/README.md +++ b/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 HTTP 8889, CE HTTP 8888, CE KV receivers 8900+rank +python serve/check_pd.py --ref http://:8888 --pd http://:8888 # token-identical? +serve/pd_run.sh stop +``` + +`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`. diff --git a/serve/check_pd.py b/serve/check_pd.py new file mode 100644 index 0000000..2203325 --- /dev/null +++ b/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) diff --git a/serve/nccl_xtray_test.py b/serve/nccl_xtray_test.py new file mode 100644 index 0000000..543bf87 --- /dev/null +++ b/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= 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() diff --git a/serve/pd_run.sh b/serve/pd_run.sh new file mode 100755 index 0000000..f61a451 --- /dev/null +++ b/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 [extra pd_server args...] +# serve/pd_run.sh stop +# serve/pd_run.sh log # tail the tray's log +# +# Ports: PE HTTP 8889, CE HTTP (user-facing) 8888, CE KV receivers 8900+rank. Logs: ~/pd_.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 diff --git a/serve/pd_server.py b/serve/pd_server.py new file mode 100644 index 0000000..d94672e --- /dev/null +++ b/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://: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()