mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
d8f7a2e898
- serve/: the minimal OpenAI-compatible server of deepseek-ai/DeepInfer#147, with a Neuron worker setup (basic mode env, per-rank NeuronCore, fixed KV pool via --num-kv-blocks, no MTP lookahead) and a /v1/completions endpoint (prompt text or token ids, no chat template): what vllm bench serve --backend openai drives, as vllm-neuron#71's run_bench.py does. - examples/neuron/bench_neuron.py: the same TTFT/TPOT matrix offline. - static MoE: one packed fp32 all-gather instead of three (each 64-rank collective is ~2 ms of fixed latency a layer); a capacity-grouped expert path for prefill chunks (DEEPINFER_NEURON_MOE_CAPACITY, dense rerun on overflow; not yet validated on device). - Sub-region timings at split() boundaries under DEEPINFER_NEURON_REGION_TIMING. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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)
|