#!/bin/bash
# Offload-free 512k lane boot. Container glm-512k-fit, port 18996.
# Usage: boot-512k-fit.sh [MAXLEN] [IMAGE] [EXTRA vllm args...]
# Env knobs: MODEL_DIR (checkpoint path), CACHE_DIR (vllm cache),
#            SEQS (2), BATCHED (2048), UTIL (0.98), MM ("0" -> tower skipped)
MODEL_DIR="${MODEL_DIR:-$PWD/GLM-5.3-Flash-NVFP4-FP8ATTN-512K}"
CACHE_DIR="${CACHE_DIR:-$HOME/.cache/vllm-glm53}"
set -e
MAXLEN="${1:-524288}"
IMAGE="${2:-local/vllm-glm53:fp8attn-r4}"
shift 2 2>/dev/null || shift $# 2>/dev/null || true
MM_ARGS=()
if [ "${MM:-0}" = "0" ]; then
  MM_ARGS=(--limit-mm-per-prompt '{"image":0,"video":0}')
fi
SPEC_ARGS=(--speculative-config '{"method":"mtp","num_speculative_tokens":1}')
if [ "${SPEC:-1}" = "0" ]; then
  SPEC_ARGS=()
fi
exec docker run --name glm-512k-fit --init --rm \
  --gpus all \
  --runtime nvidia \
  --ipc=host \
  --network host \
  --shm-size=32g \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -e VLLM_ENGINE_READY_TIMEOUT_S=3600 \
  -e HF_HUB_OFFLINE=1 \
  -e VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS="${ESTG:-1}" \
  -v "$MODEL_DIR":/model:ro \
  -v "$CACHE_DIR":/root/.cache \
  "$IMAGE" \
  /model \
  --served-model-name glm-5.3-flash-fp8attn \
  --host 0.0.0.0 \
  --port 18996 \
  --tensor-parallel-size 2 \
  --max-model-len "$MAXLEN" \
  --gpu-memory-utilization "${UTIL:-0.98}" \
  --kv-cache-dtype fp8 \
  --max-num-seqs "${SEQS:-2}" \
  --max-num-batched-tokens "${BATCHED:-2048}" \
  --enable-prefix-caching \
  --no-enable-flashinfer-autotune \
  --enable-auto-tool-choice \
  --tool-call-parser glm47 \
  --reasoning-parser glm45 \
  --trust-remote-code \
  "${SPEC_ARGS[@]}" \
  "${MM_ARGS[@]}" \
  "$@"