feat(#000057): bench/watt_calibrate.py — separate prefill vs decode energy
fox 2026-05-21: account wattage for input and output separately. Prefill (process all prompt tokens, parallel/compute-bound) and decode (generate output, autoregressive/bandwidth-bound) are different GPU ops with different J/token — a single per-token number can't represent both. Slope calibration (no sub-request power alignment): sweep prompt length at tiny max_tokens -> prefill J/input-tok (fixed overhead cancels in the slope); fix a tiny prompt and sweep forced output length (ignore_eos) -> decode J/output-tok. Prefill kept COLD (unique filler so cached_tokens=0). Reuses watt_bench probes. Bad points (context overflow) skip, not abort. Measured qwen-nothink/4090 @$0.33/kWh: prefill 0.175 J/tok ($0.016/M-input-tok), decode 6.16 J/tok ($0.564/M-output-tok) — decode 35x dearer per token. Predicts measured substrate J/q within ~5%. 14 tests (+ slope). Validated live.
This commit is contained in:
parent
459060c774
commit
892d9ed037
2 changed files with 230 additions and 0 deletions
218
bench/watt_calibrate.py
Normal file
218
bench/watt_calibrate.py
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Prefill-vs-decode energy calibration (#000057).
|
||||
|
||||
fox 2026-05-21: account wattage for INPUT and OUTPUT separately. An LLM
|
||||
request is prefill (process all prompt tokens, parallel/compute-bound) +
|
||||
decode (generate output tokens, autoregressive/bandwidth-bound) — two
|
||||
different GPU operations with different J/token. A single per-token
|
||||
number hides that.
|
||||
|
||||
Method — SLOPE CALIBRATION (no sub-request power alignment needed):
|
||||
|
||||
* prefill J/input-tok — sweep prompt length at a small fixed
|
||||
max_tokens; the slope of (GPU joules / request) vs measured
|
||||
prompt_tokens isolates prefill. Fixed per-request + decode overhead
|
||||
is constant across the sweep, so it cancels in the slope.
|
||||
* decode J/output-tok — fix a tiny prompt, sweep FORCED output length
|
||||
(``ignore_eos`` + max_tokens); slope of (J/request) vs completion
|
||||
tokens isolates decode.
|
||||
|
||||
Then any query's GPU energy ≈ prefill_J·prompt_tok + decode_J·completion_tok,
|
||||
giving SEPARATE $/input-tok and $/output-tok.
|
||||
|
||||
Prefill is kept COLD: each request gets unique filler so the server's
|
||||
prompt cache (cached_tokens) can't shortcut it — we measure true
|
||||
uncached prefill (at our scale there is no prefill-cache reuse policy).
|
||||
|
||||
Run ON or AGAINST the dedicated GPU box; power is sampled there. Reuses
|
||||
bench.watt_bench probes (DRY).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from bench.watt_bench import ( # noqa: E402
|
||||
MODELS, LocalProbe, RemoteProbe, _make_client,
|
||||
)
|
||||
|
||||
_J_PER_KWH = 3.6e6
|
||||
|
||||
|
||||
_WORDS = ("time year people way day man thing woman life child world school "
|
||||
"state family student group country problem hand part place case "
|
||||
"week company system program work water room money story month").split()
|
||||
|
||||
|
||||
def _filler(approx_tokens: int, salt: int) -> str:
|
||||
"""~approx_tokens of UNIQUE text (salt prefix keeps prefill cache cold).
|
||||
|
||||
Common words tokenize ~1 token each, so the requested size ≈ real
|
||||
prompt tokens — keeps sweep points safely under the context window.
|
||||
Real token count from usage is still what the slope uses."""
|
||||
return f"u{salt}: " + " ".join(
|
||||
_WORDS[(w + salt) % len(_WORDS)] for w in range(approx_tokens))
|
||||
|
||||
|
||||
def _slope(xs: list[float], ys: list[float]) -> float:
|
||||
"""Least-squares slope dy/dx (≥2 points)."""
|
||||
n = len(xs)
|
||||
mx = sum(xs) / n
|
||||
my = sum(ys) / n
|
||||
num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
|
||||
den = sum((x - mx) ** 2 for x in xs)
|
||||
return num / den if den else 0.0
|
||||
|
||||
|
||||
def _drive(client, model, extra, reps, build_messages, max_tokens):
|
||||
"""Drive `reps` requests; return (sum_prompt_tok, sum_completion_tok,
|
||||
sum_cached_tok)."""
|
||||
sp = sc = scache = ok = 0
|
||||
for i in range(reps):
|
||||
try:
|
||||
ans = client.chat_completion(build_messages(i), model=model,
|
||||
max_tokens=max_tokens, extra_body=extra)
|
||||
except Exception as e: # noqa: BLE001 — skip a bad point (e.g. ctx overflow)
|
||||
print(f" request {i} failed ({type(e).__name__}); skipping")
|
||||
continue
|
||||
u = getattr(client, "last_usage", None) or {}
|
||||
sp += int(u.get("prompt_tokens") or 0)
|
||||
sc += int(u.get("completion_tokens") or len(ans) // 4)
|
||||
scache += int((u.get("prompt_tokens_details") or {}).get(
|
||||
"cached_tokens") or 0)
|
||||
ok += 1
|
||||
return sp, sc, scache, ok
|
||||
|
||||
|
||||
def _new_probe(a):
|
||||
if a.remote_gpu_host.strip():
|
||||
return RemoteProbe(a.remote_gpu_host.strip(), a.gpu_index,
|
||||
a.gpu_label, a.sample_interval)
|
||||
return LocalProbe(a.gpu_index, a.sample_interval, None)
|
||||
|
||||
|
||||
def _sweep_point(a, client, model, extra, reps, build_messages, max_tokens):
|
||||
"""One sweep point: measure GPU joules over `reps` driven requests."""
|
||||
probe = _new_probe(a)
|
||||
probe.start()
|
||||
sp, sc, scache, ok = _drive(client, model, extra, reps,
|
||||
build_messages, max_tokens)
|
||||
probe.stop()
|
||||
if ok == 0:
|
||||
return None # whole point failed (e.g. context overflow) — skip
|
||||
j = probe.gpu_stats().get("joules")
|
||||
return {"gpu_joules": j, "ok_reps": ok,
|
||||
"j_per_req": (j / ok if j else None),
|
||||
"prompt_tok_per_req": sp / ok, "completion_tok_per_req": sc / ok,
|
||||
"cached_tok_per_req": scache / ok}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--remote-gpu-host", default="")
|
||||
ap.add_argument("--gpu-index", type=int, default=0)
|
||||
ap.add_argument("--gpu-label", default="gpu")
|
||||
ap.add_argument("--model-key", default="qwen-nothink",
|
||||
help="MODELS key (endpoint+model+extra)")
|
||||
ap.add_argument("--endpoint", default="")
|
||||
ap.add_argument("--price-per-kwh", type=float, default=0.33)
|
||||
ap.add_argument("--reps", type=int, default=8)
|
||||
ap.add_argument("--prefill-tokens", default="512,4096,16384",
|
||||
help="prompt-length sweep points (≈ real tokens; keep "
|
||||
"under the model context window)")
|
||||
ap.add_argument("--decode-tokens", default="32,128,512",
|
||||
help="forced output-length sweep points")
|
||||
ap.add_argument("--sample-interval", type=float, default=0.25)
|
||||
ap.add_argument("--out-dir", default="bench/qa_results")
|
||||
a = ap.parse_args()
|
||||
|
||||
cfg = MODELS[a.model_key]
|
||||
model = cfg["model"]
|
||||
base_extra = dict(cfg.get("extra") or {})
|
||||
endpoint = a.endpoint or cfg["endpoint"]
|
||||
client = _make_client(endpoint)
|
||||
ts = time.strftime("%Y-%m-%dT%H-%M-%SZ", time.gmtime())
|
||||
print(f"watt_calibrate model={a.model_key} endpoint={endpoint} "
|
||||
f"gpu={a.gpu_label} reps={a.reps}")
|
||||
|
||||
# --- prefill sweep: vary prompt length, tiny fixed decode ----------
|
||||
prefill_pts = []
|
||||
for P in [int(x) for x in a.prefill_tokens.split(",")]:
|
||||
def build(i, P=P):
|
||||
return [{"role": "user", "content": _filler(P, salt=i + P * 1000)}]
|
||||
pt = _sweep_point(a, client, model, base_extra, a.reps, build,
|
||||
max_tokens=4)
|
||||
if pt is None:
|
||||
print(f" prefill P~{P}: all reps failed (context overflow?) — skipped")
|
||||
continue
|
||||
pt["sweep"] = "prefill"; pt["target_prompt_tokens"] = P
|
||||
prefill_pts.append(pt)
|
||||
print(f" prefill P~{P}: prompt={pt['prompt_tok_per_req']:.0f}tok/req "
|
||||
f"cached={pt['cached_tok_per_req']:.0f} "
|
||||
f"J/req={pt['j_per_req']}")
|
||||
|
||||
# --- decode sweep: tiny prompt, forced output length ---------------
|
||||
decode_extra = {**base_extra, "ignore_eos": True}
|
||||
decode_pts = []
|
||||
for M in [int(x) for x in a.decode_tokens.split(",")]:
|
||||
def build(i):
|
||||
return [{"role": "user", "content": f"Count upward. (seed {i})"}]
|
||||
pt = _sweep_point(a, client, model, decode_extra, a.reps, build,
|
||||
max_tokens=M)
|
||||
if pt is None:
|
||||
print(f" decode M~{M}: all reps failed — skipped")
|
||||
continue
|
||||
pt["sweep"] = "decode"; pt["target_completion_tokens"] = M
|
||||
decode_pts.append(pt)
|
||||
print(f" decode M~{M}: completion={pt['completion_tok_per_req']:.0f}"
|
||||
f"tok/req J/req={pt['j_per_req']}")
|
||||
client.close()
|
||||
|
||||
# --- slopes = per-token energy -------------------------------------
|
||||
if len(prefill_pts) < 2 or len(decode_pts) < 2:
|
||||
print(f"\nINSUFFICIENT POINTS: prefill={len(prefill_pts)} "
|
||||
f"decode={len(decode_pts)} (need ≥2 each for a slope). "
|
||||
"Adjust --prefill-tokens / --decode-tokens.")
|
||||
return 1
|
||||
prefill_j_per_tok = _slope([p["prompt_tok_per_req"] for p in prefill_pts],
|
||||
[p["j_per_req"] for p in prefill_pts])
|
||||
decode_j_per_tok = _slope([p["completion_tok_per_req"] for p in decode_pts],
|
||||
[p["j_per_req"] for p in decode_pts])
|
||||
price = a.price_per_kwh
|
||||
|
||||
def usd_per_mtok(j_per_tok):
|
||||
return round(j_per_tok / _J_PER_KWH * price * 1e6, 4)
|
||||
|
||||
report = {
|
||||
"ts": ts, "model_key": a.model_key, "model": model,
|
||||
"endpoint": endpoint, "gpu_label": a.gpu_label, "reps": a.reps,
|
||||
"price_per_kwh": price,
|
||||
"prefill_points": prefill_pts, "decode_points": decode_pts,
|
||||
"prefill_j_per_input_tok": round(prefill_j_per_tok, 5),
|
||||
"decode_j_per_output_tok": round(decode_j_per_tok, 5),
|
||||
"prefill_usd_per_M_input_tok": usd_per_mtok(prefill_j_per_tok),
|
||||
"decode_usd_per_M_output_tok": usd_per_mtok(decode_j_per_tok),
|
||||
"decode_to_prefill_ratio": (round(decode_j_per_tok / prefill_j_per_tok, 2)
|
||||
if prefill_j_per_tok else None),
|
||||
}
|
||||
outp = Path(a.out_dir) / f"watt_calibrate_{a.gpu_label}_{ts}.json"
|
||||
outp.parent.mkdir(parents=True, exist_ok=True)
|
||||
outp.write_text(json.dumps(report, indent=2))
|
||||
|
||||
print("\n=== SEPARATED ENERGY COGS @ "
|
||||
f"${price}/kWh ({a.gpu_label}, {a.model_key}) ===")
|
||||
print(f" prefill (input) : {report['prefill_j_per_input_tok']} J/tok "
|
||||
f"= ${report['prefill_usd_per_M_input_tok']}/M-input-tok")
|
||||
print(f" decode (output) : {report['decode_j_per_output_tok']} J/tok "
|
||||
f"= ${report['decode_usd_per_M_output_tok']}/M-output-tok")
|
||||
print(f" decode is {report['decode_to_prefill_ratio']}x prefill per token")
|
||||
print(f" query energy ≈ prefill·input_tok + decode·output_tok")
|
||||
print(f" saved: {outp}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
|
@ -8,9 +8,21 @@ the serving-floor marginal, and the data-derived band split.
|
|||
See docs/calculator-test-patterns.md.
|
||||
"""
|
||||
from bench.watt_bench import energy_cogs
|
||||
from bench.watt_calibrate import _slope
|
||||
from bench.watt_probe import classify_power_bands
|
||||
|
||||
|
||||
def test_slope_recovers_known_rate():
|
||||
# y = 0.175*x + 60 (decode-like: per-token slope + fixed overhead).
|
||||
xs = [32.0, 128.0, 512.0]
|
||||
ys = [0.175 * x + 60 for x in xs]
|
||||
assert round(_slope(xs, ys), 4) == 0.175 # slope ignores intercept
|
||||
|
||||
|
||||
def test_slope_two_points():
|
||||
assert _slope([100.0, 200.0], [10.0, 30.0]) == 0.2
|
||||
|
||||
|
||||
# ---- energy_cogs -------------------------------------------------------
|
||||
|
||||
def test_gross_cogs_one_kwh():
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue