diff --git a/scripts/bench-spec-decode.py b/scripts/bench-spec-decode.py new file mode 100755 index 0000000..15f9485 --- /dev/null +++ b/scripts/bench-spec-decode.py @@ -0,0 +1,232 @@ +#!/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() diff --git a/scripts/serve-frank-large.sh b/scripts/serve-frank-large.sh index 62efeea..77fb539 100755 --- a/scripts/serve-frank-large.sh +++ b/scripts/serve-frank-large.sh @@ -92,26 +92,31 @@ if [[ ! -x "$LLAMA_BIN" ]]; then fi # --- speculative decoding (MTP draft head) -------------------------------- -# The MTP head is loaded as a separate GGUF sidecar via --model-draft. +# The MTP draft head is loaded as a separate GGUF sidecar via --model-draft. +# ngram-only specs (ngram-mod, etc.) do NOT need a draft model file. # Set SPEC_TYPE=none in the env to disable speculative decoding entirely. SPEC_ARGS=() if [[ -n "$SPEC_TYPE" && "$SPEC_TYPE" != "none" ]]; then - DRAFT_MODEL="$DRAFT_MODEL_DIR/$DRAFT_MODEL_FILE" - if [[ ! -r "$DRAFT_MODEL" ]]; then - echo "missing or unreadable draft model: $DRAFT_MODEL" >&2 - exit 1 - fi SPEC_ARGS=( --spec-type "$SPEC_TYPE" - --model-draft "$DRAFT_MODEL" --spec-draft-n-max "$SPEC_DRAFT_N_MAX" --spec-draft-n-min "$SPEC_DRAFT_N_MIN" --spec-draft-p-min "$SPEC_DRAFT_P_MIN" - --spec-draft-ngl "$SPEC_DRAFT_NGL" - --spec-draft-device "$SPEC_DRAFT_DEVICE" - --spec-draft-type-k "$SPEC_DRAFT_TYPE_K" - --spec-draft-type-v "$SPEC_DRAFT_TYPE_V" ) + if [[ "$SPEC_TYPE" == *draft-mtp* ]]; then + DRAFT_MODEL="$DRAFT_MODEL_DIR/$DRAFT_MODEL_FILE" + if [[ ! -r "$DRAFT_MODEL" ]]; then + echo "missing or unreadable draft model: $DRAFT_MODEL" >&2 + exit 1 + fi + SPEC_ARGS+=( + --model-draft "$DRAFT_MODEL" + --spec-draft-ngl "$SPEC_DRAFT_NGL" + --spec-draft-device "$SPEC_DRAFT_DEVICE" + --spec-draft-type-k "$SPEC_DRAFT_TYPE_K" + --spec-draft-type-v "$SPEC_DRAFT_TYPE_V" + ) + fi fi SERVER_DIR="$(cd -- "$(dirname -- "$LLAMA_BIN")" && pwd)"