k1_dev/fullgraph_smoke.py

"""Test the one untested PPL-neutral lever: full-graph CUDA capture.
A = baseline (--performance-mode interactivity, as frontier).
B = + COMPILATION_CONFIG cudagraph_mode=FULL_DECODE_ONLY.
C = + cudagraph_mode=FULL.
Measure steady decode tok/s; STEPTIME for the verify GPU span. PPL-neutral
(capture changes scheduling, not math) so we only need TPS + boot success."""
import json, os, 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 {p.returncode}")
            for ln in open('/tmp/srv.log',errors='replace').readlines()[-22:]: 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 essay on distributed systems consensus, 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 e: print(f"  req{i} FAIL {e!r}")
    tps=toks/secs if secs else 0
    print(f"[{label}] decode={tps:.1f} tok/s")
    ag=[l for l in open('/tmp/srv.log',errors='replace') if '[steptime] agg' in l and 'kind=exec' in l][-1:]
    for l in ag: print("  [steptime]",l.split('agg',1)[-1].strip()[:140])
    return tps

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

res={}
res['A_base']=run("A baseline",{})
res['B_fulldec']=run("B FULL_DECODE_ONLY",{"COMPILATION_CONFIG":'{"cudagraph_mode":"FULL_DECODE_ONLY"}'})
res['C_full']=run("C FULL",{"COMPILATION_CONFIG":'{"cudagraph_mode":"FULL"}'})
print("\n[FG]"," ".join(f"{k}={v:.1f}" for k,v in res.items() if v))
if res.get('A_base'):
    for k in ('B_fulldec','C_full'):
        if res.get(k): print(f"[FG] {k} vs A: {100*(res[k]-res['A_base'])/res['A_base']:+.1f}%")