#!/usr/bin/env bash
# GLM-5.3-Flash W4A16 + DFlash2 speculative decoding — TP2 across two DGX Spark GB10 (SM121) nodes.
# Adapted from tonyd2wild/GLM-5.3-Flash-NVFP4-DFlash2-2x-DGX-Spark/launch-glm53-vllm-tp2-dflash2.sh
# (their 46.9 tok/s config) for a two-Spark RoCE fabric:
# - W4A16 weights (/home/pcozz/models/glm-5.3-w4a16-mtp/) instead of NVFP4
# - fabric: rocep1s0f1 / enp1s0f1np1 / 192.168.102.0/24
# - proven W4A16 serve args (EP on, fp8_e4m3 KV, block 2304, seqs 6)
# - wedge-safety: graceful docker stop -t 30 (NEVER rm -f a GPU-active container), --memory=118g cap
# Usage: launch_dflash2_tp2.sh <0|1> — worker (1) FIRST, wait 25s, then head (0).
# Env knobs (defaults = 262K serve):
# MAX_MODEL_LEN (default 262144; 1048576 = architecture max, needs KV_CACHE_MEM raise)
# KV_CACHE_MEM (default 3221225472 = 3 GiB -> 310,292 tok pool, ~10.38 KB/tok fp8;
# 1M pool needs ~10.9 GiB -> use 11811160064 = 11 GiB -> ~1.09M tok;
# safer 1M value: 9663676416 = 9 GiB -> ~1.36M)
# MODEL_DIR (default /home/pcozz/models/glm-5.3-w4a16-mtp = W4A16)
set -euo pipefail
MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}"
KV_CACHE_MEM="${KV_CACHE_MEM:-3221225472}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-6}"
MODEL_DIR="${MODEL_DIR:-/home/pcozz/models/glm-5.3-w4a16-mtp}"
# DFlash2 drafter lever: override num_speculative_tokens (default 7; any other count boot-wedges).
SPEC_NUM_TOKENS="${SPEC_NUM_TOKENS:-7}"
# EAGER=1 (default) = --enforce-eager; current production passes EAGER=0 GRAPHS=1 (CUDA graphs ON).
EAGER="${EAGER:-1}"
if [ "$EAGER" = "1" ]; then EAGER_FLAG="--enforce-eager"; else EAGER_FLAG=""; fi
# Note: engine may raise block-size to >= mamba page size (see boot log).
BLOCK_SIZE="${BLOCK_SIZE:-2304}"
# Banked default 0.795; current production passes GMU=0.90.
GMU="${GMU:-0.795}"
# GRAPHS=1 + EAGER=0 = FULL_AND_PIECEWISE capture sizes 1-32 (current production).
GRAPH_ARGS=()
if [ "${GRAPHS:-0}" = "1" ]; then GRAPH_ARGS=(--compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,4,8,16,24,32]}'); fi
NODE_RANK="${1:?usage: launch_dflash2_tp2.sh <0|1>}"
[[ "$NODE_RANK" == "0" || "$NODE_RANK" == "1" ]] || { echo "rank must be 0 or 1" >&2; exit 2; }
IMAGE="radixark/vllm-glm53-flash:sm121-v11-dflash2"
NAME="vllm_node"
MODEL_HOST_PATH="$MODEL_DIR"
MODEL_PATH="$MODEL_DIR/"
DRAFTER_HOST_PATH="${DRAFTER_HOST_PATH:-/var/tmp/models/GLM-5.3-Flash-DFlash2}"
CACHE_HOST_PATH="/var/tmp/glm53-vllm-cache"
HEAD_IP="${HEAD_IP:-192.168.102.1}"
MPORT="${MPORT:-29521}"
PORT="8000"
case "$NODE_RANK" in
0) HOST_IP="${HOST_IP_RANK0:-192.168.102.1}"; HEADLESS="" ;;
1) HOST_IP="${HOST_IP_RANK1:-192.168.102.2}"; HEADLESS="--headless" ;;
esac
test -f "$MODEL_HOST_PATH/config.json"
test -f "$DRAFTER_HOST_PATH/config.json"
test -f "$HOME/patches/sparse_attn_indexer_kpool.py"
mkdir -p "$CACHE_HOST_PATH"
# Wedge-safe replace: graceful stop first (never rm -f a GPU-active container — UVM wedge).
if docker ps --format '{{.Names}}' | grep -qx "$NAME"; then
echo "gracefully stopping running $NAME ..."
docker stop -t 30 "$NAME" >/dev/null || true
fi
docker rm "$NAME" 2>/dev/null || true
docker run --gpus all -d \
--name "$NAME" --restart no \
--memory=118g \
--network host --ipc host --shm-size 32g \
--ulimit memlock=-1:-1 --cap-add IPC_LOCK \
--device /dev/infiniband:/dev/infiniband \
-v /home/pcozz/models:/home/pcozz/models \
-v "$CACHE_HOST_PATH:/cache" \
-v "$CACHE_HOST_PATH/flashinfer:/root/.cache/flashinfer" \
-v "$CACHE_HOST_PATH/tilelang:/root/.tilelang" \
-e VLLM_HOST_IP=$HOST_IP \
-e HF_HOME=/cache/huggingface \
-e HF_HUB_OFFLINE=1 -e TRANSFORMERS_OFFLINE=1 \
-e VLLM_ENGINE_READY_TIMEOUT_S=3600 \
-e VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS="${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-3600}" \
-e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
-e TORCH_CUDA_ARCH_LIST=12.1a -e FLASHINFER_CUDA_ARCH_LIST=12.1a \
-e FLASHINFER_DISABLE_VERSION_CHECK=1 \
-e NCCL_NET=IB -e NCCL_IB_DISABLE=0 \
-e NCCL_IB_HCA=rocep1s0f1 -e NCCL_IB_GID_INDEX=3 \
-e NCCL_IB_ROCE_VERSION_NUM=2 -e NCCL_IB_ADDR_FAMILY=AF_INET \
-e NCCL_IB_ADDR_RANGE=192.168.102.0/24 \
-e NCCL_SOCKET_IFNAME=enp1s0f1np1 -e GLOO_SOCKET_IFNAME=enp1s0f1np1 \
-e TP_SOCKET_IFNAME=enp1s0f1np1 -e MN_IF_NAME=enp1s0f1np1 \
-e NCCL_NVLS_ENABLE=0 -e NCCL_CROSS_NIC=0 -e NCCL_IB_MERGE_NICS=0 \
-e NCCL_CUMEM_ENABLE=0 -e NCCL_IGNORE_CPU_AFFINITY=1 -e NCCL_DEBUG=WARN \
-e NCCL_BLOCKING_WAIT=0 \
-e TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \
-e TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=1800 \
-e TORCH_NCCL_DISABLE_WATCHDOG=1 \
-v $HOME/patches/sparse_attn_indexer_kpool.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/sparse_attn_indexer_kpool.py:ro \
-v $DRAFTER_HOST_PATH:/models/dflash2-draft:ro \
"$IMAGE" \
"$MODEL_PATH" \
--served-model-name glm-5.3-flash \
--host 0.0.0.0 --port "$PORT" \
--trust-remote-code \
--tensor-parallel-size 2 \
--enable-expert-parallel \
--gpu-memory-utilization "$GMU" \
--max-model-len "$MAX_MODEL_LEN" \
--max-num-seqs "$MAX_NUM_SEQS" --block-size "$BLOCK_SIZE" \
--speculative-config '{"method":"dflash","model":"/models/dflash2-draft","num_speculative_tokens":'"$SPEC_NUM_TOKENS"'}' \
--kv-cache-dtype fp8_e4m3 --kv-cache-memory "$KV_CACHE_MEM" \
$EAGER_FLAG \
"${GRAPH_ARGS[@]}" \
--tool-call-parser glm47 --enable-auto-tool-choice \
--reasoning-parser glm45 \
--default-chat-template-kwargs '{"enable_thinking":true}' \
--distributed-executor-backend mp \
--nnodes 2 --node-rank "$NODE_RANK" \
--master-addr "$HEAD_IP" --master-port "$MPORT" \
$HEADLESS
echo "launched $NAME rank=$NODE_RANK host=$HOST_IP"
sleep 2
docker ps --format '{{.Names}} {{.Status}}' | grep "$NAME" || {
echo "$NAME exited; inspect with: docker logs $NAME" >&2
exit 1
}