#!/usr/bin/env bash
# bench-pipeline-64k.sh — 4k-context benchmark through a speculative pipeline on
# a single backend.
#
# Default spec chain is draft-dflash + ngram-mod. Pass a second argument to
# override --spec-type (e.g. "ngram-mod" to bench ngram only, or "none" for no
# speculation). The DFlash draft model is only referenced when the spec type
# includes "dflash".
#
# Reads authoritative prefill / generation tok/s and draft acceptance from the
# server's own `slot print_timing` lines.
#
# Usage:
#   ./bench-pipeline-64k.sh <ROCm0|Vulkan0> [spec-type]
set -euo pipefail

DEVICE="${1:?usage: bench-pipeline-64k.sh <ROCm0|Vulkan0> [spec-type]}"
SPEC_TYPE="${2:-draft-dflash,ngram-mod}"
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=8192        # 8k ctx: room for a ~4k prompt + generation
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 "=== 4k pipeline bench: device=$DEVICE spec=$SPEC_TYPE ==="

# Only add the draft-model argument when the spec type references dflash.
DRAFT_ARGS=()
if [[ "$SPEC_TYPE" == *dflash* ]]; then
    DRAFT_ARGS+=(--spec-draft-model "$DRAFT")
fi

"$LLAMA_SERVER" \
    --model "$MODEL" \
    "${DRAFT_ARGS[@]}" \
    --spec-type "$SPEC_TYPE" \
    --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" \
    --ignore-eos \
    >> "$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
crashed=0
for i in $(seq 1 200); do
    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
        # server process died -> backend abort (e.g. Vulkan NONE-op crash)
        crashed=1
        break
    fi
    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
if [ "$crashed" = "1" ]; then
    echo " BACKEND_CRASHED (server died before ready)"
    tail -30 "$LOG"
    exit 2
fi
[ "$ready" = "1" ] || { echo " FAIL (timeout)"; tail -25 "$LOG"; exit 1; }

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

# Build a ~4k filler prompt, then append an incomplete continuation so the
# model produces real generation (pure repetitive filler stops early/EOS).
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 * 120)  # ~16k chars -> ~4k tokens
# Append a clear continuation target that forces real generation.
prompt += ("\n\nNow, in a detailed and flowing paragraph, continue this essay about "
           "the philosophy of language and its influence on modern computation. ")

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())

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