mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 23:37:28 +08:00
c6faea3bc2
serve/openai_server.py exposes /v1/chat/completions (streaming, tools, usage) on top of LLMEngine: one engine process per GPU (attention DP + MoE EP, the layout of examples/benchmark.py), prompts rendered with the checkpoint's own encoding.py, conversation-affinity routing, a watchdog that fails open requests if an engine rank dies, and /deepinfer/metrics. serve/server.sh runs it as a process group so a stop also reaps the engine's helper processes, which otherwise keep the GPUs. It integrates the host prefix cache from deepinfer/hicache when present (--hicache-size-gb, --hicache-state-slots, per-request "prefix_cache": false bypass; serve/check_prefix_cache.py is the A/B check) and serves without it otherwise. It also sets DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL to 2 by default (--prefill-any-interval): with the engine default of 100, a rank whose peers are idle admits a waiting prompt only every 100 decode steps, a ~1.4 s TTFT floor on every request in a serving setting. Used to drive the InferenceX AgentX replay (aiperf) on 4x GB300. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK
88 lines
3.6 KiB
Python
88 lines
3.6 KiB
Python
"""A/B check of the host prefix cache on a running serve/openai_server.py.
|
|
|
|
Builds a multi-turn conversation with a long shared prefix, sends every turn twice (cache bypassed via
|
|
"prefix_cache": false, then with the cache) with greedy sampling, and compares outputs and TTFT.
|
|
|
|
python serve/check_prefix_cache.py --port 8888 --turns 3 --prefix-tokens 6000
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
|
|
import requests
|
|
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--port", type=int, default=8888)
|
|
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=6000, help="approximate size of the shared document")
|
|
p.add_argument("--max-tokens", type=int, default=64)
|
|
args = p.parse_args()
|
|
URL = f"http://localhost:{args.port}/v1/chat/completions"
|
|
|
|
# A deterministic pseudo-document; ~1 token per word for this tokenizer is close enough.
|
|
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.",
|
|
"How many distinct word stems are used in the document? Answer briefly.",
|
|
"Quote the first five words of the document.",
|
|
"What is the last word of the document?",
|
|
]
|
|
|
|
|
|
def ask(messages: list[dict], use_cache: bool) -> tuple[str, float, float, dict]:
|
|
body = {
|
|
"model": args.model,
|
|
"messages": messages,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
"temperature": 0,
|
|
"max_tokens": args.max_tokens,
|
|
"seed": 1,
|
|
"prefix_cache": use_cache,
|
|
}
|
|
t0 = time.perf_counter()
|
|
ttft = None
|
|
text, usage = [], {}
|
|
with requests.post(URL, 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"):
|
|
usage = obj["usage"]
|
|
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, usage
|
|
|
|
|
|
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)]})
|
|
base, ttft_b, tot_b, usage_b = ask(messages, use_cache=False)
|
|
cached, ttft_c, tot_c, usage_c = ask(messages, use_cache=True)
|
|
same = base == cached
|
|
mismatches += not same
|
|
print(
|
|
f"turn {turn + 1}: prompt {usage_c.get('prompt_tokens')} tok | TTFT no-cache {ttft_b * 1e3:7.0f} ms, cache {ttft_c * 1e3:7.0f} ms "
|
|
f"({ttft_b / max(ttft_c, 1e-6):.1f}x) | outputs {'identical' if same else 'DIFFER'}"
|
|
)
|
|
if not same:
|
|
print(" no-cache:", repr(base[:200]))
|
|
print(" cache :", repr(cached[:200]))
|
|
messages.append({"role": "assistant", "content": base})
|
|
|
|
stats = requests.get(f"http://localhost:{args.port}/deepinfer/metrics", timeout=30).text
|
|
print("\n".join(line for line in stats.splitlines() if "hicache:hits" in line or "hicache:lookups" in line or "hit_rate" in line))
|
|
raise SystemExit(1 if mismatches else 0)
|