| 1 | #!/usr/bin/env bash |
| 2 | set -euxo pipefail |
| 3 | |
| 4 | # args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] |
| 5 | THRESHOLD=${1:-0.25} |
| 6 | NUM_Q=${2:-1319} |
| 7 | PORT=${3:-8010} |
| 8 | OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled} |
| 9 | mkdir -p "${OUT_DIR}" |
| 10 | |
| 11 | wait_for_server() { |
| 12 | local port=$1 |
| 13 | timeout "${VLLM_ENGINE_READY_TIMEOUT_S:-600}" bash -c ' |
| 14 | until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do |
| 15 | sleep 1 |
| 16 | done' |
| 17 | } |
| 18 | |
| 19 | MODEL="deepseek-ai/DeepSeek-V2-lite" |
| 20 | |
| 21 | # Set BACKENDS based on platform |
| 22 | if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then |
| 23 | # ROCm platform |
| 24 | BACKENDS=("allgather_reducescatter") |
| 25 | # Disable MOE padding for ROCm since it is causing eplb to fail |
| 26 | export VLLM_ROCM_MOE_PADDING=0 |
| 27 | else |
| 28 | # Non-ROCm platform (CUDA/other) |
| 29 | BACKENDS=("deepep_high_throughput" "deepep_low_latency") |
| 30 | fi |
| 31 | |
| 32 | cleanup() { |
| 33 | if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then |
| 34 | kill "${SERVER_PID}" 2>/dev/null || true |
| 35 | for _ in {1..20}; do |
| 36 | kill -0 "${SERVER_PID}" 2>/dev/null || break |
| 37 | sleep 0.5 |
| 38 | done |
| 39 | kill -9 "${SERVER_PID}" 2>/dev/null || true |
| 40 | fi |
| 41 | } |
| 42 | trap cleanup EXIT |
| 43 | |
| 44 | for BACK in "${BACKENDS[@]}"; do |
| 45 | VLLM_DEEP_GEMM_WARMUP=skip \ |
| 46 | vllm serve "$MODEL" \ |
| 47 | --enforce-eager \ |
| 48 | --tensor-parallel-size 2 \ |
| 49 | --data-parallel-size 2 \ |
| 50 | --enable-expert-parallel \ |
| 51 | --enable-eplb \ |
| 52 | --eplb-config '{"use_async": false}' \ |
| 53 | --trust-remote-code \ |
| 54 | --max-model-len 2048 \ |
| 55 | --all2all-backend "$BACK" \ |
| 56 | --port "$PORT" & |
| 57 | SERVER_PID=$! |
| 58 | wait_for_server "$PORT" |
| 59 | |
| 60 | TAG=$(echo "$MODEL" | tr '/: \\n' '_____') |
| 61 | OUT="${OUT_DIR}/${TAG}_${BACK}.json" |
| 62 | python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}" |
| 63 | python3 - <<PY |
| 64 | import json; acc=json.load(open('${OUT}'))['accuracy'] |
| 65 | print(f"${MODEL} ${BACK}: accuracy {acc:.3f}") |
| 66 | assert acc >= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}" |
| 67 | PY |
| 68 | |
| 69 | cleanup |
| 70 | SERVER_PID= |
| 71 | sleep 1 |
| 72 | PORT=$((PORT+1)) |
| 73 | done |