decode.py 1.7 KB

123456789101112131415161718192021222324252627
  1. # Per-workload decode with DFlash2 acceptance: thinking off, temperature 0,
  2. # 512 tokens, median of 3 after one discarded warm-up. Same prompts as the
  3. # 2026-09-06..23 measurements.
  4. import json, os, re, statistics, sys, time, urllib.request
  5. BASE = os.environ.get("GATE_URL", "http://127.0.0.1:8002")
  6. MODEL = os.environ.get("GATE_MODEL", "glm53")
  7. W = [("structured", "Count from 1 to 200, comma separated. No commentary."),
  8. ("code", "Write a red-black tree in Python with insert, delete and rebalancing. Code only."),
  9. ("prose", "Explain how a hash map works, in flowing prose. No code, no lists.")]
  10. def metrics():
  11. raw = urllib.request.urlopen(BASE + "/metrics", timeout=30).read().decode()
  12. get = lambda k: sum(float(v) for v in re.findall(rf"^vllm:{k}\S*\s+([0-9.e+]+)$", raw, re.M))
  13. return get("spec_decode_num_accepted_tokens_total"), get("spec_decode_num_draft_tokens_total")
  14. def run(prompt, n=512):
  15. body = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": prompt}],
  16. "max_tokens": n, "temperature": 0, "chat_template_kwargs": {"thinking": False}}).encode()
  17. r = urllib.request.Request(BASE + "/v1/chat/completions", data=body, headers={"Content-Type": "application/json"})
  18. t = time.time(); d = json.loads(urllib.request.urlopen(r, timeout=600).read()); e = time.time() - t
  19. return d["usage"]["completion_tokens"] / e
  20. print("%-12s %10s %8s %s" % ("workload", "tok/s", "accept", "runs"))
  21. for name, p in W:
  22. run(p, 16)
  23. a0, d0 = metrics()
  24. rates = [run(p) for _ in range(3)]
  25. a1, d1 = metrics()
  26. acc = (a1 - a0) / (d1 - d0) if d1 > d0 else float("nan")
  27. print("%-12s %10.1f %7.1f%% %s" % (name, statistics.median(rates), acc * 100, " ".join("%.1f" % r for r in rates)), flush=True)