bench: spec decode sweep + fix serve script ngram-only path
serve-frank-large.sh: only pass --model-draft when SPEC_TYPE contains 'draft-mtp'. ngram-only configs (ngram-mod, ngram-simple, etc.) do not need a draft model file — loading one wastes VRAM and may error. bench-spec-decode.py: new benchmark script. Sweeps 4 configs (none, ngram-mod, draft-mtp, draft-mtp+ngram-mod), generates 16k tokens per config at temp 0 from a minimal prompt. Uses HTTP API (llama-bench cannot target --spec-type). Follows benchmarking policy: stop+mask the production service, bench, unmask+restart, health check.
This commit is contained in:
Frank
Executable
+232
@@ -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()
|
||||
@@ -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)"
|
||||
|
||||
Reference in New Issue
Block a user