#!/usr/bin/env bash
# GLM-5.2 hybrid NVFP4+AQLM on 4x RTX PRO 6000 Blackwell (SM120, 96GB).
# Usage: ./run_rtx6000.sh [1m|500k|250k]   (default: 1m)
#
# One-time setup (fresh box) — see SETUP steps below the config block.
set -euo pipefail

VARIANT="${1:-1m}"
HF_BASE="jarrelscy/GLM-5.2-NVFP4-AQLM-hybrid"   # weights + code/ bundle
VLLM_FORK="https://github.com/jarrelscy/vllm-glm52-sm120"  # full fork, patches committed
VLLM_BRANCH="glm52-sm120"
WORK="${WORK:-$HOME/glm52}"

case "$VARIANT" in
  1m)   HF_REPO="$HF_BASE";      MAXLEN=1048576 ;;  # 30% hot experts, 292 GB
  500k) HF_REPO="$HF_BASE-500k"; MAXLEN=524288  ;;  # 48% hot experts, 334 GB
  250k) HF_REPO="$HF_BASE-250k"; MAXLEN=262144  ;;  # 57% hot experts, 354 GB
  *) echo "usage: $0 [1m|500k|250k]"; exit 1 ;;
esac
MODEL_DIR="$WORK/models/$VARIANT"

# ---------- one-time setup ----------
if [ ! -d "$WORK/vllm/.venv" ]; then
  mkdir -p "$WORK" && cd "$WORK"
  command -v uv >/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
  export PATH="$HOME/.local/bin:$PATH"

  # vLLM changes now live in the fork (no patches to apply)
  git clone --depth 1 -b "$VLLM_BRANCH" "$VLLM_FORK" vllm && cd vllm
  uv venv --python 3.12 .venv && source .venv/bin/activate
  uv pip install vllm --torch-backend=auto   # bootstrap deps fast
  VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto
  uv pip install ninja "huggingface_hub[cli]"

  # flashinfer git main: SM120 sparse-MLA needs kv_scale_format=arbitrary_fp32
  uv pip install "flashinfer-python @ git+https://github.com/flashinfer-ai/flashinfer.git"

  # JIT toolchain: pin nvcc/crt/nvvm/cccl to torch's CUDA minor, lib64 link
  TORCH_CU=$(python -c "import torch;print(torch.version.cuda.split('.')[0]+'.'+torch.version.cuda.split('.')[1])")
  uv pip install "nvidia-cuda-nvcc==${TORCH_CU}.*" "nvidia-cuda-crt==${TORCH_CU}.*" \
                 "nvidia-nvvm==${TORCH_CU}.*" "nvidia-cuda-cccl==${TORCH_CU}.*" || true
  CU13="$WORK/vllm/.venv/lib/python3.12/site-packages/nvidia/cu13"
  ln -sfn lib "$CU13/lib64" 2>/dev/null || true
  for f in "$CU13"/lib/lib*.so.*; do b=$(basename "$f"); s=${b%%.so.*};
    [ -e "$CU13/lib/$s.so" ] || ln -s "$b" "$CU13/lib/$s.so"; done
fi

# ---------- fetch weights ----------
if [ ! -f "$MODEL_DIR/config.json" ]; then
  source "$WORK/vllm/.venv/bin/activate"
  hf download "$HF_REPO" --exclude "code/*" --local-dir "$MODEL_DIR"
fi

# ---------- serve ----------
cd "$WORK/vllm" && source .venv/bin/activate
export CUDA_HOME="$WORK/vllm/.venv/lib/python3.12/site-packages/nvidia/cu13"
export VLLM_PP_LAYER_PARTITION="21,19,19,19"
export NCCL_MAX_NCHANNELS=4 NCCL_BUFFSIZE=1048576
export VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=256

exec vllm serve "$MODEL_DIR" \
  --pipeline-parallel-size 4 \
  --gpu-memory-utilization 0.95 \
  --kv-cache-dtype fp8_ds_mla \
  --max-model-len "$MAXLEN" \
  --max-num-seqs 2 \
  --max-num-batched-tokens 2048 \
  --enforce-eager \
  --port 8000