#!/usr/bin/env bash
# Serve letechlead/MiMo-V2.6-Distill-Qwen-9B-INT8-W8A16-AutoRound as
# frank-small on Dankclaw's RTX 3090 via SGLang.
#
# Topology: the whole checkpoint (12.0 GiB, INT8 W8A16 AutoRound) lives in
# VRAM; the 3090's 24 GiB holds weights + a ~11.8 GiB fp8 KV pool at
# --mem-fraction-static 0.98. The host CPU only tokenizes.
#
# Architecture: qwen3_5 hybrid — 32 layers, 8 full-attention + 24 Gated
# DeltaNet, native 262144 window, and a vision tower (qwen3_5_vision)
# baked into the safetensors: no separate projector file, the image path
# is in the checkpoint. The config declares an MTP layer, but the shipped
# weights carry none (no mtp.* tensors in the index), so this service runs
# WITHOUT speculative decoding — a NEXTN draft would be an uninitialised
# 4.24 GiB ghost model drafting at 0% acceptance.
#
# Context: CONTEXT_LENGTH x MAX_RUNNING slots must fit the fp8 KV pool,
# which the engine autofits around the weights AND the DeltaNet ssm state
# pool (see frank-small.env). The mamba pool defaults to ~5.1 GB, so it is
# pinned (bf16, 32 slots). With the grid at its measured fit the card has
# no headroom for prefill CUDA-graph capture, so prefill runs eager
# (--disable-prefill-cuda-graph); decode graphs are untouched.
#
# Sampling: greedy by default (the fleet contract); clients that send their
# own temperature get their own values.
#
# Tunables come from frank-small.env unless FRANK_SMALL_ENV_FILE overrides it.
set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${FRANK_SMALL_ENV_FILE:-$SCRIPT_DIR/frank-small.env}"
if [[ -f "$ENV_FILE" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "$ENV_FILE"
    set +a
fi

# --- defaults (used when the env file does not set them) -------------------
: "${SGLANG_BIN:=/home/dank/GitPackages/sglang-cu128/bin/sglang}"
: "${MODEL_DIR:=/home/dank/models/MiMo-V2.6-Distill-Qwen-9B-INT8-W8A16-AutoRound}"
: "${ALIAS:=frank-small}"
: "${BIND_ADDR:=0.0.0.0}"
: "${LISTEN_PORT:=7771}"
: "${MEM_FRACTION:=0.98}"
: "${CONTEXT_LENGTH:=151552}"
: "${MAX_TOTAL_TOKENS:=606208}"
: "${MAX_RUNNING:=4}"
: "${KV_DTYPE:=fp8_e4m3}"
: "${QUANTIZATION:=auto-round}"
: "${REASONING_PARSER:=qwen3}"
: "${TOOLCALL_PARSER:=qwen3_coder}"
: "${MAMBA_SSM_DTYPE:=bfloat16}"
: "${MAX_MAMBA_CACHE_SIZE:=32}"

# Fail loudly rather than half-loading a model into VRAM: a missing file here
# is a download that has not finished, and the unit's restart policy would
# otherwise turn that into a crash loop that reads like a runtime fault.
for path in "$SGLANG_BIN" "$MODEL_DIR/model.safetensors.index.json" \
    "$(cd -- "$(dirname -- "$SGLANG_BIN")/.." && pwd)/cuda-nvcc/usr/local/cuda-13.0/bin/nvcc"; do
    [[ -e "$path" ]] || { echo "frank-small: missing $path" >&2; exit 1; }
done
shard_count=$(find "$MODEL_DIR" -maxdepth 1 -name 'model-*.safetensors' | wc -l)
if [[ "$shard_count" -ne 8 ]]; then
    echo "frank-small: expected 8 safetensors shards in $MODEL_DIR, found $shard_count" >&2
    exit 1
fi

# The UVM device node is created lazily by nvidia-modprobe on the first CUDA
# open after boot; make sure it exists before the server needs it.
nvidia-modprobe -u -c=0 2>/dev/null || true

# FlashInfer JITs its attention kernels with nvcc at boot. The host's
# /usr/bin/nvcc is CUDA 12.4 and dies on `--compress-mode` (a 12.8+
# option), so the venv carries a matching CUDA 13.0 nvcc (extracted from
# the cuda-nvcc-13-0 deb); it must shadow the system one on PATH.
export PATH="$(cd -- "$(dirname -- "$SGLANG_BIN")/.." && pwd)/cuda-nvcc/usr/local/cuda-13.0/bin:$PATH"

# Sampling: SGLang has no server-wide sampling defaults; the harness sends
# its per-model recipe from models.yml (compat.extraBody), and other clients
# get the engine's own defaults.
exec "$SGLANG_BIN" serve \
    --model-path "$MODEL_DIR" \
    --served-model-name "$ALIAS" \
    --host "$BIND_ADDR" \
    --port "$LISTEN_PORT" \
    --mem-fraction-static "$MEM_FRACTION" \
    --context-length "$CONTEXT_LENGTH" \
    --max-running-requests "$MAX_RUNNING" \
    --kv-cache-dtype "$KV_DTYPE" \
    --quantization "$QUANTIZATION" \
    --reasoning-parser "$REASONING_PARSER" \
    --tool-call-parser "$TOOLCALL_PARSER" \
    --mamba-ssm-dtype "$MAMBA_SSM_DTYPE" \
    --max-mamba-cache-size "$MAX_MAMBA_CACHE_SIZE" \
    --disable-prefill-cuda-graph \
    --max-total-tokens "$MAX_TOTAL_TOKENS"