run-glm-optimized.sh

#!/usr/bin/env bash
# SlimServe — optimized GLM-5.2-Vision-GGUF server for MI300X.
#
#   ./run-glm-optimized.sh [--tp N] [--quant NAME] [--ctx N] [--port N]
#
# Defaults target the common case: many short (~2k) requests, with a very large
# ceiling available for the occasional huge one. max_model_len is a ceiling, not
# a per-request reservation — KV is a shared pool, so a big ceiling is free when
# requests are small. The ceiling defaults to 512k on 2 GPUs and 1M on 4+.
set -euo pipefail

MODELS=/home/hotaisle/models
VENV=/home/hotaisle/.venv/bin/python

# Run from a neutral cwd: launching from the repo's parent puts the checkout
# root on sys.path, where `vllm` resolves to a namespace package and fails.
cd /

TP=
DP=1
EP=0
GPUS=
QUANT=Q2_K
CTX=                 # default depends on --tp; see below
PORT=8000
MAX_SEQS=32

while [[ $# -gt 0 ]]; do
    case "$1" in
        --tp)    TP="$2"; shift 2 ;;
        --gpus)  GPUS="$2"; shift 2 ;;
        --dp)    DP="$2"; shift 2 ;;
        --ep)    EP=1; shift ;;
        --quant) QUANT="$2"; shift 2 ;;
        --ctx)   CTX="$2"; shift 2 ;;
        --port)  PORT="$2"; shift 2 ;;
        --max-seqs) MAX_SEQS="$2"; shift 2 ;;
        *) echo "unknown arg: $1" >&2; exit 2 ;;
    esac
done

# --gpus 0,1,2,3 pins the run to those devices. Without it the server takes
# every GPU it can see. TP defaults to the number of devices selected (2 when
# neither flag is given), so `--gpus 0,1,2,3` alone means TP4.
if [[ -n "$GPUS" ]]; then
    export HIP_VISIBLE_DEVICES="$GPUS"
    export CUDA_VISIBLE_DEVICES="$GPUS"
    NUM_GPUS=$(awk -F, '{print NF}' <<<"$GPUS")
    : "${TP:=$NUM_GPUS}"
    if (( TP * DP > NUM_GPUS )); then
        echo "error: tp($TP) x dp($DP) = $((TP*DP)) exceeds the $NUM_GPUS GPU(s)" \
             "given in --gpus $GPUS" >&2
        exit 2
    fi
fi
: "${TP:=2}"

case "$QUANT" in
    Q4_K)   MODEL="$MODELS/GLM-5.2-Vision-GGUF/antirez-routed/GLM-5.2-UD-Q4_K_RoutedQ4K-00001-of-00010.gguf" ;;
    Q2_K)   MODEL="$MODELS/GLM-5.2-Vision-GGUF/antirez-routed/GLM-5.2-UD-Q2_K_RoutedQ2K-00001-of-00006.gguf" ;;
    IQ2_XXS) MODEL="$MODELS/GLM-5.2-Vision-GGUF/antirez-routed/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K-00001-of-00005.gguf" ;;
    *) echo "unknown quant: $QUANT (Q4_K|Q2_K|IQ2_XXS)" >&2; exit 2 ;;
esac
[[ -f "$MODEL" ]] || { echo "missing model: $MODEL" >&2; exit 1; }

# Exact GLM-5.2 DSpark speculator. A bare repo id is resolved through the HF
# cache, so the 5.9 GB draft downloads on first run; a local checkout is
# preferred when present.
DEFAULT_DRAFT_REPO=RedHatAI/GLM-5.2-speculator.dspark
if [[ -d "$MODELS/GLM-5.2-speculator.dspark" ]]; then
    DRAFT="$MODELS/GLM-5.2-speculator.dspark"
else
    DRAFT="$DEFAULT_DRAFT_REPO"
fi

# KV pool, sized explicitly. Auto-sizing (gpu_memory_utilization) misjudges
# this model at large max_model_len: it over-requests and OOMs at init, and
# being generous instead OOMs at execution because the sparse-MLA and indexer
# top-k workspaces also scale with max_model_len. Budget = free VRAM after
# weights, minus workspace headroom. Measured cost/token: 46.6 KB target,
# Draft KV now uses FP8; historical pool sizing needs requalification.
#
# TP2 holds ~129 GiB of weights+state per GPU, leaving ~63 GiB; TP4/TP8 shard
# the weights so far more is free (the MLA latent itself does NOT shard, so
# extra GPUs buy batch and headroom, not per-token cost).
case "$TP" in
    2) KV_GIB=47 ;;
    4) KV_GIB=90 ;;
    8) KV_GIB=120 ;;
    *) KV_GIB=47 ;;
esac
# Expert parallel needs ~10 GiB more per rank than plain TP at the same GPU
# count (measured: tp4 --ep OOMs at the tp4 budget, 187.6 GiB resident before
# profiling), so give the KV pool back that much.
[[ "$EP" == "1" ]] && KV_GIB=$((KV_GIB - 12))
KV_BYTES=$((KV_GIB * 1024 * 1024 * 1024))

# Default context ceiling by GPU count. On 2 GPUs the weights leave only ~63
# GiB per card, and a 1M ceiling needs ~52 GiB of KV *plus* GiB-scale
# sparse-MLA/indexer workspace that also scales with max_model_len — so 1M on
# TP2 cannot fit the required DSpark/FP8 path at a 1M ceiling. From 4
# GPUs up the weights shard and there is ample room, so 1M is the default there.
if [[ -z "$CTX" ]]; then
    if (( TP >= 4 )); then CTX=1048576; else CTX=524288; fi
fi
if (( TP < 4 )) && (( CTX > 524288 )); then
    echo "error: ${CTX}-token ceiling on tp=$TP does not fit the required " \
         "DSpark/FP8 path; use --ctx 524288 or --tp 4." >&2
    exit 2
fi

# num_speculative_tokens=3 measured best at batch scale: mean accept length
# only grows 2.70 -> 3.07 from 3 to 7 draft tokens while verify width doubles,
# so 7 lost ~23% throughput in the historical batch-64 experiment.
# Draft KV uses FP8 under the current serving policy.
SPEC_CFG=$(cat <<JSON
{"model": "$DRAFT", "method": "dspark", "num_speculative_tokens": 3,
 "kv_cache_dtype": "fp8"}
JSON
)

echo "SlimServe: $QUANT  tp=$TP  dp=$DP  ep=$EP  gpus=${GPUS:-all}  ctx=$CTX" \
     "kv=${KV_GIB}GiB  spec=dspark/fp8"
echo "  draft: $DRAFT"

# AITER is required: the target's sparse-MLA indexer has no non-AITER ROCm path.
export VLLM_ROCM_USE_AITER=1

ARGS=(
    --model "$MODEL"
    # No --hf-config-path / --tokenizer: GGUFConfigParser builds the whole
    # Glm5vConfig from glm-dsa.* plus the mmproj's clip.vision.* metadata, and
    # the tokenizer comes from the GGUF's own tokenizer.ggml.* fields.
    --trust-remote-code
    --served-model-name GLM-5.2-Vision
    --tensor-parallel-size "$TP"
    --data-parallel-size "$DP"
    --max-model-len "$CTX"
    --max-num-seqs "$MAX_SEQS"
    # Prefill chunk size. Measured on 2 GPUs at this KV budget, 16384 is the
    # only value that works: 8192 fails torch.compile range setup at startup,
    # 32768 kills the engine mid-request with an AssertionError in the MoE
    # shared-experts runner (its output slot is still occupied on the next
    # forward), and 65536 OOMs on activations. Re-measure before changing.
    --max-num-batched-tokens 16384
    --kv-cache-memory-bytes "$KV_BYTES"
    --block-size 64                     # sparse MLA + sparse SWA require a multiple of 64
    --enable-prefix-caching
    --enable-chunked-prefill
    --limit-mm-per-prompt '{"vision_chunk": 1}'
    # REQUIRED. Short prompts otherwise take sparse MLA's dense forward_mha
    # path, which the ROCm sparse backend does not implement — the engine dies
    # with NotImplementedError on the first small request.
    --attention-config '{"sparse_mla_force_mqa": true}'
    --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}'
    --port "$PORT"
)
ARGS+=(--speculative-config "$SPEC_CFG")
# Expert parallel: shard the 256 routed experts across ranks instead of
# splitting every expert's matrices (which is what plain TP does).
[[ "$EP" == "1" ]] && ARGS+=(--enable-expert-parallel)

exec "$VENV" -m vllm.entrypoints.openai.api_server "${ARGS[@]}"