#!/usr/bin/env python3 """Benchmark decode throughput across speculative decoding configs. Sweeps 4 configs: none, ngram-mod, draft-mtp, draft-mtp+ngram-mod. Each config generates 16k tokens from a minimal prompt at temp 0. Usage: python3 bench-spec-decode.py [--max-tokens N] [--host H] [--port P] Requires SSH access to the target host (for service management) and HTTP access to the llama.cpp server endpoint. """ import json import os import subprocess import sys import time import urllib.request import urllib.error SERVICE = "frank-default-large" ENV_FILE = "/home/dank/serverAppConfigs/scripts/frank-default-large.env" ALIAS = "frank-default-large" DEFAULT_HOST = "localhost" DEFAULT_PORT = 7770 DEFAULT_MAX_TOKENS = 16384 HEALTH_TIMEOUT = 1200 # 20 min for cold load HEALTH_INTERVAL = 5 PROMPT = ( "Write a comprehensive, detailed essay about the complete history of " "computing, from the earliest mechanical calculators through modern AI. " "Cover every major era in detail, including key figures, machines, and " "breakthroughs. Be as thorough as possible." ) CONFIGS = [ ("none", "none"), ("ngram-mod", "ngram-mod"), ("draft-mtp", "draft-mtp"), ("draft-mtp+ngram-mod", "draft-mtp,ngram-mod"), ] def ssh(cmd): r = subprocess.run( ["ssh", "dankclaw", cmd], capture_output=True, text=True, timeout=30 ) if r.returncode != 0: print(f" ssh stderr: {r.stderr.strip()}", file=sys.stderr) return r def sed_spec_type(value): ssh(f"sudo sed -i 's/^SPEC_TYPE=.*/SPEC_TYPE={value}/' {ENV_FILE}") def wait_health(host, port, timeout=HEALTH_TIMEOUT): start = time.time() while time.time() - start < timeout: try: resp = urllib.request.urlopen( f"http://{host}:{port}/health", timeout=10 ) if resp.status == 200: return True except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError): pass time.sleep(HEALTH_INTERVAL) return False def benchmark(host, port, alias, max_tokens): body = json.dumps({ "model": alias, "messages": [{"role": "user", "content": PROMPT}], "max_tokens": max_tokens, "temperature": 0, "stream": False, }).encode() req = urllib.request.Request( f"http://{host}:{port}/v1/chat/completions", data=body, headers={"Content-Type": "application/json"}, ) t0 = time.time() resp = urllib.request.urlopen(req, timeout=600) wall = time.time() - t0 data = json.loads(resp.read()) return data, wall def main(): import argparse ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--host", default=DEFAULT_HOST) ap.add_argument("--port", type=int, default=DEFAULT_PORT) ap.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS) ap.add_argument("--ssh-host", default="dankclaw") args = ap.parse_args() # Read original SPEC_TYPE from the env file r = ssh(f"grep '^SPEC_TYPE=' {ENV_FILE}") original_spec = r.stdout.strip() print(f"Original: {original_spec}") print(f"Max tokens: {args.max_tokens}") print(f"Configs: {[c[0] for c in CONFIGS]}") print() # Stop + mask the production service print("Stopping + masking production service...") ssh(f"sudo systemctl stop {SERVICE}") ssh(f"sudo systemctl mask {SERVICE}") time.sleep(2) results = [] for name, spec_type in CONFIGS: print(f"\n{'='*60}") print(f"Config: {name} (SPEC_TYPE={spec_type})") print(f"{'='*60}") # Set SPEC_TYPE sed_spec_type(spec_type) print(f"Set SPEC_TYPE={spec_type}") # Start the service ssh(f"sudo systemctl unmask {SERVICE}") ssh(f"sudo systemctl start {SERVICE}") print("Starting service...") # Wait for health print(f"Waiting for health (up to {HEALTH_TIMEOUT}s)...") if not wait_health(args.host, args.port): print(f"FAILED: service did not become healthy for {name}") results.append({ "config": name, "spec_type": spec_type, "error": "health timeout", }) ssh(f"sudo systemctl stop {SERVICE}") ssh(f"sudo systemctl mask {SERVICE}") continue load_time = 0 # not tracking separately print("Healthy. Running benchmark...") # Run benchmark try: data, wall = benchmark( args.host, args.port, ALIAS, args.max_tokens ) except Exception as e: print(f"FAILED: {e}") results.append({ "config": name, "spec_type": spec_type, "error": str(e), }) ssh(f"sudo systemctl stop {SERVICE}") ssh(f"sudo systemctl mask {SERVICE}") continue timings = data.get("timings", {}) usage = data.get("usage", {}) pred_n = timings.get("predicted_n", usage.get("completion_tokens", 0)) pred_ms = timings.get("predicted_ms", 0) pred_ps = timings.get("predicted_per_second", 0) prompt_n = timings.get("prompt_n", usage.get("prompt_tokens", 0)) prompt_ms = timings.get("prompt_ms", 0) print(f" Generated: {pred_n} tokens") print(f" Wall time: {wall:.1f}s") print(f" Decode: {pred_ps:.2f} tok/s") print(f" Prefill: {prompt_n} tokens in {prompt_ms:.0f}ms" f" ({prompt_n / max(prompt_ms, 1) * 1000:.1f} tok/s)") results.append({ "config": name, "spec_type": spec_type, "predicted_tokens": pred_n, "predicted_ms": pred_ms, "predicted_per_second": pred_ps, "prompt_tokens": prompt_n, "prompt_ms": prompt_ms, "wall_time": wall, }) # Stop + mask between configs ssh(f"sudo systemctl stop {SERVICE}") ssh(f"sudo systemctl mask {SERVICE}") time.sleep(2) # Restore original SPEC_TYPE orig_val = original_spec.split("=", 1)[1] if "=" in original_spec else "none" sed_spec_type(orig_val) # Restart production service print(f"\nRestoring SPEC_TYPE={orig_val}") print("Unmasking + starting production service...") ssh(f"sudo systemctl unmask {SERVICE}") ssh(f"sudo systemctl start {SERVICE}") # Health check if wait_health(args.host, args.port): print("Production service healthy.") else: print("WARNING: production service not healthy after restore!") # Print results table print(f"\n{'='*70}") print("RESULTS") print(f"{'='*70}") print(f"{'Config':<30} {'Tokens':>8} {'Wall(s)':>10} {'Decode tok/s':>14} {'Error':>20}") print(f"{'-'*30} {'-'*8} {'-'*10} {'-'*14} {'-'*20}") for r in results: if "error" in r: print(f"{r['config']:<30} {'':>8} {'':>10} {'':>14} {r['error']:>20}") else: print(f"{r['config']:<30} {r['predicted_tokens']:>8} " f"{r['wall_time']:>10.1f} {r['predicted_per_second']:>14.2f} {'':>20}") print(f"{'='*70}") # Save JSON results ts = time.strftime("%Y%m%d-%H%M%S") out = f"/tmp/bench-spec-decode-{ts}.json" with open(out, "w") as f: json.dump({"results": results, "max_tokens": args.max_tokens, "timestamp": ts}, f, indent=2) print(f"Results saved: {out}") if __name__ == "__main__": main()