| 1 | #!/bin/bash |
| 2 | set -xe |
| 3 | |
| 4 | # Parse command line arguments |
| 5 | KV_BUFFER_DEVICE="cuda" # Default to cuda |
| 6 | ATTENTION_BACKEND="" # Default to empty (use vllm default) |
| 7 | CROSS_LAYERS_BLOCKS="False" |
| 8 | |
| 9 | while [[ $# -gt 0 ]]; do |
| 10 | case $1 in |
| 11 | --kv_buffer_device) |
| 12 | KV_BUFFER_DEVICE="$2" |
| 13 | shift 2 |
| 14 | ;; |
| 15 | --attention-backend) |
| 16 | ATTENTION_BACKEND="$2" |
| 17 | shift 2 |
| 18 | ;; |
| 19 | --enable-cross-layers) |
| 20 | CROSS_LAYERS_BLOCKS="True" |
| 21 | shift 1 |
| 22 | ;; |
| 23 | *) |
| 24 | echo "Unknown option $1" |
| 25 | echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]" |
| 26 | exit 1 |
| 27 | ;; |
| 28 | esac |
| 29 | done |
| 30 | |
| 31 | echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE" |
| 32 | if [[ -n "$ATTENTION_BACKEND" ]]; then |
| 33 | echo "Using attention backend: $ATTENTION_BACKEND" |
| 34 | fi |
| 35 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 36 | echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS" |
| 37 | fi |
| 38 | |
| 39 | # Empty leaves the layout to the backends, for models that constrain it. |
| 40 | PREFILLER_KV_LAYOUT=${PREFILLER_KV_LAYOUT-"HND"} |
| 41 | DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT-"HND"} # Default to HND, optional NHD |
| 42 | if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then |
| 43 | KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"' |
| 44 | else |
| 45 | KV_CONFIG_HETERO_LAYOUT='' |
| 46 | fi |
| 47 | |
| 48 | if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then |
| 49 | # Cross-layer blocks are exercised via the block-outermost layout. |
| 50 | PREFILLER_KV_LAYOUT="BLHNC" |
| 51 | DECODER_KV_LAYOUT="BLHNC" |
| 52 | fi |
| 53 | |
| 54 | # Connector: default pull NixlConnector; NixlPushConnector enables PP prefill. |
| 55 | KV_CONNECTOR=${KV_CONNECTOR:-NixlConnector} |
| 56 | |
| 57 | # Build the kv-transfer-config for P and D |
| 58 | if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then |
| 59 | KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}'}' |
| 60 | KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}'}' |
| 61 | else |
| 62 | KV_CONFIG_P="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}" |
| 63 | KV_CONFIG_D="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}" |
| 64 | fi |
| 65 | |
| 66 | # Models to run |
| 67 | MODEL_NAMES=${MODEL_NAMES:-} |
| 68 | if [[ -n "$MODEL_NAMES" ]]; then |
| 69 | MODELS=("$MODEL_NAMES") |
| 70 | else |
| 71 | MODELS=( |
| 72 | "Qwen/Qwen3-0.6B" |
| 73 | ) |
| 74 | fi |
| 75 | |
| 76 | # Number of prefill and decode instances to create |
| 77 | NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} # Default to 1 |
| 78 | NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1 |
| 79 | PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1} |
| 80 | PREFILLER_PP_SIZE=${PREFILLER_PP_SIZE:-1} # >1 requires NixlPushConnector |
| 81 | DECODER_TP_SIZE=${DECODER_TP_SIZE:-1} |
| 82 | GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2} |
| 83 | PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128} |
| 84 | DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128} |
| 85 | ENFORCE_EAGER=${ENFORCE_EAGER:-1} |
| 86 | # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048) |
| 87 | VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} |
| 88 | # Pin concurrent prefiller and non-DP decoder engines to separate internal |
| 89 | # port windows. DP decoder ranks retain their existing internal port selection. |
| 90 | PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000} |
| 91 | DECODER_INTERNAL_PORT_BASE=${DECODER_INTERNAL_PORT_BASE:-30000} |
| 92 | INTERNAL_PORT_STRIDE=${INTERNAL_PORT_STRIDE:-100} |
| 93 | |
| 94 | # Resolve the repository root from the script location instead of `.git`. |
| 95 | # The ROCm CI image copies `/vllm-workspace` without the Git metadata, so |
| 96 | # `git rev-parse --show-toplevel` is not reliable at runtime. |
| 97 | SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" |
| 98 | GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" |
| 99 | |
| 100 | SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") |
| 101 | |
| 102 | # Trap the SIGINT signal (triggered by Ctrl+C) |
| 103 | trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT |
| 104 | |
| 105 | # Waits for vLLM to start. |
| 106 | wait_for_server() { |
| 107 | local port=$1 |
| 108 | timeout 1200 bash -c " |
| 109 | until curl -s localhost:${port}/v1/completions > /dev/null; do |
| 110 | sleep 1 |
| 111 | done" && return 0 || return 1 |
| 112 | } |
| 113 | |
| 114 | # Function to clean up previous instances |
| 115 | wait_for_gpu_memory_release() { |
| 116 | if [[ "$SMI_BIN" == *"rocm"* ]]; then |
| 117 | PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_memory_to_settle; wait_for_memory_to_settle()" |
| 118 | fi |
| 119 | } |
| 120 | |
| 121 | cleanup_instances() { |
| 122 | echo "Cleaning up any running vLLM instances..." |
| 123 | pkill -f "toy_proxy_server.py" || true |
| 124 | pkill -TERM -f "vllm serve" || true |
| 125 | sleep 3 |
| 126 | pkill -9 -f "vllm serve" || true |
| 127 | sleep 2 |
| 128 | wait_for_gpu_memory_release |
| 129 | } |
| 130 | |
| 131 | get_num_gpus() { |
| 132 | if [[ "$SMI_BIN" == *"nvidia"* ]]; then |
| 133 | $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l |
| 134 | elif [[ "$SMI_BIN" == *"rocm"* ]]; then |
| 135 | $SMI_BIN -l | grep -c GPU |
| 136 | else |
| 137 | # works for non-cuda platforms, |
| 138 | # assuming at least 1 device and |
| 139 | # let system to decide which card to use |
| 140 | echo "1" |
| 141 | fi |
| 142 | } |
| 143 | |
| 144 | # Function to run tests for a specific model |
| 145 | run_tests_for_model() { |
| 146 | local model_name=$1 |
| 147 | cleanup_instances |
| 148 | echo "================================" |
| 149 | echo "Testing model: $model_name" |
| 150 | echo "================================" |
| 151 | |
| 152 | # Arrays to store all hosts and ports |
| 153 | PREFILL_HOSTS=() |
| 154 | PREFILL_PORTS=() |
| 155 | DECODE_HOSTS=() |
| 156 | DECODE_PORTS=() |
| 157 | |
| 158 | # Start prefill instances |
| 159 | for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do |
| 160 | # Calculate GPU ID - we'll distribute across available GPUs |
| 161 | GPU_START=$((i % $(get_num_gpus))) |
| 162 | GPU_ID=$GPU_START |
| 163 | NEXT_GPU=$GPU_START |
| 164 | # Reserve TP*PP GPUs for the prefiller (TP shards across PP stages). |
| 165 | PREFILLER_WORLD_SIZE=$((PREFILLER_TP_SIZE * PREFILLER_PP_SIZE)) |
| 166 | for (( j=1; j < PREFILLER_WORLD_SIZE; j++ )); do |
| 167 | NEXT_GPU=$(((GPU_START + j) % $(get_num_gpus))) |
| 168 | GPU_ID="${GPU_ID},${NEXT_GPU}" |
| 169 | done |
| 170 | |
| 171 | # Calculate port number (base port + instance number) |
| 172 | PORT=$((8100 + i)) |
| 173 | # Calculate side channel port. Avoid clash with with TP workers. |
| 174 | SIDE_CHANNEL_PORT=$((5559 + i)) |
| 175 | INTERNAL_PORT=$((PREFILLER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE)) |
| 176 | |
| 177 | echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT" |
| 178 | |
| 179 | # Build the command with or without model-specific args |
| 180 | BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ |
| 181 | ${PREFILLER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$PREFILLER_KV_LAYOUT} \ |
| 182 | VLLM_PORT=$INTERNAL_PORT \ |
| 183 | UCX_NET_DEVICES=all \ |
| 184 | VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ |
| 185 | vllm serve $model_name \ |
| 186 | --port $PORT \ |
| 187 | --block-size ${PREFILL_BLOCK_SIZE} \ |
| 188 | --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ |
| 189 | --tensor-parallel-size $PREFILLER_TP_SIZE \ |
| 190 | --pipeline-parallel-size $PREFILLER_PP_SIZE \ |
| 191 | --kv-transfer-config '$KV_CONFIG_P'" |
| 192 | if [[ "$ENFORCE_EAGER" == "1" ]]; then |
| 193 | BASE_CMD="${BASE_CMD} --enforce-eager" |
| 194 | fi |
| 195 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 196 | IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" |
| 197 | for arg in "${extra_args[@]}"; do |
| 198 | BASE_CMD="${BASE_CMD} $arg" |
| 199 | done |
| 200 | fi |
| 201 | |
| 202 | # Add attention backend config if specified |
| 203 | if [[ -n "$ATTENTION_BACKEND" ]]; then |
| 204 | BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND" |
| 205 | fi |
| 206 | |
| 207 | |
| 208 | FULL_CMD="$BASE_CMD" |
| 209 | eval "$FULL_CMD &" |
| 210 | |
| 211 | # Store host and port for proxy configuration |
| 212 | PREFILL_HOSTS+=("localhost") |
| 213 | PREFILL_PORTS+=("$PORT") |
| 214 | done |
| 215 | |
| 216 | # Start decode instances |
| 217 | for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do |
| 218 | # Calculate GPU ID - we'll distribute across available GPUs, starting from after prefill GPUs |
| 219 | DECODE_START=$(((i + NEXT_GPU + 1) % $(get_num_gpus))) |
| 220 | GPU_ID=$DECODE_START |
| 221 | NEXT_GPU=$DECODE_START |
| 222 | # If DECODER_TP_SIZE is more than 1 |
| 223 | for (( j=1; j < DECODER_TP_SIZE; j++ )); do |
| 224 | NEXT_GPU=$(((DECODE_START + j) % $(get_num_gpus))) |
| 225 | GPU_ID="${GPU_ID},${NEXT_GPU}" |
| 226 | done |
| 227 | # Calculate port number (base port + instance number) |
| 228 | PORT=$((8200 + i)) |
| 229 | # Calculate side channel port |
| 230 | SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE)) |
| 231 | INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE)) |
| 232 | # For non-DP mode, set VLLM_PORT to pin the internal port; |
| 233 | # For DP mode, set VLLM_DP_MASTER_PORT instead to avoid race condition. |
| 234 | if [[ -z "${DP_EP:-}" ]]; then |
| 235 | DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT" |
| 236 | else |
| 237 | DECODER_INTERNAL_PORT_ENV="VLLM_DP_MASTER_PORT=$INTERNAL_PORT" |
| 238 | fi |
| 239 | |
| 240 | echo "Starting decode instance $i on GPU $GPU_ID, port $PORT" |
| 241 | |
| 242 | # Build the command with or without model-specific args |
| 243 | BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ |
| 244 | ${DECODER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT} \ |
| 245 | $DECODER_INTERNAL_PORT_ENV \ |
| 246 | UCX_NET_DEVICES=all \ |
| 247 | VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ |
| 248 | vllm serve $model_name \ |
| 249 | --port $PORT \ |
| 250 | --block-size ${DECODE_BLOCK_SIZE} \ |
| 251 | --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ |
| 252 | --kv-transfer-config '$KV_CONFIG_D'" |
| 253 | if [[ "$ENFORCE_EAGER" == "1" ]]; then |
| 254 | BASE_CMD="${BASE_CMD} --enforce-eager" |
| 255 | fi |
| 256 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 257 | IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" |
| 258 | for arg in "${extra_args[@]}"; do |
| 259 | BASE_CMD="${BASE_CMD} $arg" |
| 260 | done |
| 261 | fi |
| 262 | |
| 263 | # Add attention backend config if specified |
| 264 | if [[ -n "$ATTENTION_BACKEND" ]]; then |
| 265 | BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND" |
| 266 | fi |
| 267 | |
| 268 | |
| 269 | # DP-EP attention mode |
| 270 | if [[ -z "$DP_EP" ]]; then |
| 271 | BASE_CMD="${BASE_CMD} --tensor-parallel-size $DECODER_TP_SIZE" |
| 272 | else |
| 273 | echo "DP-EP Attention enabled, deploying with dp=DECODER_TP_SIZE and tp=1" |
| 274 | BASE_CMD="${BASE_CMD} --data-parallel-size $DECODER_TP_SIZE \ |
| 275 | --tensor-parallel-size 1 --enable-expert-parallel" |
| 276 | fi |
| 277 | |
| 278 | FULL_CMD="$BASE_CMD" |
| 279 | |
| 280 | eval "$FULL_CMD &" |
| 281 | |
| 282 | # Store host and port for proxy configuration |
| 283 | DECODE_HOSTS+=("localhost") |
| 284 | DECODE_PORTS+=("$PORT") |
| 285 | done |
| 286 | |
| 287 | # Wait for all instances to start |
| 288 | for PORT in "${PREFILL_PORTS[@]}"; do |
| 289 | echo "Waiting for prefill instance on port $PORT to start..." |
| 290 | wait_for_server "$PORT" |
| 291 | done |
| 292 | |
| 293 | for PORT in "${DECODE_PORTS[@]}"; do |
| 294 | echo "Waiting for decode instance on port $PORT to start..." |
| 295 | wait_for_server "$PORT" |
| 296 | done |
| 297 | |
| 298 | # Build the command for the proxy server with all the hosts and ports |
| 299 | PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port 8192" |
| 300 | |
| 301 | # Add all prefill hosts and ports |
| 302 | PROXY_CMD+=" --prefiller-hosts ${PREFILL_HOSTS[*]}" |
| 303 | PROXY_CMD+=" --prefiller-ports ${PREFILL_PORTS[*]}" |
| 304 | |
| 305 | # Add all decode hosts and ports |
| 306 | PROXY_CMD+=" --decoder-hosts ${DECODE_HOSTS[*]}" |
| 307 | PROXY_CMD+=" --decoder-ports ${DECODE_PORTS[*]}" |
| 308 | |
| 309 | # Start the proxy server |
| 310 | echo "Starting proxy server with command: $PROXY_CMD" |
| 311 | $PROXY_CMD & |
| 312 | |
| 313 | # Wait for the proxy to start |
| 314 | sleep 5 |
| 315 | |
| 316 | # Run lm eval for this model |
| 317 | echo "Running tests for $model_name" |
| 318 | TEST_MODEL=$model_name python3 -m pytest -s -x "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_accuracy.py |
| 319 | |
| 320 | # Clean up before running next model |
| 321 | cleanup_instances |
| 322 | } |
| 323 | |
| 324 | # Run tests for each model |
| 325 | for model in "${MODELS[@]}"; do |
| 326 | run_tests_for_model "$model" |
| 327 | done |
| 328 | |
| 329 | echo "All tests completed!" |