| 1234567891011121314151617181920212223242526272829303132 |
- """Cold-cache prefill throughput: unique filler each time so nothing is cached."""
- import json, sys, time, urllib.request, uuid
- BASE, MODEL = sys.argv[1], sys.argv[2]
- def once(target):
- nonce = uuid.uuid4().hex
- unit = f"Reading {nonce} note {{n}}: the pump was serviced and the filter replaced.\n"
- text = "".join(unit.format(n=i) for i in range(target // 16))
- body = {"model": MODEL, "messages": [{"role": "user", "content": text + "\nReply with OK."}],
- "max_tokens": 4, "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=3600) 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", []):
- d = ch.get("delta") or {}
- if (d.get("content") or d.get("reasoning") or d.get("reasoning_content")) and ttft is None:
- ttft = time.time() - t0
- n = (usage or {}).get("prompt_tokens", 0)
- if ttft:
- print(f" prompt={n:>7} ttft={ttft:6.1f}s prefill={n/ttft:8.0f} tok/s", flush=True)
- else:
- print(f" prompt={n:>7} no first token seen (usage={usage})", flush=True)
- for t in ([int(a) for a in sys.argv[3:]] or [8000, 42000, 128000]):
- once(t)
|
