| 1 | #!/bin/bash |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | # |
| 5 | # Accuracy test driver for the p2p connector |
| 6 | # (OffloadingConnector + TieringOffloadingSpec + p2p tier). |
| 7 | # |
| 8 | # Mirrors tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh: |
| 9 | # brings up N prefillers + M decoders on the local host, fronts them |
| 10 | # with p2p_connector_proxy.py, then runs the connector-agnostic |
| 11 | # test_accuracy.py (gsm8k via lm_eval) against the proxy. |
| 12 | # |
| 13 | # Knobs (env vars unless flagged otherwise): |
| 14 | # MODEL_NAMES space-separated model list (default: Llama-3.2-1B-Instruct) |
| 15 | # NUM_PREFILL_INSTANCES default 1 |
| 16 | # NUM_DECODE_INSTANCES default 1 |
| 17 | # PREFILLER_TP_SIZE default 1 |
| 18 | # DECODER_TP_SIZE default 1 |
| 19 | # DP_EP when set, deploy the decoder in DP-EP attention |
| 20 | # mode (dp=DECODER_TP_SIZE, tp=1, |
| 21 | # --enable-expert-parallel) instead of tensor |
| 22 | # parallel; the proxy round-robins decode across |
| 23 | # the DP replicas. Mirrors |
| 24 | # nixl_integration/run_accuracy_test.sh. |
| 25 | # GPU_MEMORY_UTILIZATION default 0.45 |
| 26 | # MAX_MODEL_LEN default 512 |
| 27 | # PREFILL_BLOCK_SIZE default 128 |
| 28 | # DECODE_BLOCK_SIZE default 128 |
| 29 | # CPU_BYTES default 209715200 (200 MB) |
| 30 | # VLLM_SERVE_EXTRA_ARGS comma-separated extra args for vllm serve |
| 31 | # --decoder-first toggle decoder-first proxy mode |
| 32 | # |
| 33 | # Examples: |
| 34 | # bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh |
| 35 | # NUM_PREFILL_INSTANCES=2 NUM_DECODE_INSTANCES=2 \ |
| 36 | # bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh |
| 37 | # bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh --decoder-first |
| 38 | # DP_EP=1 DECODER_TP_SIZE=2 \ |
| 39 | # bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh |
| 40 | |
| 41 | set -xe |
| 42 | |
| 43 | # --------------------------------------------------------------------------- |
| 44 | # Args |
| 45 | # --------------------------------------------------------------------------- |
| 46 | DECODER_FIRST="false" |
| 47 | |
| 48 | while [[ $# -gt 0 ]]; do |
| 49 | case $1 in |
| 50 | --decoder-first) |
| 51 | DECODER_FIRST="true" |
| 52 | shift 1 |
| 53 | ;; |
| 54 | *) |
| 55 | echo "Unknown option $1" |
| 56 | echo "Usage: $0 [--decoder-first]" |
| 57 | exit 1 |
| 58 | ;; |
| 59 | esac |
| 60 | done |
| 61 | |
| 62 | # --------------------------------------------------------------------------- |
| 63 | # Models |
| 64 | # --------------------------------------------------------------------------- |
| 65 | MODEL_NAMES=${MODEL_NAMES:-} |
| 66 | if [[ -n "$MODEL_NAMES" ]]; then |
| 67 | # shellcheck disable=SC2206 |
| 68 | MODELS=($MODEL_NAMES) |
| 69 | else |
| 70 | MODELS=( |
| 71 | "meta-llama/Llama-3.2-1B-Instruct" |
| 72 | ) |
| 73 | fi |
| 74 | |
| 75 | # --------------------------------------------------------------------------- |
| 76 | # Defaults |
| 77 | # --------------------------------------------------------------------------- |
| 78 | NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} |
| 79 | NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} |
| 80 | PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1} |
| 81 | DECODER_TP_SIZE=${DECODER_TP_SIZE:-1} |
| 82 | # DP-EP attention mode (see header). When DP_EP is set the decoder uses |
| 83 | # data-parallel replicas (dp=DECODER_TP_SIZE) that the proxy round-robins over. |
| 84 | DP_EP=${DP_EP:-} |
| 85 | if [[ -n "$DP_EP" ]]; then |
| 86 | DECODER_DP_SIZE=${DECODER_TP_SIZE} |
| 87 | else |
| 88 | DECODER_DP_SIZE=1 |
| 89 | fi |
| 90 | GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.45} |
| 91 | MAX_MODEL_LEN=${MAX_MODEL_LEN:-512} |
| 92 | PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128} |
| 93 | DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128} |
| 94 | CPU_BYTES=${CPU_BYTES:-209715200} |
| 95 | VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} |
| 96 | |
| 97 | # Base ports — per-instance offsets layered on top. |
| 98 | PREFILL_HTTP_BASE=8100 |
| 99 | DECODE_HTTP_BASE=8200 |
| 100 | PREFILL_PD_BASE=7777 |
| 101 | DECODE_PD_BASE=$((PREFILL_PD_BASE + NUM_PREFILL_INSTANCES)) |
| 102 | PROXY_PORT=8192 |
| 103 | P2P_HOST=127.0.0.1 |
| 104 | |
| 105 | # --------------------------------------------------------------------------- |
| 106 | # Resolve repo root + venv (works in .venv and /workspace/venv pods) |
| 107 | # --------------------------------------------------------------------------- |
| 108 | SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" |
| 109 | GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../../.." && pwd -P)}" |
| 110 | |
| 111 | if [[ -z "${VLLM_BIN:-}" ]]; then |
| 112 | if [[ -x "${GIT_ROOT}/.venv/bin/vllm" ]]; then |
| 113 | VLLM_BIN="${GIT_ROOT}/.venv/bin/vllm" |
| 114 | elif [[ -x "/workspace/venv/bin/vllm" ]]; then |
| 115 | VLLM_BIN="/workspace/venv/bin/vllm" |
| 116 | else |
| 117 | VLLM_BIN="$(command -v vllm)" |
| 118 | fi |
| 119 | fi |
| 120 | if [[ -z "${PYTHON_BIN:-}" ]]; then |
| 121 | if [[ -x "${GIT_ROOT}/.venv/bin/python" ]]; then |
| 122 | PYTHON_BIN="${GIT_ROOT}/.venv/bin/python" |
| 123 | elif [[ -x "/workspace/venv/bin/python" ]]; then |
| 124 | PYTHON_BIN="/workspace/venv/bin/python" |
| 125 | else |
| 126 | PYTHON_BIN="$(command -v python3 || command -v python)" |
| 127 | fi |
| 128 | fi |
| 129 | echo "Using vllm: ${VLLM_BIN}" |
| 130 | echo "Using python: ${PYTHON_BIN}" |
| 131 | |
| 132 | SMI_BIN=$(command -v nvidia-smi || command -v rocm-smi || echo "") |
| 133 | |
| 134 | # Trap SIGINT/SIGTERM/EXIT to kill background jobs. |
| 135 | trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT |
| 136 | |
| 137 | # --------------------------------------------------------------------------- |
| 138 | # Helpers |
| 139 | # --------------------------------------------------------------------------- |
| 140 | wait_for_server() { |
| 141 | local port=$1 |
| 142 | timeout 1200 bash -c " |
| 143 | until curl -s localhost:${port}/v1/completions > /dev/null; do |
| 144 | sleep 1 |
| 145 | done" && return 0 || return 1 |
| 146 | } |
| 147 | |
| 148 | cleanup_instances() { |
| 149 | echo "Cleaning up any running vLLM / proxy instances..." |
| 150 | pkill -f "vllm serve" || true |
| 151 | pkill -f "p2p_connector_proxy.py" || true |
| 152 | sleep 2 |
| 153 | } |
| 154 | |
| 155 | get_num_gpus() { |
| 156 | if [[ "$SMI_BIN" == *"nvidia"* ]]; then |
| 157 | $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l |
| 158 | elif [[ "$SMI_BIN" == *"rocm"* ]]; then |
| 159 | $SMI_BIN -l | grep -c GPU |
| 160 | else |
| 161 | echo "1" |
| 162 | fi |
| 163 | } |
| 164 | |
| 165 | # Build the OffloadingConnector kv-transfer-config for a given PD port. |
| 166 | # Mirrors deploy_local.sh:131. |
| 167 | build_kv_config() { |
| 168 | local pd_port=$1 |
| 169 | printf '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"spec_name":"TieringOffloadingSpec","cpu_bytes_to_use":%s,"secondary_tiers":[{"type":"p2p","host":"%s","port":%s}]}}' \ |
| 170 | "${CPU_BYTES}" "${P2P_HOST}" "${pd_port}" |
| 171 | } |
| 172 | |
| 173 | # --------------------------------------------------------------------------- |
| 174 | # Per-model run |
| 175 | # --------------------------------------------------------------------------- |
| 176 | run_tests_for_model() { |
| 177 | local model_name=$1 |
| 178 | echo "================================" |
| 179 | echo "Testing model: $model_name" |
| 180 | echo " prefillers=${NUM_PREFILL_INSTANCES} (tp=${PREFILLER_TP_SIZE})" |
| 181 | if [[ -n "$DP_EP" ]]; then |
| 182 | echo " decoders=${NUM_DECODE_INSTANCES} (dp-ep=${DECODER_TP_SIZE}, tp=1)" |
| 183 | else |
| 184 | echo " decoders=${NUM_DECODE_INSTANCES} (tp=${DECODER_TP_SIZE})" |
| 185 | fi |
| 186 | echo " decoder_first=${DECODER_FIRST}" |
| 187 | echo "================================" |
| 188 | |
| 189 | PREFILL_HOSTS=() |
| 190 | PREFILL_PORTS=() |
| 191 | PREFILL_PD_PORTS=() |
| 192 | DECODE_HOSTS=() |
| 193 | DECODE_PORTS=() |
| 194 | DECODE_PD_PORTS=() |
| 195 | |
| 196 | local num_gpus |
| 197 | num_gpus=$(get_num_gpus) |
| 198 | local next_gpu=0 |
| 199 | |
| 200 | # ---- Prefillers ---- |
| 201 | for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do |
| 202 | local gpu_id=$((i * PREFILLER_TP_SIZE % num_gpus)) |
| 203 | local cuda_devs="${gpu_id}" |
| 204 | for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do |
| 205 | cuda_devs="${cuda_devs},$(((gpu_id + j) % num_gpus))" |
| 206 | done |
| 207 | next_gpu=$(((gpu_id + PREFILLER_TP_SIZE) % num_gpus)) |
| 208 | |
| 209 | local http_port=$((PREFILL_HTTP_BASE + i)) |
| 210 | local pd_port=$((PREFILL_PD_BASE + i)) |
| 211 | local kv_cfg |
| 212 | kv_cfg=$(build_kv_config "${pd_port}") |
| 213 | |
| 214 | echo "Prefiller $i: gpu=[${cuda_devs}] http=${http_port} pd=${pd_port}" |
| 215 | |
| 216 | BASE_CMD="CUDA_VISIBLE_DEVICES=${cuda_devs} \ |
| 217 | PYTHONHASHSEED=42 \ |
| 218 | ${VLLM_BIN} serve ${model_name} \ |
| 219 | --port ${http_port} \ |
| 220 | --enforce-eager \ |
| 221 | --block-size ${PREFILL_BLOCK_SIZE} \ |
| 222 | --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \ |
| 223 | --max-model-len ${MAX_MODEL_LEN} \ |
| 224 | --tensor-parallel-size ${PREFILLER_TP_SIZE} \ |
| 225 | --kv-transfer-config '${kv_cfg}'" |
| 226 | |
| 227 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 228 | IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" |
| 229 | for arg in "${extra_args[@]}"; do |
| 230 | BASE_CMD="${BASE_CMD} $arg" |
| 231 | done |
| 232 | fi |
| 233 | |
| 234 | eval "${BASE_CMD} &" |
| 235 | |
| 236 | PREFILL_HOSTS+=("${P2P_HOST}") |
| 237 | PREFILL_PORTS+=("${http_port}") |
| 238 | PREFILL_PD_PORTS+=("${pd_port}") |
| 239 | done |
| 240 | |
| 241 | # ---- Decoders ---- |
| 242 | for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do |
| 243 | local gpu_id=$(((next_gpu + i * DECODER_TP_SIZE) % num_gpus)) |
| 244 | local cuda_devs="${gpu_id}" |
| 245 | for (( j=1; j < DECODER_TP_SIZE; j++ )); do |
| 246 | cuda_devs="${cuda_devs},$(((gpu_id + j) % num_gpus))" |
| 247 | done |
| 248 | |
| 249 | local http_port=$((DECODE_HTTP_BASE + i)) |
| 250 | local pd_port=$((DECODE_PD_BASE + i)) |
| 251 | local kv_cfg |
| 252 | kv_cfg=$(build_kv_config "${pd_port}") |
| 253 | |
| 254 | echo "Decoder $i: gpu=[${cuda_devs}] http=${http_port} pd=${pd_port}" |
| 255 | |
| 256 | BASE_CMD="CUDA_VISIBLE_DEVICES=${cuda_devs} \ |
| 257 | PYTHONHASHSEED=42 \ |
| 258 | ${VLLM_BIN} serve ${model_name} \ |
| 259 | --port ${http_port} \ |
| 260 | --enforce-eager \ |
| 261 | --block-size ${DECODE_BLOCK_SIZE} \ |
| 262 | --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \ |
| 263 | --max-model-len ${MAX_MODEL_LEN} \ |
| 264 | --kv-transfer-config '${kv_cfg}'" |
| 265 | |
| 266 | # DP-EP attention mode: data-parallel + expert-parallel decode replicas |
| 267 | # (dp=DECODER_TP_SIZE, tp=1) instead of tensor parallel. Mirrors |
| 268 | # nixl_integration/run_accuracy_test.sh. |
| 269 | if [[ -z "$DP_EP" ]]; then |
| 270 | BASE_CMD="${BASE_CMD} --tensor-parallel-size ${DECODER_TP_SIZE}" |
| 271 | else |
| 272 | echo "DP-EP Attention enabled, deploying decoder with dp=${DECODER_TP_SIZE} and tp=1" |
| 273 | BASE_CMD="${BASE_CMD} --data-parallel-size ${DECODER_TP_SIZE} \ |
| 274 | --tensor-parallel-size 1 --enable-expert-parallel" |
| 275 | fi |
| 276 | |
| 277 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 278 | IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" |
| 279 | for arg in "${extra_args[@]}"; do |
| 280 | BASE_CMD="${BASE_CMD} $arg" |
| 281 | done |
| 282 | fi |
| 283 | |
| 284 | eval "${BASE_CMD} &" |
| 285 | |
| 286 | DECODE_HOSTS+=("${P2P_HOST}") |
| 287 | DECODE_PORTS+=("${http_port}") |
| 288 | DECODE_PD_PORTS+=("${pd_port}") |
| 289 | done |
| 290 | |
| 291 | # ---- Wait for HTTP readiness ---- |
| 292 | for port in "${PREFILL_PORTS[@]}"; do |
| 293 | echo "Waiting for prefill instance on port $port to start..." |
| 294 | wait_for_server "$port" |
| 295 | done |
| 296 | for port in "${DECODE_PORTS[@]}"; do |
| 297 | echo "Waiting for decode instance on port $port to start..." |
| 298 | wait_for_server "$port" |
| 299 | done |
| 300 | |
| 301 | # ---- Proxy ---- |
| 302 | # The proxy currently advertises a single prefiller PD address to decoders. |
| 303 | # For the 1xM and matched NxM common cases the first prefiller's PD coords |
| 304 | # are the right pick; multi-prefiller PD round-robin is a follow-up. |
| 305 | PROXY_CMD="${PYTHON_BIN} ${SCRIPT_DIR}/p2p_connector_proxy.py \ |
| 306 | --port ${PROXY_PORT} \ |
| 307 | --host ${P2P_HOST} \ |
| 308 | --prefiller-hosts ${PREFILL_HOSTS[*]} \ |
| 309 | --prefiller-ports ${PREFILL_PORTS[*]} \ |
| 310 | --decoder-hosts ${DECODE_HOSTS[*]} \ |
| 311 | --decoder-ports ${DECODE_PORTS[*]} \ |
| 312 | --p2p-connector-host ${P2P_HOST} \ |
| 313 | --p2p-connector-port ${PREFILL_PD_PORTS[0]} \ |
| 314 | --decoder-p2p-connector-host ${P2P_HOST} \ |
| 315 | --decoder-p2p-connector-port ${DECODE_PD_PORTS[0]} \ |
| 316 | --decoder-dp-size ${DECODER_DP_SIZE}" |
| 317 | |
| 318 | if [[ "${DECODER_FIRST}" == "true" ]]; then |
| 319 | PROXY_CMD="${PROXY_CMD} --decoder-first" |
| 320 | fi |
| 321 | |
| 322 | echo "Starting proxy: ${PROXY_CMD}" |
| 323 | eval "${PROXY_CMD} &" |
| 324 | |
| 325 | sleep 5 |
| 326 | |
| 327 | # ---- Run accuracy test (reused from nixl_integration) ---- |
| 328 | echo "Running tests for $model_name" |
| 329 | TEST_MODEL=$model_name "${PYTHON_BIN}" -m pytest -s -x \ |
| 330 | "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_accuracy.py" |
| 331 | |
| 332 | cleanup_instances |
| 333 | sleep 3 |
| 334 | } |
| 335 | |
| 336 | # --------------------------------------------------------------------------- |
| 337 | # Drive |
| 338 | # --------------------------------------------------------------------------- |
| 339 | for model in "${MODELS[@]}"; do |
| 340 | run_tests_for_model "$model" |
| 341 | done |
| 342 | |
| 343 | echo "All tests completed!" |