File size: 2,356 Bytes
ee52468 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """Decode speed with experts streamed from the SSD, for several prefetch settings: per answer the decode tok/s and the
time the NPU waited for expert reads per token. Each setting in a fresh process (cold cache), settings interleaved.
usage: python bench_decode.py MODEL_DIR GB K1,K2,... [ROUNDS=2]"""
import json, os, subprocess, sys
CHILD = r'''
import json, sys, time
from onw.chat import ChatEngine
e = ChatEngine(sys.argv[1], "NPU", pld=False)
qs = ["NPUとGPUの違いを、身近なたとえを使って説明してください。",
"ある商品を定価の2割引きで買うと960円でした。定価はいくらですか？途中の式も書いてください。",
"Explain the difference between a process and a thread, with a short Python example."]
out = []
for q in qs:
e.checkpoint = None
b0, st = e.model.bank.stats(), None
toks = []
for x in e.stream_chat([{"role": "user", "content": q}], 192):
if isinstance(x, dict):
st = x
else:
toks.append(x)
b1 = e.model.bank.stats()
n = st["completion_tokens"]
out.append({"tok_s": round(st["decode_tok_s"], 2), "wait_ms_per_tok": round((b1["read_s"] - b0["read_s"]) * 1000 / n, 1),
"hit": round(b1["hit_rate"], 3), "text": "".join(toks)[:30]})
print("RESULT " + json.dumps(out, ensure_ascii=False), flush=True)
'''
def main():
d, gb, ks = sys.argv[1], sys.argv[2], sys.argv[3].split(",")
rounds = int(sys.argv[4]) if len(sys.argv) > 4 else 2
for r in range(rounds):
for k in ks:
env = {**os.environ, "PYTHONUTF8": "1", "ONW_EXPERT_GB": gb, "ONW_PREFETCH_K": k}
p = subprocess.run([sys.executable, "-c", CHILD, d], capture_output=True, text=True, encoding="utf8",
errors="replace", env=env)
line = next((l for l in p.stdout.splitlines() if l.startswith("RESULT ")), None)
if not line:
print(f"round {r} K={k}: FAILED {p.stderr[-300:]}", flush=True)
continue
res = json.loads(line[7:])
print(f"round {r} K={k:>2}: " + " | ".join(f"{x['tok_s']:5.2f} tok/s, wait {x['wait_ms_per_tok']:5.1f} ms/tok"
for x in res) + f" | {res[0]['text']!r}", flush=True)
if __name__ == "__main__":
main()
|