#!/usr/bin/env bash
# bench-spec-matrix.sh - Qwen3.8-27B UD-Q8_K_XL speculative-decode matrix.
# Matrix: spec_type {draft-mtp,ngram-mod | draft-mtp | ngram-mod}
#         x --spec-draft-n-max {2..10} x --spec-draft-p-min {off, 0.8}
# = 54 configs. Each: start llama-server, wait ready, 3 generations of a
# 1100-token prompt (n_predict 200, default sampling), record timings.
set -u

MODEL=/home/dank/models/Qwen3.8-27B-UD-Q8_K_XL.gguf
LLAMA_SERVER=/home/dank/llama.cpp/build/bin/llama-server
PORT=7890
DEVICE=ROCm0
OUT=/tmp/spec_matrix_results.tsv
RUNLOG=/tmp/spec_matrix_run.log
PROMPT=$(python3 -c "print(\"The fleet of inference servers must remain available at all times. Operators monitor memory, throughput, and queue depth. \"*50)")
N_OUT=200
RUNS=3

: > "$OUT"
printf 'spec_type\tn_max\tp_min\trun\ttps\tpps\tdraft_acc\tdraft_n\tstop\n' >> "$OUT"

cfg=0
for spec in "draft-mtp,ngram-mod" "draft-mtp" "ngram-mod"; do
  for n_max in 2 3 4 5 6 7 8 9 10; do
    for p_min in none 0.8; do
      cfg=$((cfg+1))
      echo "=== CONFIG $cfg/54: spec=$spec n_max=$n_max p_min=$p_min $(date +%T) ===" | tee -a "$RUNLOG"
      ARGS=(--model "$MODEL" --spec-type "$spec" --spec-draft-n-max "$n_max")
      # Qwen3.8 bundles its MTP drafter inside the target GGUF; no separate
      # --spec-draft-model file (llama.cpp auto-detects, b10411+).
      if [[ "$p_min" != none ]]; then
        ARGS+=(--spec-draft-p-min "$p_min")
      fi
      ARGS+=(--device "$DEVICE" --gpu-layers all --cache-type-k q8_0 --cache-type-v f16
             --ctx-size 2048 --parallel 1 --batch-size 512 --ubatch-size 128
             --temp 1.0 --top-k 20 --top-p 0.95 --min-p 0.0
             --host 127.0.0.1 --port "$PORT")
      "$LLAMA_SERVER" "${ARGS[@]}" > /tmp/matrix_server.log 2>&1 &
      SPID=$!
      ready=0
      for i in $(seq 1 150); do
        if curl -sf "http://127.0.0.1:$PORT/health" > /dev/null 2>&1; then ready=1; break; fi
        if ! kill -0 "$SPID" 2>/dev/null; then break; fi
        sleep 2
      done
      if [[ "$ready" != 1 ]]; then
        echo "NOT_READY: spec=$spec n_max=$n_max p_min=$p_min" | tee -a "$RUNLOG"
        tail -4 /tmp/matrix_server.log >> "$RUNLOG"
        kill "$SPID" 2>/dev/null; wait "$SPID" 2>/dev/null
        continue
      fi
      for r in 1 2 3; do
        curl -s -m 300 -X POST "http://127.0.0.1:$PORT/completion" -H "Content-Type: application/json" \
          -d "{\"prompt\": \"$PROMPT\", \"n_predict\": $N_OUT, \"stream\": false, \"cache_prompt\": false}" \
          | python3 -c "
import json,sys
try:
    d=json.load(sys.stdin)
    t=d.get('timings',{})
    row=[\"$spec\",\"$n_max\",\"$p_min\",\"$r\",
         str(round(t.get('predicted_per_second',0),2)),
         str(round(t.get('prompt_per_second',0),2)),
         str(t.get('draft_n_accepted','')),
         str(t.get('draft_n','')),
         str(d.get('stop_type',''))]
    with open(\"$OUT\",\"a\") as f:
        f.write('\t'.join(row)+'\n')
except Exception as e:
    with open(\"$OUT\",\"a\") as f:
        f.write('\t'.join([\"$spec\",\"$n_max\",\"$p_min\",\"$r\",\"ERR\",str(e),\"\",\"\",\"\"])+'\n')
"
      done
      kill "$SPID" 2>/dev/null; wait "$SPID" 2>/dev/null
      sleep 3
    done
  done
done
echo "=== MATRIX DONE $(date) ===" | tee -a "$RUNLOG"