#!/usr/bin/env bash
# One benchmark configuration: start vLLM serve, run the single-stream client, stop the server.
# usage: run_mtp_serve.sh MODEL_VIEW TAG K [DRAFT_VOCAB]
M=$1; TAG=$2; K=$3; V=${4:-}
cd /root/q38prune
mkdir -p results/mtp
while [ "$(awk '$1=="shmem"{print $2}' /sys/fs/cgroup/memory.stat)" -gt 2000000000 ]; do sleep 10; done
SPEC=()
if [ "$K" -gt 0 ]; then
  if [ -n "$V" ]; then
    export Q38_DRAFT_VOCAB=$V Q38_DRAFT_HEAD_Q=marlin Q38_DRAFT_MTP_Q=marlin
    SPEC=(--speculative-config "{\"method\":\"mtp\",\"num_speculative_tokens\":$K,\"use_local_argmax_reduction\":true}")
  else
    unset Q38_DRAFT_VOCAB
    SPEC=(--speculative-config "{\"method\":\"mtp\",\"num_speculative_tokens\":$K}")
  fi
fi
./serve_vllm.sh "$M" bench 8000 --kv-cache-dtype auto --max-model-len 16384 --max-num-seqs 8 --block-size 48 \
  --no-enable-prefix-caching --limit-mm-per-prompt '{"image":0,"video":0}' "${SPEC[@]}" > logs/mtp_serve_$TAG.log 2>&1 &
SP=$!
for i in $(seq 1 120); do curl -s localhost:8000/v1/models > /dev/null && break; sleep 10; done
/root/wdm/envs/jev/bin/python draft/bench_mtp_client.py --model bench --out results/mtp/$TAG.json > logs/mtp_client_$TAG.log 2>&1
grep RESULT logs/mtp_client_$TAG.log | cut -c1-600
grep -o "\[q38-draft\].*" logs/mtp_serve_$TAG.log | head -1 | cut -c1-400
kill $(pgrep -f "served-model-name benc[h]") 2>/dev/null
wait $SP 2>/dev/null
sleep 15