#!/bin/bash
# 内核版 llama.cpp 的最终 A/B：llama-bench + 五种任务在三种投机配置下的对比
set -u
export PATH=/usr/local/cuda-13.1/bin:/usr/local/bin:/usr/bin:/bin
export LD_LIBRARY_PATH=/usr/local/cuda-13.1/lib64:${LD_LIBRARY_PATH:-}
BIN=/home/zyw/src/work/llama.cpp-9a9394a/build/bin
OUT=/home/zyw/ab
M=/home/zyw/bonsai2-mtp/Ternary-Bonsai-2-27B-PTQ1_0-mtp.gguf
mkdir -p "$OUT"

echo "############ 1) llama-bench（原版 fork 基线 PTQ1_0 = 45.12 tok/s）############"
$BIN/llama-bench -m "$M" -ngl 99 -fa on -p 512,4096 -n 128 -r 3 2>&1 | tee "$OUT/bench-patched.txt"

run_case() {
  local label="$1"; shift
  pkill -x llama-server 2>/dev/null; sleep 3
  cd /home/zyw
  env GGML_CUDA_BATCH_INVARIANT=1 nohup "$BIN/llama-server" -m "$M" \
    -ngl 99 -fa on -c 16384 -np 1 -ctk q4_0 -ctv q4_0 --jinja --no-context-shift \
    --chat-template-kwargs '{"enable_thinking":false}' "$@" \
    --host 127.0.0.1 --port 8083 > "$OUT/server-$label.log" 2>&1 &
  for i in $(seq 1 50); do
    curl -s --noproxy '*' -m 2 http://127.0.0.1:8083/health 2>/dev/null | grep -q ok && break
    sleep 1
  done
  echo
  echo "############ $label ############"
  grep -iE "adding speculative|speculative decoding context" "$OUT/server-$label.log" | head -3
  python3 /home/zyw/spec_test2.py 8083 "$label"
}

run_case "内核-关投机"
run_case "内核-MTP" --spec-type draft-mtp --spec-draft-n-max 1
run_case "内核-MTP+ngram" --spec-type draft-mtp,ngram-map-k4v --spec-draft-n-max 1

pkill -x llama-server 2>/dev/null
echo
echo "AB_FINAL_DONE"