vh
4a5c3fcccf
Closes the one axis of the original premise left unverified. Measured cold (cache-busted) on both builds under matching serve configs: ~6.7k-token prompt 3,206 -> 6,334 tok/s prefill (+98%) ~27k-token prompt 2,862 -> 5,085 tok/s prefill (+78%) TTFT on a ~27k doc 9.43 -> 5.31 s (-44%) Prefill gains far exceed the +18% decode gain, and that ordering is the expected one: decode at bs=1 is memory-bandwidth-bound and the weights are 4-bit under either scheme, so little changes; prefill is compute-bound, which is where native Blackwell FP4 tensor cores replace the Marlin dequant-to-BF16 path. The summarizer aliases are the consumers that feel this. Adds bench/prefill_bench.py plus the raw JSON. The harness deliberately uses SystemRandom: a seeded nonce regenerates the previous run's prompts verbatim, prefix caching then serves them, and the first attempt read ~41k tok/s of cache-hit rather than ~5k of actual prefill.
53 lines
2.3 KiB
Python
53 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Prefill throughput via TTFT on long prompts (prefill-dominated)."""
|
|
import argparse, json, random, statistics, time, urllib.request
|
|
|
|
def ttft(base, model, prompt):
|
|
req = urllib.request.Request(base + "/v1/chat/completions",
|
|
data=json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 16, "temperature": 0, "stream": True,
|
|
"chat_template_kwargs": {"enable_thinking": False}}).encode(),
|
|
headers={"Content-Type": "application/json"})
|
|
t0 = time.perf_counter()
|
|
with urllib.request.urlopen(req, timeout=900) as r:
|
|
for line in r:
|
|
if line.startswith(b"data: ") and b"[DONE]" not in line:
|
|
dl = json.loads(line[6:])["choices"][0].get("delta", {})
|
|
if dl.get("content") or dl.get("reasoning"):
|
|
return time.perf_counter() - t0
|
|
return None
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--base", required=True)
|
|
ap.add_argument("--model", required=True)
|
|
ap.add_argument("--tag", default="run")
|
|
ap.add_argument("--out", default=None)
|
|
a = ap.parse_args()
|
|
# NO fixed seed: a seeded nonce reproduces the previous run's prompts verbatim,
|
|
# so --enable-prefix-caching serves them from cache and you measure cache-hit
|
|
# TTFT, not prefill. (First attempt read 41k tok/s that way.) Lengths stay
|
|
# identical across builds; only the nonce differs.
|
|
rnd = random.SystemRandom()
|
|
res = {"tag": a.tag}
|
|
body = ""
|
|
for reps, label in [(4000, "~6.5k"), (16000, "~26k")]:
|
|
ts = []
|
|
for _ in range(3):
|
|
n = rnd.randint(10**12, 10**13)
|
|
body = f"[doc {n}] The archived engineering report details subsystem telemetry. " * (reps // 12)
|
|
v = ttft(a.base, a.model, body + "\n\nSummarize in one sentence.")
|
|
if v:
|
|
ts.append(v)
|
|
tok = len(body) // 4
|
|
med = statistics.median(ts)
|
|
rate = tok / med
|
|
res[label] = {"approx_prompt_tokens": tok, "ttft_median_s": med, "prefill_tok_s": rate}
|
|
print(f" {label} prompt (~{tok} tok): TTFT {med:.3f}s -> ~{rate:,.0f} tok/s prefill")
|
|
if a.out:
|
|
json.dump(res, open(a.out, "w"), indent=2)
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|