coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 00:36:20 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/serve/pd_server.py
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

827 lines
37 KiB
Python
Raw Blame History

"""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
# --kv-transport mooncake: MC_MAGIC + HDR; CE replies one JSON line {"status": ok|busy|error, "addr", "segment"};
# after the TransferEngine write the PE sends b"go" (or b"xx" on failure) and the CE answers b"ok".
MC_MAGIC = b"MCK1" # cannot start a TCP-path header (rids are hex)
MC_RETRY_S = 0.005
MC_WRITE_ATTEMPTS = 3 # then fall back to TCP for that transfer, so a failing peer never strands a request
def _mooncake_link(args: argparse.Namespace, device): # noqa: ANN001, ANN202
from deepinfer.hicache.pd_mooncake import GpuArena, MooncakeLink
arena = GpuArena(int(args.pd_arena_gb * 2**30), device)
host = args.mooncake_host or (args.kv_host if args.role == "ce" and args.kv_host else socket.gethostbyname(socket.gethostname()))
return MooncakeLink(host, args.mooncake_protocol, arena, args.mooncake_device)
def _recv_line(sock: socket.socket) -> bytes:
out = bytearray()
while not out.endswith(b"\n"):
chunk = sock.recv(1)
if not chunk:
raise ConnectionError("connection closed")
out += chunk
return bytes(out)
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)
p.add_argument("--eplb-interval", type=int, default=0, help="see openai_server.py; 0 (default) disables EPLB")
p.add_argument("--eplb-after-prefill", type=int, default=0, choices=[0, 1], help="see openai_server.py")
p.add_argument("--tokenizer-backend", default="auto", choices=["auto", "fastokens", "segment", "hf"], help="see openai_server.py")
p.add_argument("--tokenizer-verify-rate", type=float, default=0.005, help="see openai_server.py")
p.add_argument(
"--kv-transport",
default="tcp",
choices=["tcp", "mooncake"],
help="KV/state data plane PE -> CE. tcp: pinned host buffers over the KV socket. mooncake: GPU arena to GPU "
"arena through a Mooncake TransferEngine (nvlink_intra within a node, rdma/efa across hosts); the KV socket "
"carries only control messages",
)
p.add_argument("--mooncake-protocol", default="nvlink_intra", help="TransferEngine protocol (nvlink_intra, rdma, efa, tcp)")
p.add_argument("--mooncake-device", default="", help="TransferEngine device name(s) for rdma/efa (empty: auto)")
p.add_argument("--mooncake-host", default=None, help="address peers use for this rank's TransferEngine handshake")
p.add_argument("--pd-arena-gb", type=float, default=4, help="mooncake: GPU staging arena per rank (PE send / CE receive)")
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))
if args.eplb_interval == 0:
# Same as openai_server.py: without EPLB the redundant expert slots only add expert weight per rank.
os.environ.setdefault("DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS", "0")
setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False)
if args.eplb_interval > 0:
os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval)
os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill)
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)
link = _mooncake_link(args, transfer.device) if args.kv_transport == "mooncake" else None
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
region = None
if link is not None:
from deepinfer.hicache.pd_mooncake import region_layout
nb = transfer.num_blocks(prompt_len)
state_off, total = region_layout(nb, transfer.page_bytes, transfer.state_bytes)
off = link.arena.alloc(total)
if off is not None: # else: arena full, this one goes over TCP
region = (off, total)
pages = link.arena.view(off, nb * transfer.page_bytes).view(nb, transfer.page_bytes)
state = link.arena.view(off + state_off, transfer.state_bytes)
ev = transfer.export_into(seq, prompt_len, pages, state)
if region is None:
pages, state, ev = transfer.export(seq, prompt_len)
item = {
"pages": pages,
"state": state,
"region": region,
"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 send_mooncake(rid: str, item: dict, host: str, port: int) -> None:
pages, state = item["pages"], item["state"]
hdr = MC_MAGIC + HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel())
off, total = item["region"]
deadline = time.monotonic() + 600
while True:
with socket.create_connection((host, port), timeout=60) as s:
s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
s.sendall(hdr)
resv = json.loads(_recv_line(s))
if resv["status"] == "ok":
for attempt in range(MC_WRITE_ATTEMPTS):
try:
link.write(resv["segment"], off, resv["addr"], total)
break
except Exception as e: # noqa: BLE001
print(
f"[pe rank {dp_rank}] mooncake write to {resv['segment']} failed (attempt {attempt + 1}): {e!r}", flush=True
)
time.sleep(0.01 * (attempt + 1))
else:
s.sendall(b"xx")
return False # caller falls back to TCP
s.sendall(b"go")
ack = s.recv(2)
if ack != b"ok":
raise RuntimeError(f"CE rejected transfer: {ack!r}")
return True
if resv["status"] != "busy" or time.monotonic() > deadline:
raise RuntimeError(f"CE reservation failed: {resv}")
time.sleep(MC_RETRY_S) # CE arena full: retry once it has admitted something
def sender():
if link is not None:
import torch
torch.cuda.set_device(transfer.device) # TransferEngine copies use the calling thread's device
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()
if item["region"] is not None:
try:
sent = send_mooncake(rid, item, host, port)
if not sent: # copy the region to pinned host buffers and use the TCP path below
item["pages"] = item["pages"].to("cpu").pin_memory()
item["state"] = item["state"].to("cpu").pin_memory()
finally:
link.arena.release(*item["region"])
if sent:
continue
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, arena region or None)
link = _mooncake_link(args, transfer.device) if args.kv_transport == "mooncake" else None
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_mooncake(conn: socket.socket, rid: str, first_token: int, prompt_len: int, num_pages: int) -> None:
from deepinfer.hicache.pd_mooncake import region_layout
state_off, total = region_layout(num_pages, transfer.page_bytes, transfer.state_bytes)
off = link.arena.alloc(total) if link is not None else None
if off is None:
status = "busy" if link is not None and total <= link.arena.nbytes else "error"
conn.sendall(json.dumps({"status": status}).encode() + b"\n")
return
try:
conn.sendall(json.dumps({"status": "ok", "addr": link.arena.base + off, "segment": link.segment}).encode() + b"\n")
go = bytearray(2)
_recv_exact(conn, memoryview(go))
if go != b"go":
raise RuntimeError(f"PE aborted the transfer: {bytes(go)!r}")
except BaseException:
link.arena.release(off, total)
raise
pages = link.arena.view(off, num_pages * transfer.page_bytes).view(num_pages, transfer.page_bytes)
state = link.arena.view(off + state_off, transfer.state_bytes)
recv_q.put((rid, first_token, prompt_len, pages, state, (off, total)))
conn.sendall(b"ok")
def receive_one(conn: socket.socket):
try:
with conn:
hdr = bytearray(HDR.size)
_recv_exact(conn, memoryview(hdr)[: len(MC_MAGIC)])
mooncake = bytes(hdr[: len(MC_MAGIC)]) == MC_MAGIC
_recv_exact(conn, memoryview(hdr)[0 if mooncake else len(MC_MAGIC) :])
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"
if mooncake:
receive_mooncake(conn, rid, first_token, prompt_len, num_pages)
return
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, None))
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, tuple] = {} # rid -> (admit event, pages, state, arena region or None)
def release(region: tuple[int, int] | None) -> None:
if region is not None:
link.arena.release(*region)
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)
gone = arrived.pop(rid, None)
if gone is not None:
release(gone[4])
if rid in pending:
pending.remove(rid)
else:
engine.abort_request(rid)
while True:
try:
rid, first_token, prompt_len, pages, state, region = recv_q.get_nowait()
except queue.Empty:
break
arrived[rid] = (first_token, prompt_len, pages, state, region)
if rid in expected:
pending.append(rid)
while pending:
rid = pending[0]
spec, (first_token, prompt_len, pages, state, region) = 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)
release(region)
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"))
release(region)
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, region)
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, k in keep.items() if k[0].query()]:
release(keep.pop(rid)[3])
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()
Reference in New Issue View Git Blame Copy Permalink