diff --git a/bench/watt_calibrate.py b/bench/watt_calibrate.py new file mode 100644 index 0000000..4db41e5 --- /dev/null +++ b/bench/watt_calibrate.py @@ -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()) diff --git a/tests/test_watt_cogs.py b/tests/test_watt_cogs.py index 7f83d4e..0b0dad2 100644 --- a/tests/test_watt_cogs.py +++ b/tests/test_watt_cogs.py @@ -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():