| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950 |
- """Decode rate over N short-prompt runs, plus DFlash2 draft acceptance.
- Acceptance length, not rate, is what predicts throughput for a k-token drafter.
- """
- import json, re, statistics, sys, time, urllib.request
- BASE, MODEL, N = sys.argv[1], sys.argv[2], int(sys.argv[3])
- def metrics():
- raw = urllib.request.urlopen(BASE.replace("/v1", "") + "/metrics", timeout=30).read().decode()
- out = {}
- for key in ("spec_decode_num_accepted_tokens_total", "spec_decode_num_draft_tokens_total"):
- m = re.search(rf"^vllm:{key}\S*\s+([0-9.e+]+)$", raw, re.M)
- if m:
- out[key] = float(m.group(1))
- return out
- def once(prompt, max_tokens):
- body = {"model": MODEL, "messages": [{"role": "user", "content": prompt}],
- "max_tokens": max_tokens, "temperature": 0, "stream": True,
- "stream_options": {"include_usage": True}}
- req = urllib.request.Request(BASE + "/chat/completions", json.dumps(body).encode(),
- {"Content-Type": "application/json"})
- t0 = time.time(); ttft = None; usage = None
- with urllib.request.urlopen(req, timeout=1800) as r:
- for raw in r:
- line = raw.decode("utf-8", "replace").strip()
- if not line.startswith("data:"): continue
- d = line[5:].strip()
- if d == "[DONE]": break
- ev = json.loads(d)
- if ev.get("usage"): usage = ev["usage"]
- for ch in ev.get("choices", []):
- delta = ch.get("delta", {})
- if (delta.get("content") or delta.get("reasoning")) and ttft is None:
- ttft = time.time() - t0
- total = time.time() - t0
- completion = (usage or {}).get("completion_tokens", 0)
- return (completion - 1) / (total - ttft) if ttft and total > ttft else 0
- before = metrics()
- rates = [once("Explain unified memory on GB10 in two sentences.", 200) for _ in range(N)]
- after = metrics()
- acc = after.get("spec_decode_num_accepted_tokens_total", 0) - before.get("spec_decode_num_accepted_tokens_total", 0)
- draft = after.get("spec_decode_num_draft_tokens_total", 0) - before.get("spec_decode_num_draft_tokens_total", 0)
- print(f" decode tok/s over {N} runs: " + " ".join(f"{r:.1f}" for r in sorted(rates)))
- print(f" median {statistics.median(rates):.1f} mean {statistics.mean(rates):.1f} "
- f"spread {max(rates)-min(rates):.1f}")
- if draft:
- print(f" draft: accepted {acc:.0f} of {draft:.0f} = {acc/draft:.3f}; "
- f"accepted per step {acc/(draft/7):.2f} of 7")
|
