arborist/bench/watt_calibrate.py

#!/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())