mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
8e488bbd55
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
82 lines
3.2 KiB
Python
82 lines
3.2 KiB
Python
"""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)
|