| 1 | #!/bin/bash |
| 2 | set -xe |
| 3 | |
| 4 | # Parse command line arguments |
| 5 | ATTENTION_BACKEND="" # Default to empty (use vllm default) |
| 6 | |
| 7 | while [[ $# -gt 0 ]]; do |
| 8 | case $1 in |
| 9 | --attention-backend) |
| 10 | ATTENTION_BACKEND="$2" |
| 11 | shift 2 |
| 12 | ;; |
| 13 | *) |
| 14 | echo "Unknown option $1" |
| 15 | echo "Usage: $0 [--attention-backend <backend>]" |
| 16 | exit 1 |
| 17 | ;; |
| 18 | esac |
| 19 | done |
| 20 | |
| 21 | if [[ -n "$ATTENTION_BACKEND" ]]; then |
| 22 | echo "Using attention backend: $ATTENTION_BACKEND" |
| 23 | fi |
| 24 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 25 | echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS" |
| 26 | fi |
| 27 | |
| 28 | # Mooncake discovers peers via its bootstrap server, so no side-channel |
| 29 | # or layout env vars are needed. |
| 30 | |
| 31 | # Comma-separated whitelist restricting RDMA topology discovery to named HCAs. |
| 32 | # This default is specific to the b200-k8s CI pool: those nodes expose four |
| 33 | # InfiniBand HCAs that have no network device next to a single RoCE port |
| 34 | # (mlx5_12). With all five visible the two sides can settle on different link |
| 35 | # layers, and the QP handshake then fails to reach RTR. Set to a different NIC |
| 36 | # on other hardware, or to the empty string to discover every device. |
| 37 | MOONCAKE_DEVICE_NAME=${MOONCAKE_DEVICE_NAME-mlx5_12} |
| 38 | kv_config() { |
| 39 | printf '{"kv_connector":"MooncakeConnector","kv_role":"%s",' "$1" |
| 40 | printf '"kv_connector_extra_config":{"device_name":"%s"}}' "${MOONCAKE_DEVICE_NAME}" |
| 41 | } |
| 42 | KV_CONFIG_P=$(kv_config kv_producer) |
| 43 | KV_CONFIG_D=$(kv_config kv_consumer) |
| 44 | |
| 45 | # Models to run |
| 46 | MODEL_NAME=${MODEL_NAME:-} |
| 47 | if [[ -n "$MODEL_NAME" ]]; then |
| 48 | MODEL="$MODEL_NAME" |
| 49 | else |
| 50 | MODEL="Qwen/Qwen3-0.6B" |
| 51 | fi |
| 52 | |
| 53 | # Number of prefill and decode instances to create |
| 54 | NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} # Default to 1 |
| 55 | NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1 |
| 56 | PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1} |
| 57 | DECODER_TP_SIZE=${DECODER_TP_SIZE:-1} |
| 58 | GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2} |
| 59 | BLOCK_SIZE=${BLOCK_SIZE:-128} |
| 60 | # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048) |
| 61 | VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} |
| 62 | |
| 63 | # Resolve the repository root from the script location instead of `.git`. |
| 64 | SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" |
| 65 | GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" |
| 66 | |
| 67 | SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") |
| 68 | |
| 69 | # Trap the SIGINT signal (triggered by Ctrl+C) |
| 70 | trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT |
| 71 | |
| 72 | # Waits for vLLM to start. |
| 73 | wait_for_server() { |
| 74 | local port=$1 |
| 75 | timeout 1200 bash -c " |
| 76 | until curl -s localhost:${port}/v1/completions > /dev/null; do |
| 77 | sleep 1 |
| 78 | done" && return 0 || return 1 |
| 79 | } |
| 80 | |
| 81 | # b200-k8s pods use host networking, so a fixed port can already be taken by |
| 82 | # another job on the node. Bind all the ports at once so they are distinct. |
| 83 | free_ports() { |
| 84 | python3 - "$1" <<'EOF' |
| 85 | import socket |
| 86 | import sys |
| 87 | |
| 88 | socks = [socket.socket() for _ in range(int(sys.argv[1]))] |
| 89 | for s in socks: |
| 90 | s.bind(("", 0)) |
| 91 | print(" ".join(str(s.getsockname()[1]) for s in socks)) |
| 92 | EOF |
| 93 | } |
| 94 | |
| 95 | # Function to clean up previous instances |
| 96 | cleanup_instances() { |
| 97 | echo "Cleaning up any running vLLM instances..." |
| 98 | pkill -f "vllm serve" || true |
| 99 | sleep 2 |
| 100 | } |
| 101 | |
| 102 | get_num_gpus() { |
| 103 | if [[ "$SMI_BIN" == *"nvidia"* ]]; then |
| 104 | $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l |
| 105 | elif [[ "$SMI_BIN" == *"rocm"* ]]; then |
| 106 | $SMI_BIN -l | grep -c GPU |
| 107 | else |
| 108 | # works for non-cuda platforms, |
| 109 | # assuming at least 1 device and |
| 110 | # let system to decide which card to use |
| 111 | echo "1" |
| 112 | fi |
| 113 | } |
| 114 | |
| 115 | # Function to run tests for a specific model |
| 116 | run_tests_for_model() { |
| 117 | local model_name=$1 |
| 118 | echo "================================" |
| 119 | echo "Testing model: $model_name" |
| 120 | echo "================================" |
| 121 | |
| 122 | # Proxy args: repeated --prefill URL BOOTSTRAP_PORT and --decode URL. |
| 123 | PROXY_ARGS=() |
| 124 | # Per prefiller: server and bootstrap. Per decoder: server. Then the proxy. |
| 125 | read -r -a PORTS <<< "$(free_ports $((2 * NUM_PREFILL_INSTANCES + NUM_DECODE_INSTANCES + 1)))" |
| 126 | PROXY_PORT=${PORTS[${#PORTS[@]} - 1]} |
| 127 | |
| 128 | # Start prefill instances |
| 129 | for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do |
| 130 | # Calculate GPU ID - we'll distribute across available GPUs |
| 131 | GPU_ID=$((i % $(get_num_gpus))) |
| 132 | NEXT_GPU=${GPU_ID} |
| 133 | # If PREFILLER_TP_SIZE is more than 1 |
| 134 | for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do |
| 135 | NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus))) |
| 136 | GPU_ID="${GPU_ID},${NEXT_GPU}" |
| 137 | done |
| 138 | |
| 139 | PORT=${PORTS[$((2 * i))]} |
| 140 | BOOTSTRAP_PORT=${PORTS[$((2 * i + 1))]} |
| 141 | |
| 142 | echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT, bootstrap $BOOTSTRAP_PORT" |
| 143 | |
| 144 | BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ |
| 145 | VLLM_MOONCAKE_BOOTSTRAP_PORT=$BOOTSTRAP_PORT \ |
| 146 | vllm serve $model_name \ |
| 147 | --port $PORT \ |
| 148 | --enforce-eager \ |
| 149 | --block-size ${BLOCK_SIZE} \ |
| 150 | --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ |
| 151 | --tensor-parallel-size $PREFILLER_TP_SIZE \ |
| 152 | --kv-transfer-config '$KV_CONFIG_P'" |
| 153 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 154 | IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" |
| 155 | for arg in "${extra_args[@]}"; do |
| 156 | BASE_CMD="${BASE_CMD} $arg" |
| 157 | done |
| 158 | fi |
| 159 | |
| 160 | # Add attention backend config if specified |
| 161 | if [[ -n "$ATTENTION_BACKEND" ]]; then |
| 162 | BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND" |
| 163 | fi |
| 164 | |
| 165 | eval "$BASE_CMD &" |
| 166 | |
| 167 | PROXY_ARGS+=(--prefill "http://localhost:${PORT}" "$BOOTSTRAP_PORT") |
| 168 | PREFILL_PORTS+=("$PORT") |
| 169 | done |
| 170 | |
| 171 | # Start decode instances |
| 172 | for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do |
| 173 | # Calculate GPU ID - distribute starting from after prefill GPUs |
| 174 | GPU_ID=$(((i + NEXT_GPU + 1) % $(get_num_gpus))) |
| 175 | # If DECODER_TP_SIZE is more than 1 |
| 176 | for (( j=1; j < DECODER_TP_SIZE; j++ )); do |
| 177 | NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus))) |
| 178 | GPU_ID="${GPU_ID},${NEXT_GPU}" |
| 179 | done |
| 180 | PORT=${PORTS[$((2 * NUM_PREFILL_INSTANCES + i))]} |
| 181 | |
| 182 | echo "Starting decode instance $i on GPU $GPU_ID, port $PORT" |
| 183 | |
| 184 | BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ |
| 185 | vllm serve $model_name \ |
| 186 | --port $PORT \ |
| 187 | --enforce-eager \ |
| 188 | --block-size ${BLOCK_SIZE} \ |
| 189 | --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ |
| 190 | --tensor-parallel-size $DECODER_TP_SIZE \ |
| 191 | --kv-transfer-config '$KV_CONFIG_D'" |
| 192 | if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then |
| 193 | IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" |
| 194 | for arg in "${extra_args[@]}"; do |
| 195 | BASE_CMD="${BASE_CMD} $arg" |
| 196 | done |
| 197 | fi |
| 198 | |
| 199 | # Add attention backend config if specified |
| 200 | if [[ -n "$ATTENTION_BACKEND" ]]; then |
| 201 | BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND" |
| 202 | fi |
| 203 | |
| 204 | eval "$BASE_CMD &" |
| 205 | |
| 206 | PROXY_ARGS+=(--decode "http://localhost:${PORT}") |
| 207 | DECODE_PORTS+=("$PORT") |
| 208 | done |
| 209 | |
| 210 | # Wait for all instances to start |
| 211 | for PORT in "${PREFILL_PORTS[@]}"; do |
| 212 | echo "Waiting for prefill instance on port $PORT to start..." |
| 213 | wait_for_server "$PORT" |
| 214 | done |
| 215 | |
| 216 | for PORT in "${DECODE_PORTS[@]}"; do |
| 217 | echo "Waiting for decode instance on port $PORT to start..." |
| 218 | wait_for_server "$PORT" |
| 219 | done |
| 220 | |
| 221 | # Start the proxy server |
| 222 | echo "Starting proxy server on port $PROXY_PORT" |
| 223 | python3 "${GIT_ROOT}/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py" \ |
| 224 | --port "$PROXY_PORT" "${PROXY_ARGS[@]}" & |
| 225 | |
| 226 | # Wait for the proxy to start |
| 227 | sleep 5 |
| 228 | |
| 229 | # Run lm eval for this model |
| 230 | echo "Running tests for $model_name" |
| 231 | TEST_MODEL=$model_name PROXY_PORT=$PROXY_PORT python3 -m pytest -s -x \ |
| 232 | "${GIT_ROOT}/tests/v1/kv_connector/mooncake_integration/test_accuracy.py" |
| 233 | |
| 234 | # Clean up before running next model |
| 235 | cleanup_instances |
| 236 | sleep 3 |
| 237 | } |
| 238 | |
| 239 | run_tests_for_model "$MODEL" |
| 240 | |
| 241 | echo "All tests completed!" |