#!/usr/bin/env bash
# One Flint replica on one Hopper H200 GPU (141 GB). FA4 and the trtllm MoE kernels are
# Blackwell-only, so this uses FA3 attention and Triton MoE. 80 GB H100s cannot hold the BF16
# weights on one GPU.
#
#   MODEL=/path/to/flint CUDA_VISIBLE_DEVICES=0 PORT=30000 bash serve_h200.sh
#
# Extra arguments are passed through to sglang.launch_server.
set -euo pipefail

MODEL="${MODEL:-RecursiveSI/flint}"
PORT="${PORT:-30000}"
revision_args=()
if [[ -n "${REVISION:-}" ]]; then
  revision_args=(--revision "$REVISION")
fi

export SGLANG_EXTERNAL_MODEL_PACKAGE="${SGLANG_EXTERNAL_MODEL_PACKAGE:-flint_sglang}"
export FLASHINFER_USE_CUDA_NORM="${FLASHINFER_USE_CUDA_NORM:-1}"
# One Triton cache per replica: replicas compiling the same kernels into one cache directory can race.
export TRITON_CACHE_DIR="${FLINT_TRITON_CACHE_BASE:-${HOME:-/tmp}/.triton/cache}/flint-$(hostname)-${PORT}"

exec python3 -m sglang.launch_server \
  --model-path "$MODEL" "${revision_args[@]}" \
  --served-model-name flint \
  --host 0.0.0.0 --port "$PORT" \
  --tp-size 1 \
  --attention-backend fa3 \
  --moe-runner-backend triton \
  --enable-fp32-lm-head \
  --context-length 262144 \
  --chunked-prefill-size 4096 \
  --swa-full-tokens-ratio 0.4 \
  --mem-fraction-static 0.9 \
  --tool-call-parser flint \
  --reasoning-parser flint \
  --grammar-backend flint \
  "$@"