#!/usr/bin/env bash
# bench-pipeline-64k.sh — 64k-context benchmark through the FULL speculative
# pipeline (draft-dflash + ngram-mod) on a single backend.
#
# Starts llama-server with the DFlash+ngram chain, sends a ~64k-token prompt,
# and reads the authoritative prefill / generation tok/s and draft acceptance
# from the server's own `slot print_timing` lines.
#
# Usage:
#   ./bench-pipeline-64k.sh <ROCm0|Vulkan0>
set -euo pipefail

DEVICE="${1:?usage: bench-pipeline-64k.sh <ROCm0|Vulkan0>}"
MODEL_DIR=/home/dank/models
MODEL_FILE=Qwen3.6-27B-Uncensored-HauhauCS-Aggressive-Q8_K_P.gguf
DRAFT_FILE=Qwen3.6-27B-DFlash-Q8_0.gguf
LLAMA_SERVER=/home/dank/llama.cpp/build/bin/llama-server
MODEL="$MODEL_DIR/$MODEL_FILE"
DRAFT="$MODEL_DIR/$DRAFT_FILE"

PORT=7891
CTX=65536
PARALLEL=1
BATCH=2048
UBATCH=1024
OUT_TOKENS=256
LOG="/tmp/pipeline-$DEVICE.log"

for f in "$MODEL" "$DRAFT"; do
    [[ -f "$f" ]] || { echo "missing: $f" >&2; exit 1; }
done

: > "$LOG"
echo "=== 64k full-pipeline bench: device=$DEVICE ==="

"$LLAMA_SERVER" \
    --model "$MODEL" \
    --spec-draft-model "$DRAFT" \
    --spec-type "draft-dflash,ngram-mod" \
    --spec-draft-n-max 8 \
    --spec-ngram-mod-n-min 48 \
    --spec-ngram-mod-n-max 64 \
    --spec-ngram-mod-n-match 24 \
    --device "$DEVICE" \
    --gpu-layers all \
    --cache-type-k q8_0 \
    --cache-type-v f16 \
    --ctx-size "$CTX" \
    --parallel "$PARALLEL" \
    --batch-size "$BATCH" \
    --ubatch-size "$UBATCH" \
    --alias bench \
    --host 127.0.0.1 \
    --port "$PORT" \
    >> "$LOG" 2>&1 &
SERVER_PID=$!
trap 'kill $SERVER_PID 2>/dev/null; wait $SERVER_PID 2>/dev/null' EXIT

echo "server pid $SERVER_PID on 127.0.0.1:$PORT"
echo -n "waiting for server"
ready=0
for i in $(seq 1 400); do
    if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
        echo " [ready]"
        ready=1
        break
    fi
    echo -n "."
    sleep 2
done
[ "$ready" = "1" ] || { echo " FAIL"; tail -25 "$LOG"; exit 1; }

# --- drive a ~64k-token generation, then pull timings from the server log ---
python3 - "$DEVICE" "$PORT" "$OUT_TOKENS" "$LOG" <<'PY'
import sys, time, json, urllib.request, os
device, port, out_tokens, log = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]

# Build a prompt that actually occupies ~64k tokens. Tokenizer demo data:
# this model averages ~0.25 tokens/char; ~250k chars -> ~62k tokens.
unit = ("The careful study of natural language reveals patterns of meaning that "
        "emerge from the deliberate arrangement of words and phrases, and the "
        "patient reader soon learns to recognize the structure beneath the surface. ")
prompt = unit * 1900  # ~250k chars

body = json.dumps({
    "prompt": prompt,
    "max_tokens": out_tokens,
    "temperature": 0.0,
    "stream": False,
}).encode()
req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions",
                             data=body, headers={"Content-Type": "application/json"})
t0 = time.time()
with urllib.request.urlopen(req, timeout=3600) as resp:
    out = json.load(resp)
dt = time.time() - t0
pt = out.get("usage", {}).get("prompt_tokens", 0)
ct = out.get("usage", {}).get("completion_tokens", 0)
print(f"request done: prompt_tokens={pt} completion_tokens={ct} wall={dt:.1f}s", flush=True)

# The server log already printed authoritative per-slot timings; surface them.
lines = open(log, encoding="utf-8", errors="ignore").read().splitlines()
prompt_evals = [l for l in lines if "prompt eval time" in l]
evals       = [l for l in lines if "      eval time" in l and "prompt eval" not in l]
accepts     = [l for l in lines if "draft acceptance" in l]

def grab(lines, label):
    if lines:
        print(label + ": " + lines[-1].strip())

grab(prompt_evals, "PREFILL(timing)")
grab(evals, "GENERATION(timing)")
grab(accepts, "DRAFT_ACCEPTANCE")
print(f"prompt_tokens={pt} completion_tokens={ct} wall={dt:.1f}s")
PY