#!/usr/bin/env bash
# bench-spec-stack.sh — compare MTP-only vs MTP+ngram-mod spec on the v3
# CIRU prompt-forge stack (Dankclaw, gfx1151), at 8k / 16k / 64k context.
#
# Reads authoritative prefill/generation tok/s from llama-server's own
# `print_timing` lines (prompt eval time / eval time). Runs one server
# instance per (spec, ctx) on a scratch port.
#
# Usage:
#   ./scripts/bench-spec-stack.sh <draft-mtp|draft-mtp,ngram-mod> <ctx> [out_tokens]
# Example:
#   ./scripts/bench-spec-stack.sh draft-mtp,ngram-mod 16384 256
set -euo pipefail

SPEC="${1:?usage: bench-spec-stack.sh <spec> <ctx> [out_tokens]}"
CTX="${2:?usage: bench-spec-stack.sh <spec> <ctx> [out_tokens]}"
OUT_TOKENS="${3:-256}"

ENV_FILE=/home/dank/vllmScripts/scripts/frank-default-large.env
source "$ENV_FILE"

# Derive the PromptForge sidecar paths from the env base values.
export PROMPTFORGE_SIDECAR="$MODEL_DIR/$FFN_PFS"
export PROMPTFORGE_GDN_SIDECAR="$MODEL_DIR/$GDN_PFS"
export PROMPTFORGE_MTP_OUTPUT_K8_PROXY="$MODEL_DIR/$OUTPUT_K8_PFS"

MODEL="$MODEL_DIR/$MODEL_FILE"
LLAMA_BIN=/home/dank/ciru-qwen38/ROCmFPX/build-promptforge/bin/llama-server
PORT=7893
PARALLEL=1
BATCH=4096
UBATCH=4096
LOG="/tmp/specstack-$(echo "$SPEC" | tr ',' '-')-ctx$CTX.log"
: > "$LOG"

for f in "$MODEL" "$PROMPTFORGE_SIDECAR" "$PROMPTFORGE_GDN_SIDECAR" "$PROMPTFORGE_MTP_OUTPUT_K8_PROXY" "$MODEL_DIR/$MMPROJ_FILE"; do
    [[ -f "$f" ]] || { echo "missing: $f" >&2; exit 1; }
done
[[ -x "$LLAMA_BIN" ]] || { echo "missing: $LLAMA_BIN" >&2; exit 1; }

export PROMPTFORGE_MODE="m2048_fused_tail1476"
export GGML_CUDA_GRAPH_OPT=0
export PROMPTFORGE_ENABLE_SMALLM_W8=1
export PROMPTFORGE_MTP_OUTPUT_K8=1
export PROMPTFORGE_MTP_OUTPUT_K8_VALIDATE=0
export PROMPTFORGE_MTP_OUTPUT_K8_PROXY
unset LLAMA_MTP_CPU_ARGMAX_FASTPATH LLAMA_MTP_BACKEND_GREEDY || true
unset PROMPTFORGE_IU4_SIDECAR PROMPTFORGE_GDN_IU4_SIDECAR || true
unset PROMPTFORGE_ENABLE_NGRAM_MOD PROMPTFORGE_ENABLE_NGRAM_M65_IU4 || true

export ROCM_PATH="$ROCM_DEVEL"
export LD_LIBRARY_PATH="${ROCM_LIB}/lib:${ROCM_CORE}/lib:${ROCM_DEVEL}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"

echo "=== spec-stack bench: spec=$SPEC ctx=$CTX model=$MODEL_FILE ==="
echo "=== env: OUTPUT_K8=$PROMPTFORGE_MTP_OUTPUT_K8 smallm=$PROMPTFORGE_ENABLE_SMALLM_W8 ==="

# ngram-mod tuning matches the retired fallback line: n_min 48 / n_max 64 / n_match 24
"$LLAMA_BIN" \
    --model "$MODEL" \
    --mmproj "$MODEL_DIR/$MMPROJ_FILE" \
    --alias bench \
    --host 127.0.0.1 \
    --port "$PORT" \
    -dev ROCm0 \
    -ngl 999 \
    -c "$CTX" \
    -b "$BATCH" \
    -ub "$UBATCH" \
    -fa on \
    -ctk "$CACHE_TYPE_K" \
    -ctv "$CACHE_TYPE_V" \
    -np "$PARALLEL" \
    -ctxcp "$CTX_CHECKPOINTS" \
    --cache-ram "$CACHE_RAM" \
    --cache-prompt \
    --no-mmap \
    --spec-type "$SPEC" \
    --spec-draft-device ROCm0 \
    --spec-draft-ngl 999 \
    --spec-draft-type-k f16 \
    --spec-draft-type-v f16 \
    --spec-draft-n-max 4 \
    --spec-draft-n-min 0 \
    --spec-draft-p-min 0.0 \
    --spec-draft-p-split 0.10 \
    --spec-draft-backend-sampling \
    --spec-ngram-mod-n-min 48 \
    --spec-ngram-mod-n-max 64 \
    --spec-ngram-mod-n-match 24 \
    --temp 0.7 \
    --top-p 0.95 \
    -lv 4 \
    >> "$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 900); do
    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
        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
    [ $((i % 10)) = 0 ] && 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; }

# Print the implementations that initialized (feasibility evidence)
grep -E "adding speculative implementation|ngram_mod|draft-mtp" "$LOG" | head -6

python3 - "$SPEC" "$CTX" "$PORT" "$OUT_TOKENS" "$LOG" <<'PY'
import sys, time, json, urllib.request
spec, ctx, port, out_tokens, log = sys.argv[1], int(sys.argv[2]), sys.argv[3], int(sys.argv[4]), sys.argv[5]

# Flowing prose (varied) at ~75% of ctx so MTP acceptance is realistic.
VAR = (
    "narrative nonfiction blends observation with reflection, ",
    "and a careful writer lets each clause earn its place ",
    "in the sentence's rising line, so meaning accumulates slowly. ",
    "The reader is not hurried; the reader is led. Loose and periodic, ",
    "the style turns on concrete nouns and verbs that carry weight, ",
    "and it resists the abstract noun that hides a vague claim. ",
    "A paragraph is not a list. It is a single turn of thought that ",
    "opens, develops, and lands. This is the craft of prose. Choose, ",
    "arrange, and above all omit. Omission is the discipline that ",
    "keeps the whole clear, and clarity is the courtesy we owe the ",
    "willing reader."
)
n = len(VAR)
parts = []
t = 0
while t * 76 < (ctx * 3) // 4:
    for k in range(5):
        parts.append(VAR[(t + k) % n])
    parts.append(f" This is cued by note {t}. ")
    t += 1
prompt = "".join(parts)
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=7200) 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)

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]

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

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