#!/usr/bin/env bash
#
# Start the FP8 checkpoint on an RTX 3070 8 GB under WSL2.
#
# Defaults were re-validated on 2026-08-07. Two changes matter:
#
#   1. --dtype is now bfloat16, NOT half.
#      With half, the server starts, /health returns 200, the health check
#      reports healthy - and every response is "!!!!!!!!". Dequantizing FP8
#      E4M3 into FP16 overflows the FP16 range, the logits become NaN, and
#      argmax picks token 0. It fails silently. bfloat16 is what config.json
#      declares, Ampere supports it natively, and it costs the same memory.
#
#   2. The context is 8192 instead of 2048, reached by storing the KV cache
#      in fp8 (KV_CACHE_DTYPE). With fp16 KV, vLLM reports that only 5,680
#      tokens fit in the memory this card can spare.
#
# Measured with these defaults:
#     Model loading took 5.51 GiB
#     GPU KV cache size: 9,088 tokens
#     process footprint ~7.8 GiB
#
# Always send a real request after changing anything here - a health check
# cannot detect the NaN failure mode.

set -euo pipefail

MODEL_ID="${MODEL_ID:-hsmin92/internvl35-fp8}"
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-internvl35-fp8}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8000}"

# --- dtype -------------------------------------------------------------------
# Do not set this to half. See the note at the top of this file.
DTYPE="${DTYPE:-bfloat16}"

# --- capacity ----------------------------------------------------------------
MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}"

# 8 is roughly what a 9,088-token cache sustains for typical requests
# (a 960x544 frame plus a normalized crop is about 1,084 prompt tokens).
# vLLM's default of 128 lets more requests in than the cache can hold, which
# triggers preemption and recompute and ends up slower than a smaller limit.
MAX_NUM_SEQS="${MAX_NUM_SEQS:-8}"

# Also sizes the multimodal encoder cache - this is NOT a VRAM knob.
# vLLM's default of 2048 rejects large images outright:
#   "image item with 2816 embedding tokens, which exceeds the
#    pre-allocated encoder cache size 2048"
# InternVL dynamic tiling reaches 3,329 embedding tokens for one image.
MAX_BATCHED_TOKENS="${MAX_BATCHED_TOKENS:-4096}"

MAX_IMAGES="${MAX_IMAGES:-2}"

# --- memory ------------------------------------------------------------------
# This fraction applies to the free memory the PROCESS can see, which is not
# what host nvidia-smi reports. Under WSL2 the gap was about 0.5 GiB, making
# 0.866 the practical ceiling on this card. Startup is refused if the fraction
# exceeds what is actually free.
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.85}"

# Pin the cache instead of letting utilization decide, so startup is
# deterministic. 640M of fp8 KV = 9,088 tokens, enough for one full-length
# 8,192-token request or about 8 typical ones.
KV_CACHE_MEMORY_BYTES="${KV_CACHE_MEMORY_BYTES:-640M}"

# fp8 halves the per-token cost from 144 KiB to 72 KiB. Without it an
# 8,192-token context does not fit on this card.
KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}"

# CUDA graphs need roughly another 0.5-1 GiB that this build does not have.
# Set ENFORCE_EAGER=0 only if you have measured the headroom.
ENFORCE_EAGER="${ENFORCE_EAGER:-1}"

# Required for the server to ACCEPT requests carrying a tool_choice field
# (Open WebUI sends tool_choice: "auto" by default). Actual tool calling still
# does not work - this checkpoint's chat template has no tool rendering.
TOOL_FLAGS="${TOOL_FLAGS:-1}"

if [[ -z "${VIRTUAL_ENV:-}" ]]; then
  echo "ERROR: Activate the vLLM virtual environment first." >&2
  exit 1
fi

if ! command -v vllm >/dev/null 2>&1; then
  echo "ERROR: vllm is not installed in the active environment." >&2
  exit 1
fi

if [[ "$DTYPE" == "half" || "$DTYPE" == "float16" ]]; then
  echo "ERROR: DTYPE=$DTYPE produces NaN logits with this FP8 checkpoint." >&2
  echo "The server would start and pass health checks while answering" >&2
  echo "every request with '!!!!'. Use bfloat16." >&2
  exit 1
fi

if [[ ! -f /usr/lib/wsl/lib/libcuda.so ]]; then
  echo "ERROR: /usr/lib/wsl/lib/libcuda.so was not found." >&2
  echo "This script is intended for NVIDIA GPU passthrough in WSL2." >&2
  exit 1
fi

NVRTC_LIB_DIR="$({ python - <<'PY'
import site
from pathlib import Path

candidates = []
for site_dir in site.getsitepackages():
    nvidia_dir = Path(site_dir) / "nvidia"
    if not nvidia_dir.exists():
        continue
    for lib_dir in nvidia_dir.glob("cu*/lib"):
        if any(lib_dir.glob("libnvrtc-builtins.so*")) and any(lib_dir.glob("libnvrtc.so*")):
            candidates.append(lib_dir)

if not candidates:
    raise SystemExit(1)

print(sorted(candidates)[-1])
PY
} 2>/dev/null)" || {
  echo "ERROR: pip-installed NVRTC libraries were not found." >&2
  exit 1
}

CUDA_DRIVER_LIB_DIR="/usr/lib/wsl/lib"

# WSL2 compatibility settings validated on RTX 3070.
export VLLM_USE_V2_MODEL_RUNNER=0
export VLLM_USE_FLASHINFER_SAMPLER=0

# DeepGEMM targets Hopper; it is not used on Ampere.
export VLLM_USE_DEEP_GEMM=0
export TOKENIZERS_PARALLELISM=false

# Compile-time linker path for -lcuda.
export LIBRARY_PATH="$CUDA_DRIVER_LIB_DIR:${LIBRARY_PATH:-}"

# Runtime paths for the WSL CUDA driver and pip-installed NVRTC.
export LD_LIBRARY_PATH="$CUDA_DRIVER_LIB_DIR:$NVRTC_LIB_DIR:${LD_LIBRARY_PATH:-}"

ARGS=(
  "$MODEL_ID"
  --served-model-name "$SERVED_MODEL_NAME"
  --host "$HOST"
  --port "$PORT"
  --trust-remote-code
  --dtype "$DTYPE"
  --max-model-len "$MAX_MODEL_LEN"
  --max-num-seqs "$MAX_NUM_SEQS"
  --max-num-batched-tokens "$MAX_BATCHED_TOKENS"
  --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION"
  --kv-cache-dtype "$KV_CACHE_DTYPE"
  --limit-mm-per-prompt "{\"image\":$MAX_IMAGES,\"video\":0}"
)

if [[ -n "$KV_CACHE_MEMORY_BYTES" ]]; then
  ARGS+=(--kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES")
fi

if [[ "$ENFORCE_EAGER" == "1" ]]; then
  ARGS+=(--enforce-eager)
fi

if [[ "$TOOL_FLAGS" == "1" ]]; then
  ARGS+=(--enable-auto-tool-choice --tool-call-parser hermes)
fi

echo "Model:            $MODEL_ID"
echo "NVRTC:            $NVRTC_LIB_DIR"
echo "dtype:            $DTYPE"
echo "Context:          $MAX_MODEL_LEN"
echo "Max sequences:    $MAX_NUM_SEQS"
echo "Batched tokens:   $MAX_BATCHED_TOKENS"
echo "GPU utilization:  $GPU_MEMORY_UTILIZATION"
echo "KV cache:         ${KV_CACHE_MEMORY_BYTES:-automatic} ($KV_CACHE_DTYPE)"

exec vllm serve "${ARGS[@]}"