[CI] Mooncake PD integration tests (#46844)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Aug 3, 2026 at 15:55 UTC
8ba87e01813271fec7841b6bec63595aa1698a91
9 files changed
+413
-4
.buildkite/scripts/install-kv-connectors.sh
+19
-3
@@ -13,7 +13,7 @@ REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv
13
14
uv pip install --system -r "${REQUIREMENTS_FILE}"
15
16
-NIXL_METADATA=$(python3 - <<'PY'
16
+KV_METADATA=$(python3 - <<'PY'
17
import importlib.metadata as metadata
18
19
import torch
@@ -22,10 +22,16 @@ cuda_version = torch.version.cuda
22
if cuda_version is None:
23
raise SystemExit("torch.version.cuda is not set")
24
25
-print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
25
+try:
26
+ mooncake_version = metadata.version("mooncake-transfer-engine")
27
+except metadata.PackageNotFoundError:
28
+ mooncake_version = ""
29
+
30
+print(cuda_version.split(".", 1)[0], metadata.version("nixl"), mooncake_version)
31
PY
32
)
28
-read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
33
+read -r CUDA_MAJOR NIXL_VERSION MOONCAKE_VERSION <<<"${KV_METADATA}"
34
+MOONCAKE_VERSION="${MOONCAKE_VERSION:-}"
35
36
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
37
# matching this CI image so nixl_ep_cpp links against the available libcudart.
@@ -42,3 +48,13 @@ for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
48
version = "not installed"
49
print(f"{package_name}: {version}")
50
PY
51
+
52
+# The default mooncake-transfer-engine PyPI wheel is built against CUDA 12; its
53
+# engine.so links libcudart.so.12, absent from the CUDA 13 runtime image. On a
54
+# CUDA 13 image, swap it for the cuda13 variant (same version), which links
55
+# libcudart.so.13. Both expose the `mooncake` package, so uninstall the CUDA 12
56
+# build first to avoid a clash.
57
+if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then
58
+ uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true
59
+ uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
60
+fi
.buildkite/test_areas/disaggregated_mooncake.yaml
new
+21
@@ -0,0 +1,21 @@
1
+group: Disaggregated Mooncake
2
+depends_on:
3
+ - image-build
4
+steps:
5
+- label: Distributed MooncakeConnector PD accuracy (4 GPUs)
6
+ key: distributed-mooncakeconnector-pd-accuracy-4-gpus
7
+ timeout_in_minutes: 10
8
+ working_dir: "/vllm-workspace/tests"
9
+ # Runs on b200-k8s because it is the only NVIDIA pool that can carry mooncake
10
+ # transfers: its pods are the only ones granted IPC_LOCK, which RDMA needs to
11
+ # pin memory, and the only ones using host networking. Elsewhere topology
12
+ # discovery finds no usable HCA and the tcp fallback then advertises the
13
+ # docker bridge address, so every transfer fails to connect.
14
+ device: b200-k8s
15
+ num_devices: 4
16
+ source_file_dependencies:
17
+ - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
18
+ - tests/v1/kv_connector/mooncake_integration/
19
+ commands:
20
+ - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
21
+ - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
docs/features/mooncake_connector_usage.md
+1
@@ -60,6 +60,7 @@ Now you can send requests to the proxy server through port 8000.
60
61
- **num_workers**: Size of thread pool for one prefiller worker to transfer KV caches by mooncake. (default 10)
62
- **mooncake_protocol**: Mooncake connector protocol. (default "rdma")
63
+- **device_name**: Comma-separated whitelist of RDMA devices (e.g. `"mlx5_0,mlx5_1"`) to restrict topology discovery to. Empty discovers every device. Useful on hosts exposing a mix of InfiniBand and RoCE ports, where both peers must settle on the same link layer.
64
65
## Example Scripts/Code
66
requirements/kv_connectors.txt
+3
-1
@@ -3,4 +3,6 @@ lmcache >= 0.3.9
3
# until a fixed newer release is verified for runtime images.
4
cupy-cuda13x < 14.1.0
5
nixl == 1.3.1
6
-mooncake-transfer-engine >= 0.3.8
6
+# CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for
7
+# the mooncake-transfer-engine-cuda13 variant of the same version.
8
+mooncake-transfer-engine >= 0.3.12
tests/v1/kv_connector/mooncake_integration/__init__.py
tests/v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
new
+53
@@ -0,0 +1,53 @@
1
+#!/usr/bin/env bash
2
+set -euo pipefail
3
+
4
+# Utility to run Mooncake P/D integration tests sequentially with varying TP
5
+# configurations.
6
+SCRIPT="v1/kv_connector/mooncake_integration/run_accuracy_test.sh"
7
+IMPORT_CANARY="v1/kv_connector/mooncake_integration/test_mooncake_imports.py"
8
+
9
+echo "=== Running Mooncake import canary ==="
10
+python3 -m pytest -s -x "${IMPORT_CANARY}"
11
+
12
+# Symmetric TP configurations (default model: Qwen/Qwen3-0.6B).
13
+tp_configs=(
14
+ "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=1"
15
+ "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2"
16
+)
17
+
18
+configs=("${tp_configs[@]}")
19
+
20
+run_tests() {
21
+ local label=$1
22
+ local extra_args=$2
23
+
24
+ echo "=== Running tests (${label}) ==="
25
+ for cfg in "${configs[@]}"; do
26
+ local -a cfg_parts extra_args_parts
27
+ read -r -a cfg_parts <<< "$cfg"
28
+ read -r -a extra_args_parts <<< "$extra_args"
29
+
30
+ echo "-> Running with ${cfg} ${extra_args:+and ${extra_args}}"
31
+ # Use 'env' to safely set variables without eval
32
+ # keep argv splitting safe and SC2086-clean via arrays.
33
+ if ! env "${cfg_parts[@]}" bash "${SCRIPT}" "${extra_args_parts[@]}"; then
34
+ echo "❌ Test failed for config: ${cfg} ${extra_args:+(${extra_args})}"
35
+ exit 1
36
+ fi
37
+ done
38
+ echo "✅ All ${label} tests passed!"
39
+}
40
+
41
+# Set backend
42
+label="default backend"
43
+cmdline_args=""
44
+if [[ -n "${ATTENTION_BACKEND:-}" ]]; then
45
+ echo "ATTENTION_BACKEND is set, running with --attention-backend ${ATTENTION_BACKEND}"
46
+ label="${ATTENTION_BACKEND} backend"
47
+ cmdline_args=" --attention-backend ${ATTENTION_BACKEND} "
48
+else
49
+ echo "running with default attention backend"
50
+fi
51
+
52
+# Run tests
53
+run_tests "${label}" "${cmdline_args}"
tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh
new
+227
@@ -0,0 +1,227 @@
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
+# Function to clean up previous instances
82
+cleanup_instances() {
83
+ echo "Cleaning up any running vLLM instances..."
84
+ pkill -f "vllm serve" || true
85
+ sleep 2
86
+}
87
+
88
+get_num_gpus() {
89
+ if [[ "$SMI_BIN" == *"nvidia"* ]]; then
90
+ $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l
91
+ elif [[ "$SMI_BIN" == *"rocm"* ]]; then
92
+ $SMI_BIN -l | grep -c GPU
93
+ else
94
+ # works for non-cuda platforms,
95
+ # assuming at least 1 device and
96
+ # let system to decide which card to use
97
+ echo "1"
98
+ fi
99
+}
100
+
101
+# Function to run tests for a specific model
102
+run_tests_for_model() {
103
+ local model_name=$1
104
+ echo "================================"
105
+ echo "Testing model: $model_name"
106
+ echo "================================"
107
+
108
+ # Proxy args: repeated --prefill URL BOOTSTRAP_PORT and --decode URL.
109
+ PROXY_ARGS=()
110
+
111
+ # Start prefill instances
112
+ for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
113
+ # Calculate GPU ID - we'll distribute across available GPUs
114
+ GPU_ID=$((i % $(get_num_gpus)))
115
+ NEXT_GPU=${GPU_ID}
116
+ # If PREFILLER_TP_SIZE is more than 1
117
+ for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do
118
+ NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus)))
119
+ GPU_ID="${GPU_ID},${NEXT_GPU}"
120
+ done
121
+
122
+ # Calculate port number (base port + instance number)
123
+ PORT=$((8100 + i))
124
+ # Bootstrap server port for this prefiller. Avoid clash across instances.
125
+ BOOTSTRAP_PORT=$((8998 + i))
126
+
127
+ echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT, bootstrap $BOOTSTRAP_PORT"
128
+
129
+ BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
130
+ VLLM_MOONCAKE_BOOTSTRAP_PORT=$BOOTSTRAP_PORT \
131
+ vllm serve $model_name \
132
+ --port $PORT \
133
+ --enforce-eager \
134
+ --block-size ${BLOCK_SIZE} \
135
+ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
136
+ --tensor-parallel-size $PREFILLER_TP_SIZE \
137
+ --kv-transfer-config '$KV_CONFIG_P'"
138
+ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
139
+ IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
140
+ for arg in "${extra_args[@]}"; do
141
+ BASE_CMD="${BASE_CMD} $arg"
142
+ done
143
+ fi
144
+
145
+ # Add attention backend config if specified
146
+ if [[ -n "$ATTENTION_BACKEND" ]]; then
147
+ BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
148
+ fi
149
+
150
+ eval "$BASE_CMD &"
151
+
152
+ PROXY_ARGS+=(--prefill "http://localhost:${PORT}" "$BOOTSTRAP_PORT")
153
+ PREFILL_PORTS+=("$PORT")
154
+ done
155
+
156
+ # Start decode instances
157
+ for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
158
+ # Calculate GPU ID - distribute starting from after prefill GPUs
159
+ GPU_ID=$(((i + NEXT_GPU + 1) % $(get_num_gpus)))
160
+ # If DECODER_TP_SIZE is more than 1
161
+ for (( j=1; j < DECODER_TP_SIZE; j++ )); do
162
+ NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus)))
163
+ GPU_ID="${GPU_ID},${NEXT_GPU}"
164
+ done
165
+ # Calculate port number (base port + instance number)
166
+ PORT=$((8200 + i))
167
+
168
+ echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
169
+
170
+ BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
171
+ vllm serve $model_name \
172
+ --port $PORT \
173
+ --enforce-eager \
174
+ --block-size ${BLOCK_SIZE} \
175
+ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
176
+ --tensor-parallel-size $DECODER_TP_SIZE \
177
+ --kv-transfer-config '$KV_CONFIG_D'"
178
+ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
179
+ IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
180
+ for arg in "${extra_args[@]}"; do
181
+ BASE_CMD="${BASE_CMD} $arg"
182
+ done
183
+ fi
184
+
185
+ # Add attention backend config if specified
186
+ if [[ -n "$ATTENTION_BACKEND" ]]; then
187
+ BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
188
+ fi
189
+
190
+ eval "$BASE_CMD &"
191
+
192
+ PROXY_ARGS+=(--decode "http://localhost:${PORT}")
193
+ DECODE_PORTS+=("$PORT")
194
+ done
195
+
196
+ # Wait for all instances to start
197
+ for PORT in "${PREFILL_PORTS[@]}"; do
198
+ echo "Waiting for prefill instance on port $PORT to start..."
199
+ wait_for_server "$PORT"
200
+ done
201
+
202
+ for PORT in "${DECODE_PORTS[@]}"; do
203
+ echo "Waiting for decode instance on port $PORT to start..."
204
+ wait_for_server "$PORT"
205
+ done
206
+
207
+ # Start the proxy server
208
+ echo "Starting proxy server on port 8192"
209
+ python3 "${GIT_ROOT}/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py" \
210
+ --port 8192 "${PROXY_ARGS[@]}" &
211
+
212
+ # Wait for the proxy to start
213
+ sleep 5
214
+
215
+ # Run lm eval for this model
216
+ echo "Running tests for $model_name"
217
+ TEST_MODEL=$model_name python3 -m pytest -s -x \
218
+ "${GIT_ROOT}/tests/v1/kv_connector/mooncake_integration/test_accuracy.py"
219
+
220
+ # Clean up before running next model
221
+ cleanup_instances
222
+ sleep 3
223
+}
224
+
225
+run_tests_for_model "$MODEL"
226
+
227
+echo "All tests completed!"
tests/v1/kv_connector/mooncake_integration/test_accuracy.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
+
5
+import lm_eval
6
+import openai
7
+
8
+BASE_URL = "http://localhost:8192/v1"
9
+NUM_CONCURRENT = 100
10
+TASK = "gsm8k"
11
+FILTER = "exact_match,strict-match"
12
+RTOL = 0.05
13
+
14
+# Model-specific expected values
15
+EXPECTED_VALUES = {
16
+ "Qwen/Qwen3-0.6B": 0.41,
17
+}
18
+
19
+SIMPLE_PROMPT = (
20
+ "The best part about working on vLLM is that I got to meet so many people across "
21
+ "various different organizations like UCB, Google, and Meta which means",
22
+)
23
+
24
+# Get model name from environment variable
25
+MODEL_NAME = os.environ.get("TEST_MODEL", "Qwen/Qwen3-0.6B")
26
+
27
+
28
+def run_simple_prompt():
29
+ client = openai.OpenAI(api_key="EMPTY", base_url=BASE_URL)
30
+ completion = client.completions.create(model=MODEL_NAME, prompt=SIMPLE_PROMPT)
31
+
32
+ print("-" * 50)
33
+ print(f"Completion results for {MODEL_NAME}:")
34
+ print(completion)
35
+ print("-" * 50)
36
+
37
+
38
+def test_accuracy():
39
+ """Run the end to end accuracy test."""
40
+ run_simple_prompt()
41
+
42
+ model_args = (
43
+ f"model={MODEL_NAME},"
44
+ f"base_url={BASE_URL}/completions,"
45
+ f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
46
+ )
47
+ results = lm_eval.simple_evaluate(
48
+ model="local-completions",
49
+ model_args=model_args,
50
+ tasks=TASK,
51
+ )
52
+
53
+ measured_value = results["results"][TASK][FILTER]
54
+ expected_value = EXPECTED_VALUES.get(MODEL_NAME)
55
+
56
+ print(f"Measured accuracy value: {measured_value}\n")
57
+ if expected_value is None:
58
+ print(
59
+ f"Warning: No expected value found for {MODEL_NAME}. "
60
+ "Skipping accuracy check."
61
+ )
62
+ return
63
+
64
+ assert (
65
+ measured_value - RTOL < expected_value
66
+ and measured_value + RTOL > expected_value
67
+ ), f"Expected: {expected_value} | Measured: {measured_value}"
tests/v1/kv_connector/mooncake_integration/test_mooncake_imports.py
new
+22
@@ -0,0 +1,22 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Mooncake import canary for CUDA wheel selection."""
4
+
5
+import importlib.metadata as metadata
6
+
7
+
8
+def test_mooncake_engine_imports():
9
+ """Fail here rather than as "Mooncake is not available" at engine startup.
10
+
11
+ The PyPI mooncake-transfer-engine wheel is built against CUDA 12, so on the
12
+ CUDA 13 image install-kv-connectors.sh replaces it with the cuda13 variant.
13
+ A mismatch surfaces as an ImportError on the compiled extension.
14
+ """
15
+ for package_name in ("mooncake-transfer-engine", "mooncake-transfer-engine-cuda13"):
16
+ try:
17
+ version = metadata.version(package_name)
18
+ except metadata.PackageNotFoundError:
19
+ version = "not installed"
20
+ print(f"{package_name}: {version}")
21
+
22
+ import mooncake.engine # noqa: F401