#!/usr/bin/env bash
# Serve gemma4-e4b-w4a16-mtp-vLLM: gemma-4-E4B-it, GPTQ W4A16 (full-precision
# lm_head), paired with the Gemma-4 MTP drafter for speculative decoding.
# Measured ~130 tok/s single-stream decode on one RTX 4060 Ti (16 GB), 131k ctx,
# correct tool calls. See ../README.md for the full writeup.
#
# Usage:
#   MODEL_DIR=/path/to/checkpoint DRAFTER_DIR=/path/to/drafter ./launch.sh
#
# Requires: vLLM 0.26.0+ built against CUDA (see requirements.txt), a CUDA GPU
# with >=12 GB free VRAM, and both directories from this repo downloaded locally
# (vLLM does not fetch a separate --speculative-config model over the Hub).
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

MODEL_DIR="${MODEL_DIR:-$HERE/../checkpoint}"
DRAFTER_DIR="${DRAFTER_DIR:-$HERE/../drafter}"

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
# vLLM's gloo process group otherwise binds whatever the default route
# interface is (a Tailscale/VPN addr, a stale docker bridge, ...) and hangs
# or dies on startup. Loopback is always correct for a single-GPU server.
export GLOO_SOCKET_IFNAME=lo
export VLLM_HOST_IP=127.0.0.1

exec vllm serve "$MODEL_DIR" \
  --served-model-name gemma-4-e4b-w4a16-mtp \
  --host "${HOST:-0.0.0.0}" --port "${PORT:-8000}" \
  --max-model-len "${MAX_MODEL_LEN:-131072}" \
  --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.90}" \
  --max-num-seqs "${MAX_NUM_SEQS:-1}" \
  --speculative-config "{\"method\":\"mtp\",\"model\":\"$DRAFTER_DIR\",\"num_speculative_tokens\":${NUM_SPECULATIVE_TOKENS:-7}}" \
  --enable-auto-tool-choice --tool-call-parser gemma4 \
  --reasoning-parser gemma4