#!/usr/bin/env bash
set -euo pipefail

IMAGE=${IMAGE:-voipmonitor/vllm:eldritch-enlightenment-v45c1582-b12xf3686b5-pc1441b5-cu132-20260704}
STANDARD_MODEL=${STANDARD_MODEL:-/root/.cache/huggingface/hub/models--deepseek-ai--DeepSeek-V4-Flash/snapshots/6976c7ff1b30a1b2cb7805021b8ba4684041f136}
DSPARK_MODEL=${DSPARK_MODEL:-/root/.cache/huggingface/hub/models--deepseek-ai--DeepSeek-V4-Flash-DSpark/snapshots/913f0657a874f76844e2e91cbe706dbcaceeb6d7}

NAME=${NAME:-ds4-v9}
PORT=${PORT:-8000}
GPUS=${GPUS:-0,1}
TP=${TP:-2}
DCP=${DCP:-1}
BACKEND=${BACKEND:-b12x-a16}        # b12x variants | lucifer-default | lucifer-cutlass | mainline-cutlass
MODE=${MODE:-dspark}                # standard-mtp0 | standard-mtp2 | standard-mtp3 | dspark
MAX_NUM_SEQS=${MAX_NUM_SEQS:-${SEQ:-64}}
MAX_MODEL_LEN=${MAX_MODEL_LEN:-262144}
MAX_BATCHED=${MAX_BATCHED:-8192}
GPU_MEM_WAS_SET=0
if [[ -n "${GPU_MEM+x}" ]]; then
  GPU_MEM_WAS_SET=1
fi
GPU_MEM=${GPU_MEM:-0.90}
KV_CACHE_MEMORY=${KV_CACHE_MEMORY:-}
GRAPH=${GRAPH:-auto}
PREFIX_CACHE=${PREFIX_CACHE:-1}
DSPARK_TOKENS=${DSPARK_TOKENS:-5}
MTP_TOKENS_WAS_SET=0
if [[ -n "${MTP_TOKENS+x}" ]]; then
  MTP_TOKENS_WAS_SET=1
fi
MTP_TOKENS=${MTP_TOKENS:-2}
SAMPLE=${SAMPLE:-probabilistic}
CACHE=${CACHE:-/root/.cache/vllm-ds4-v9/$NAME}
CONTAINER_TMP=${CONTAINER_TMP:-$CACHE/tmp}
VLLM_INTERNAL_HOST_IP=${VLLM_INTERNAL_HOST_IP:-${VLLM_HOST_IP:-127.0.0.1}}
GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-lo}
NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-lo}
ENABLE_TOPO_PIN=${ENABLE_TOPO_PIN:-0}
AOT_COMPILE=${AOT_COMPILE:-1}
BREAKABLE_CUDAGRAPH=${BREAKABLE_CUDAGRAPH:-0}
PCIE_ALLREDUCE=${PCIE_ALLREDUCE:-1}
DISABLE_CUSTOM_ALLREDUCE=${DISABLE_CUSTOM_ALLREDUCE:-0}
CUDAGRAPH_MODE=${CUDAGRAPH_MODE:-FULL_AND_PIECEWISE}
ENFORCE_EAGER=${ENFORCE_EAGER:-0}
ASYNC_SCHEDULING=${ASYNC_SCHEDULING:-1}
PREFIX_CACHE_RETENTION_INTERVAL=${PREFIX_CACHE_RETENTION_INTERVAL:-4096}
FLASHINFER_DISABLE_VERSION_CHECK=${FLASHINFER_DISABLE_VERSION_CHECK:-0}

mkdir -p \
  "$CACHE/vllm" \
  "$CACHE/tilelang/tmp" \
  "$CACHE/tvm" \
  "$CACHE/triton" \
  "$CACHE/torchinductor" \
  "$CACHE/torch_extensions" \
  "$CACHE/flashinfer" \
  "$CONTAINER_TMP"

case "$MODE" in
  standard-mtp0)
    MODEL=$STANDARD_MODEL
    SERVED_MODEL=DeepSeek-V4-Flash
    SPEC_ARGS=()
    if [[ "$GRAPH" == "auto" ]]; then GRAPH=256; fi
    ;;
  standard-mtp2|standard-mtp3)
    MODEL=$STANDARD_MODEL
    SERVED_MODEL=DeepSeek-V4-Flash
    if [[ "$MODE" == "standard-mtp3" && "$MTP_TOKENS_WAS_SET" == "0" ]]; then
      MTP_TOKENS=3
    fi
    if [[ "$BACKEND" == b12x* ]]; then
      SPEC_JSON=$(printf '{"method":"mtp","num_speculative_tokens":%s,"draft_sample_method":"%s","moe_backend":"b12x"}' "$MTP_TOKENS" "$SAMPLE")
    else
      SPEC_JSON=$(printf '{"method":"mtp","num_speculative_tokens":%s,"draft_sample_method":"%s"}' "$MTP_TOKENS" "$SAMPLE")
    fi
    SPEC_ARGS=(--speculative-config "$SPEC_JSON")
    if [[ "$GRAPH" == "auto" ]]; then GRAPH=512; fi
    ;;
  dspark)
    MODEL=$DSPARK_MODEL
    SERVED_MODEL=DeepSeek-V4-Flash-DSpark
    SPEC_JSON=$(printf '{"model":"%s","method":"dspark","num_speculative_tokens":%s,"draft_sample_method":"%s"}' "$DSPARK_MODEL" "$DSPARK_TOKENS" "$SAMPLE")
    SPEC_ARGS=(--speculative-config "$SPEC_JSON")
    if [[ "$GRAPH" == "auto" ]]; then GRAPH=512; fi
    if [[ "$GPU_MEM_WAS_SET" == "0" ]]; then GPU_MEM=0.93; fi
    ;;
  *)
    echo "Unknown MODE=$MODE" >&2
    exit 2
    ;;
esac

SERVED_MODEL=${SERVED_MODEL_OVERRIDE:-$SERVED_MODEL}

test -f "$MODEL/config.json"

b12x_common_env=(
  -e VLLM_USE_B12X_WO_PROJECTION=1
  -e VLLM_USE_B12X_MHC=1
  -e VLLM_USE_B12X_MOE=1
  -e VLLM_USE_B12X_SPARSE_INDEXER=1
  -e VLLM_ENABLE_PCIE_ALLREDUCE="$PCIE_ALLREDUCE"
  -e VLLM_PCIE_ALLREDUCE_BACKEND=b12x
  -e VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE=64KB
  -e VLLM_PCIE_ONESHOT_ALLREDUCE_MIN_SIZE="${PCIE_ONESHOT_MIN_SIZE:-0}"
  -e VLLM_PCIE_ONESHOT_SINGLE_CHANNEL="${PCIE_ONESHOT_SINGLE_CHANNEL:-1}"
  -e VLLM_DEBUG_B12X_CAPTURE_SEQUENCE="${DEBUG_B12X_CAPTURE_SEQUENCE:-0}"
  -e B12X_MLA_SM120_UNIFIED=1
  -e B12X_MHC_MAX_TOKENS=16384
  -e B12X_DENSE_SPLITK_TURBO=1
  -e B12X_W4A16_TC_DECODE=1
)

case "$BACKEND" in
  b12x|b12x-a16)
    BACKEND_ARGS=(--attention-backend B12X_MLA_SPARSE --moe-backend b12x --linear-backend b12x)
    BACKEND_ENV=(
      "${b12x_common_env[@]}"
      -e VLLM_USE_B12X_FP8_GEMM=1
      -e B12X_MOE_FORCE_A8=0
      -e B12X_MOE_FORCE_A16=1
    )
    ;;
  b12x-a8)
    BACKEND_ARGS=(--attention-backend B12X_MLA_SPARSE --moe-backend b12x --linear-backend b12x)
    BACKEND_ENV=(
      "${b12x_common_env[@]}"
      -e VLLM_USE_B12X_FP8_GEMM=1
      -e B12X_MOE_FORCE_A8=1
      -e B12X_MOE_FORCE_A16=0
    )
    ;;
  b12x-a8-dglin)
    BACKEND_ARGS=(--attention-backend B12X_MLA_SPARSE --moe-backend b12x)
    BACKEND_ENV=(
      "${b12x_common_env[@]}"
      -e VLLM_USE_B12X_FP8_GEMM=0
      -e B12X_MOE_FORCE_A8=1
      -e B12X_MOE_FORCE_A16=0
    )
    ;;
  lucifer-cutlass)
    BACKEND_ARGS=(--attention-backend FLASHINFER_MLA_SPARSE_DSV4 --kernel-config.moe_backend flashinfer_cutlass)
    BACKEND_ENV=(
      -e VLLM_ENABLE_PCIE_ALLREDUCE="$PCIE_ALLREDUCE"
      -e VLLM_PCIE_ALLREDUCE_BACKEND=b12x
      -e VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE=64KB
      -e VLLM_PCIE_ONESHOT_ALLREDUCE_MIN_SIZE="${PCIE_ONESHOT_MIN_SIZE:-0}"
    )
    ;;
  lucifer-default)
    BACKEND_ARGS=(--attention-backend FLASHINFER_MLA_SPARSE_DSV4)
    BACKEND_ENV=(
      -e VLLM_ENABLE_PCIE_ALLREDUCE="$PCIE_ALLREDUCE"
      -e VLLM_PCIE_ALLREDUCE_BACKEND=b12x
      -e VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE=64KB
      -e VLLM_PCIE_ONESHOT_ALLREDUCE_MIN_SIZE="${PCIE_ONESHOT_MIN_SIZE:-0}"
    )
    ;;
  mainline-cutlass)
    BACKEND_ARGS=(--attention-backend FLASHINFER_MLA_SPARSE_DSV4 --kernel-config.moe_backend flashinfer_cutlass)
    BACKEND_ENV=()
    ;;
  *)
    echo "Unknown BACKEND=$BACKEND" >&2
    exit 2
    ;;
esac

PREFIX_ARGS=(--enable-prefix-caching)
if [[ "$PREFIX_CACHE" != "1" ]]; then
  PREFIX_ARGS=(--no-enable-prefix-caching)
fi

EAGER_ARGS=()
if [[ "$ENFORCE_EAGER" == "1" ]]; then
  EAGER_ARGS=(--enforce-eager)
fi

ALLREDUCE_ARGS=()
if [[ "$DISABLE_CUSTOM_ALLREDUCE" == "1" ]]; then
  ALLREDUCE_ARGS=(--disable-custom-all-reduce)
fi

SCHEDULING_ARGS=(--async-scheduling)
if [[ "$ASYNC_SCHEDULING" != "1" ]]; then
  SCHEDULING_ARGS=(--no-async-scheduling)
fi

KV_MEMORY_ARGS=()
if [[ -n "$KV_CACHE_MEMORY" ]]; then
  KV_MEMORY_ARGS=(--kv-cache-memory "$KV_CACHE_MEMORY")
fi

PIN_ARGS=()
PIN_ENV=()
TOPO_CPUSET=${PIN_CPUSET_CPUS:-}
TOPO_MEMSET=${PIN_CPUSET_MEMS:-0}
if [[ "$ENABLE_TOPO_PIN" == "1" || -n "$TOPO_CPUSET" ]]; then
  if [[ -z "$TOPO_CPUSET" ]]; then
    case "$GPUS" in
      0,1) TOPO_CPUSET=0-7,64-71 ;;
      2,3) TOPO_CPUSET=8-15,72-79 ;;
      4,5) TOPO_CPUSET=16-23,80-87 ;;
      6,7) TOPO_CPUSET=24-31,88-95 ;;
      8,9) TOPO_CPUSET=32-39,96-103 ;;
      10,11) TOPO_CPUSET=40-47,104-111 ;;
      12,13) TOPO_CPUSET=48-55,112-119 ;;
      14,15) TOPO_CPUSET=56-63,120-127 ;;
      0,1,2,3) TOPO_CPUSET=0-15,64-79 ;;
      4,5,6,7) TOPO_CPUSET=16-31,80-95 ;;
      8,9,10,11) TOPO_CPUSET=32-47,96-111 ;;
      12,13,14,15) TOPO_CPUSET=48-63,112-127 ;;
      *)
        echo "ENABLE_TOPO_PIN=1 but no CPU mapping for GPUS=$GPUS; docker CPU pin disabled" >&2
        ;;
    esac
  fi
  if [[ -n "$TOPO_CPUSET" ]]; then
    PIN_ARGS=(--cpuset-cpus "$TOPO_CPUSET" --cpuset-mems "$TOPO_MEMSET")
    PIN_ENV=(-e DS4_CPUSET_CPUS="$TOPO_CPUSET" -e DS4_CPUSET_MEMS="$TOPO_MEMSET")
  fi
fi

docker rm -f "$NAME" >/dev/null 2>&1 || true
docker run -d \
  --name "$NAME" \
  --gpus all \
  --runtime nvidia \
  --ipc host \
  --shm-size 32g \
  --network host \
  --init \
  --entrypoint /bin/bash \
  "${PIN_ARGS[@]}" \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --ulimit nofile=1048576:1048576 \
  -v /root/.cache/huggingface:/root/.cache/huggingface:ro \
  -v "$CACHE:/cache:rw" \
  -v "$CONTAINER_TMP:/container-tmp:rw" \
  -e CUDA_VISIBLE_DEVICES="$GPUS" \
  -e CUDA_DEVICE_ORDER=PCI_BUS_ID \
  -e VLLM_HOST_IP="$VLLM_INTERNAL_HOST_IP" \
  -e GLOO_SOCKET_IFNAME="$GLOO_SOCKET_IFNAME" \
  -e NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" \
  -e CUTE_DSL_ARCH=sm_120a \
  -e NCCL_IB_DISABLE=1 \
  -e NCCL_P2P_LEVEL=SYS \
  -e NCCL_PROTO=LL,LL128,Simple \
  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  -e VLLM_PREFIX_CACHE_RETENTION_INTERVAL="$PREFIX_CACHE_RETENTION_INTERVAL" \
  -e VLLM_USE_AOT_COMPILE="$AOT_COMPILE" \
  -e VLLM_USE_MEGA_AOT_ARTIFACT=1 \
  -e VLLM_USE_BREAKABLE_CUDAGRAPH="$BREAKABLE_CUDAGRAPH" \
  -e VLLM_USE_V2_MODEL_RUNNER=1 \
  -e VLLM_USE_FLASHINFER_SAMPLER=1 \
  -e VLLM_MEMORY_PROFILE_INCLUDE_ATTN=1 \
  -e VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1 \
  -e SAFETENSORS_FAST_GPU=1 \
  -e TMPDIR=/container-tmp \
  -e XDG_CACHE_HOME=/cache \
  -e VLLM_CACHE_DIR=/cache/vllm \
  -e TILELANG_CACHE_DIR=/cache/tilelang \
  -e TILELANG_TMP_DIR=/cache/tilelang/tmp \
  -e TVM_CACHE_DIR=/cache/tvm \
  -e TRITON_CACHE_DIR=/cache/triton \
  -e TORCHINDUCTOR_CACHE_DIR=/cache/torchinductor \
  -e TORCH_EXTENSIONS_DIR=/cache/torch_extensions \
  -e FLASHINFER_WORKSPACE_BASE=/cache/flashinfer \
  -e FLASHINFER_DISABLE_VERSION_CHECK="$FLASHINFER_DISABLE_VERSION_CHECK" \
  "${PIN_ENV[@]}" \
  "${BACKEND_ENV[@]}" \
  "$IMAGE" \
  -lc 'unset NCCL_GRAPH_FILE NCCL_GRAPH_DUMP_FILE VLLM_B12X_MLA_EXTEND_MAX_CHUNKS; exec vllm serve "$@"' \
  -- "$MODEL" \
  --served-model-name "$SERVED_MODEL" \
  --host 0.0.0.0 \
  --port "$PORT" \
  --trust-remote-code \
  --kv-cache-dtype fp8 \
  --block-size 256 \
  --load-format auto \
  --tensor-parallel-size "$TP" \
  --decode-context-parallel-size "$DCP" \
  "${ALLREDUCE_ARGS[@]}" \
  --gpu-memory-utilization "$GPU_MEM" \
  "${KV_MEMORY_ARGS[@]}" \
  --max-model-len "$MAX_MODEL_LEN" \
  --max-num-seqs "$MAX_NUM_SEQS" \
  --max-num-batched-tokens "$MAX_BATCHED" \
  --max-cudagraph-capture-size "$GRAPH" \
  --compilation-config "$(printf '{\"cudagraph_mode\":\"%s\",\"custom_ops\":[\"all\"]}' "$CUDAGRAPH_MODE")" \
  "${EAGER_ARGS[@]}" \
  "${SCHEDULING_ARGS[@]}" \
  --no-scheduler-reserve-full-isl \
  --enable-chunked-prefill \
  --enable-flashinfer-autotune \
  --tokenizer-mode deepseek_v4 \
  --tool-call-parser deepseek_v4 \
  --reasoning-parser deepseek_v4 \
  --enable-auto-tool-choice \
  --enable-prompt-tokens-details \
  --enable-force-include-usage \
  --enable-request-id-headers \
  --default-chat-template-kwargs.thinking=true \
  --default-chat-template-kwargs.reasoning_effort=high \
  "${SPEC_ARGS[@]}" \
  "${BACKEND_ARGS[@]}" \
  "${PREFIX_ARGS[@]}"

echo "$NAME $SERVED_MODEL $BACKEND $MODE TP=$TP DCP=$DCP GPUS=$GPUS PORT=$PORT GRAPH=$GRAPH CUDAGRAPH=$CUDAGRAPH_MODE EAGER=$ENFORCE_EAGER MAX_NUM_SEQS=$MAX_NUM_SEQS PCIE_AR=$PCIE_ALLREDUCE CPUSET=${TOPO_CPUSET:-none} MEMSET=${TOPO_MEMSET:-none}"