#!/usr/bin/env bash
# Serve Qwen3.8-27B on the 4090 via llama.cpp, with built-in MTP speculative decoding.
set -euo pipefail

LLAMA_DIR="${LLAMA_DIR:-$HOME/code/llama.cpp}"
MODEL_DIR="${MODEL_DIR:-$HOME/models/qwen3.8-27b}"
SERVER="$LLAMA_DIR/build/bin/llama-server"

MODEL="$MODEL_DIR/Qwen3.8-27B-UD-Q4_K_XL.gguf"
MMPROJ="$MODEL_DIR/mmproj-F16.gguf"
DSPARK="$MODEL_DIR/Qwen3.8-27B-DSpark-Q8_0.gguf"

# 131072 measured at 22648 MiB of 24564, leaving ~1.9 GB. 163840 is the ceiling
# with MTP (23544 MiB) but the headroom gets thin; 196608 OOMs. Override with CTX=.
CTX="${CTX:-131072}"
PORT="${PORT:-8080}"
HOST="${HOST:-0.0.0.0}"
NDRAFT="${NDRAFT:-5}"
# q4_0 KV is what buys 128K. Measured 3/3 needle recall at 120K tokens, full GPU
# prompt speed, and slightly *faster* decode than q8_0 (smaller cache = less
# bandwidth per step). Not lossless, but not degraded either.
#
# Only f16, q8_0 and q4_0 have CUDA flash-attention kernels. q5_1 loads and looks
# fine on paper, then silently falls back to CPU: 18 tok/s prompt processing
# against 2185 for q8_0. Do not use it.
KVTYPE="${KVTYPE:-q4_0}"

# Thinking is on by default in llama.cpp with an *unrestricted* budget, which is
# how a 512-token reply ends up entirely inside reasoning_content with nothing
# left for an answer. 8192 is room to think without letting it run away.
# 0 disables thinking; -1 restores the unrestricted default.
REASONING_BUDGET="${REASONING_BUDGET:-8192}"
# Injected right before the end-of-thinking tag when the budget runs out, so the
# model wraps up instead of being cut off mid-thought. Unset to leave it empty.
REASONING_BUDGET_MESSAGE="${REASONING_BUDGET_MESSAGE:-Thinking budget reached. Give the best answer you have now.}"

# Qwen3-recommended sampling. Several of these differ from the llama.cpp defaults
# (temp 0.80, top_k 40, min_p 0.05); the rest are pinned so they cannot drift
# with an upstream change. These are only *defaults* — anything the client sends
# in the request wins, which is why the benchmarks here still get temperature 0.
#
# 0.7 for instruction following. Qwen documents ~0.6 for thinking-mode work and
# 1.0 for open-ended generation, so raise it with TEMP= for creative tasks.
TEMP="${TEMP:-0.7}"
TOP_P="${TOP_P:-0.95}"
TOP_K="${TOP_K:-20}"
MIN_P="${MIN_P:-0.0}"
PRESENCE_PENALTY="${PRESENCE_PENALTY:-0.0}"
REPEAT_PENALTY="${REPEAT_PENALTY:-1.0}"
# K and V can differ. K is the more quantization-sensitive of the two, so a q8_0
# K with q4_0 V is the natural middle setting — set KTYPE/VTYPE to override.
KTYPE="${KTYPE:-$KVTYPE}"
VTYPE="${VTYPE:-$KVTYPE}"
# llama-server defaults to 4 slots. They share one unified KV cache, but each
# still costs compute buffers, which is pure waste for a single-user box.
PARALLEL="${PARALLEL:-1}"

# Speculative mode: mtp (default, free) | dspark (sidecar drafter) | plain (off)
MODE="mtp"
case "${1:-}" in
    --plain)  MODE="plain"  ;;
    --dspark) MODE="dspark" ;;
    --mtp|"") MODE="mtp"    ;;
    *) echo "usage: $0 [--mtp|--dspark|--plain]" >&2; exit 2 ;;
esac

[[ -x "$SERVER" ]] || { echo "no llama-server at $SERVER — run setup/01-llama-cpp.sh" >&2; exit 1; }
[[ -f "$MODEL" ]]  || { echo "no model at $MODEL — run setup/02-models.sh" >&2; exit 1; }

args=(
    --model "$MODEL"
    --n-gpu-layers 99
    --ctx-size "$CTX"
    --flash-attn on
    --cache-type-k "$KTYPE"
    --cache-type-v "$VTYPE"
    --parallel "$PARALLEL"
    --host "$HOST"
    --port "$PORT"
    --jinja
    --alias qwen3.8-27b
    --reasoning-budget "$REASONING_BUDGET"
    --temp "$TEMP"
    --top-p "$TOP_P"
    --top-k "$TOP_K"
    --min-p "$MIN_P"
    --presence-penalty "$PRESENCE_PENALTY"
    --repeat-penalty "$REPEAT_PENALTY"
)

# Only meaningful when the budget can actually be exhausted.
if [[ -n "$REASONING_BUDGET_MESSAGE" && "$REASONING_BUDGET" != "-1" && "$REASONING_BUDGET" != "0" ]]; then
    args+=(--reasoning-budget-message "$REASONING_BUDGET_MESSAGE")
fi

# Vision tower is optional; skip with NO_VISION=1 to save ~0.9 GB.
if [[ -f "$MMPROJ" && "${NO_VISION:-0}" != "1" ]]; then
    args+=(--mmproj "$MMPROJ")
fi

case "$MODE" in
    mtp)
        # The MTP head is embedded in the main GGUF as blk.64.nextn.*, so this
        # needs no drafter file. --spec-type draft-mtp sets load_mtp on the main
        # model params; the draft context reuses the target's weights.
        args+=(--spec-type draft-mtp --spec-draft-n-max "$NDRAFT")
        ;;
    dspark)
        [[ -f "$DSPARK" ]] || { echo "no DSpark drafter at $DSPARK" >&2; exit 1; }
        args+=(--spec-draft-model "$DSPARK" --spec-type draft-dspark --spec-draft-n-max "$NDRAFT")
        ;;
    plain) ;;
esac

echo "mode=$MODE ctx=$CTX k=$KTYPE v=$VTYPE parallel=$PARALLEL draft-n=$NDRAFT think=$REASONING_BUDGET"
echo "sampling: temp=$TEMP top_p=$TOP_P top_k=$TOP_K min_p=$MIN_P presence=$PRESENCE_PENALTY repeat=$REPEAT_PENALTY"
exec "$SERVER" "${args[@]}"