k1_dev/mega_ab.py

"""Decisive A/B: my drafter megakernel ON vs OFF on the SAME base, with
STEPTIME breakdown. Answers: does the megakernel collapse the draft-chain
latency, and does that convert to TPS? Same weights, only DRAFTER_MEGAKERNEL flips."""
import json, os, re, subprocess, sys, time, urllib.request

SUB = "/tmp/sub"
base = "http://127.0.0.1:8000"


def boot(extra):
    m = json.load(open(f"{SUB}/manifest.json"))
    env = dict(os.environ)
    for k, v in m.get("env", {}).items():
        env[k] = v
    env.update(extra)
    env["PORT"] = "8000"
    env["STEPTIME"] = "1"
    log = open("/tmp/srv.log", "w")
    p = subprocess.Popen([sys.executable, f"{SUB}/serve.py"], cwd=SUB, env=env,
                         stdout=log, stderr=subprocess.STDOUT)
    dl = time.time() + 1000
    while time.time() < dl:
        if p.poll() is not None:
            print(f"  EXITED code={p.returncode}")
            for ln in open('/tmp/srv.log', errors='replace').readlines()[-25:]:
                print("  [srv]", ln.rstrip()[:200])
            return None
        try:
            urllib.request.urlopen(base + "/v1/models", timeout=2); return p
        except Exception:
            time.sleep(5)
    return None


def measure(label):
    toks, secs = 0, 0.0
    for i in range(8):
        body = json.dumps({"model": "gemma-4-e4b-it",
            "messages": [{"role": "user", "content":
                f"Write a long detailed technical essay on operating system schedulers, run {i}."}],
            "max_tokens": 512, "temperature": 0.0, "stream": False}).encode()
        req = urllib.request.Request(base + "/v1/chat/completions", data=body,
                                     headers={"content-type": "application/json"})
        t0 = time.time()
        try:
            r = json.load(urllib.request.urlopen(req, timeout=180))
            dt = time.time() - t0
            ct = r.get("usage", {}).get("completion_tokens", 0)
            if i >= 3:
                toks += ct; secs += dt
        except Exception as exc:
            print(f"  req{i} FAIL {exc!r}")
    tps = toks / secs if secs else 0
    print(f"[{label}] decode = {tps:.1f} tok/s")
    return tps


def steptime_tail(label):
    # surface the latest [steptime] agg lines + mega activation
    lines = open("/tmp/srv.log", errors="replace").readlines()
    aggs = [ln for ln in lines if "[steptime] agg" in ln][-4:]
    acts = [ln for ln in lines if "[mega] ACTIVE" in ln or "[mega] DISABLED" in ln][:2]
    for ln in acts + aggs:
        print(f"  [{label}]", ln.rstrip()[:200])


def run(label, extra):
    print(f"=== {label}: {extra} ===")
    p = boot(extra)
    if p is None:
        return None
    tps = measure(label)
    steptime_tail(label)
    p.terminate(); time.sleep(10)
    return tps


off = run("OFF onegraph", {"DRAFTER_MEGAKERNEL": "0"})
on = run("ON megakernel", {"DRAFTER_MEGAKERNEL": "1", "DRAFTER_MEGAKERNEL_REQUIRE": "1"})
if off and on:
    print(f"\n[AB] onegraph={off:.1f}  megakernel={on:.1f}  delta={on-off:+.1f} ({100*(on-off)/off:+.1f}%)")