A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 07:15 UTC
[CI] Add Nixl+OffloadingConnector e2e integration tests (#39200)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 10, 2026 at 15:40 UTC
c9dddc144b01bbbb3ee70f529575aeda95ed1778
2 files changed
+219
.buildkite/test_areas/distributed.yaml
+14
@@ -268,6 +268,20 @@ steps:
268
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
269
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
270
271
+- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
272
+ timeout_in_minutes: 30
273
+ working_dir: "/vllm-workspace/tests"
274
+ num_devices: 2
275
+ source_file_dependencies:
276
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
277
+ - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
278
+ - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
279
+ - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
280
+ - tests/v1/kv_connector/nixl_integration/
281
+ commands:
282
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
283
+ - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
284
+
285
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
286
timeout_in_minutes: 30
287
device: a100
tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
new
+205
@@ -0,0 +1,205 @@
1
+#!/bin/bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# Integration accuracy test for MultiConnector (NixlConnector + OffloadingConnector).
6
+#
7
+# Launches a P/D setup where both prefill and decode instances use MultiConnector
8
+# wrapping NixlConnector and OffloadingConnector, then runs gsm8k accuracy via
9
+# test_accuracy.py.
10
+#
11
+# By default runs two configurations:
12
+# 1. Normal KV layout (NixlConnector without cross-layer blocks)
13
+# 2. Cross-layer KV layout (NixlConnector with enable_cross_layers_blocks)
14
+#
15
+# Usage:
16
+# bash tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
17
+#
18
+# Environment variables:
19
+# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B)
20
+# GPU_MEMORY_UTILIZATION - GPU memory fraction (default: 0.6)
21
+# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve
22
+# SKIP_CROSS_LAYERS - set to 1 to skip the cross-layer layout test
23
+# SKIP_NORMAL_LAYOUT - set to 1 to skip the normal layout test
24
+set -xe
25
+
26
+# ── Configuration ────────────────────────────────────────────────────────
27
+
28
+MODEL_NAMES=${MODEL_NAMES:-}
29
+if [[ -n "$MODEL_NAMES" ]]; then
30
+ MODELS=("$MODEL_NAMES")
31
+else
32
+ MODELS=("Qwen/Qwen3-0.6B")
33
+fi
34
+
35
+GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.6}
36
+BLOCK_SIZE=${BLOCK_SIZE:-128}
37
+VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
38
+
39
+GIT_ROOT=$(git rev-parse --show-toplevel)
40
+SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
41
+
42
+# ── KV transfer configs ─────────────────────────────────────────────────
43
+
44
+# Normal layout: OffloadingConnector prefers cross-layer but NixlConnector
45
+# does not, so MultiConnector.prefer_cross_layer_blocks = False.
46
+KV_CONFIG_NORMAL='{
47
+ "kv_connector":"MultiConnector",
48
+ "kv_role":"kv_both",
49
+ "kv_connector_extra_config":{
50
+ "connectors":[
51
+ {"kv_connector":"NixlConnector","kv_role":"kv_both"},
52
+ {"kv_connector":"OffloadingConnector","kv_role":"kv_both",
53
+ "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}}
54
+ ]
55
+ }
56
+}'
57
+# Remove whitespace for CLI safety
58
+KV_CONFIG_NORMAL=$(echo "$KV_CONFIG_NORMAL" | tr -d '[:space:]')
59
+
60
+# Cross-layer layout: both connectors prefer cross-layer blocks.
61
+KV_CONFIG_CROSS_LAYERS='{
62
+ "kv_connector":"MultiConnector",
63
+ "kv_role":"kv_both",
64
+ "kv_connector_extra_config":{
65
+ "connectors":[
66
+ {"kv_connector":"NixlConnector","kv_role":"kv_both",
67
+ "kv_connector_extra_config":{"enable_cross_layers_blocks":"True"}},
68
+ {"kv_connector":"OffloadingConnector","kv_role":"kv_both",
69
+ "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}}
70
+ ]
71
+ }
72
+}'
73
+KV_CONFIG_CROSS_LAYERS=$(echo "$KV_CONFIG_CROSS_LAYERS" | tr -d '[:space:]')
74
+
75
+# ── Helpers ──────────────────────────────────────────────────────────────
76
+
77
+trap 'kill $(jobs -pr) 2>/dev/null' SIGINT SIGTERM EXIT
78
+
79
+wait_for_server() {
80
+ local port=$1
81
+ timeout 1200 bash -c "
82
+ until curl -s localhost:${port}/v1/completions > /dev/null; do
83
+ sleep 1
84
+ done" && return 0 || return 1
85
+}
86
+
87
+cleanup_instances() {
88
+ echo "Cleaning up any running vLLM instances..."
89
+ pkill -f "vllm serve" || true
90
+ sleep 2
91
+}
92
+
93
+get_num_gpus() {
94
+ if [[ "$SMI_BIN" == *"nvidia"* ]]; then
95
+ $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l
96
+ elif [[ "$SMI_BIN" == *"rocm"* ]]; then
97
+ $SMI_BIN -l | grep -c GPU
98
+ else
99
+ echo "1"
100
+ fi
101
+}
102
+
103
+# ── Run tests for one model with a given KV config ───────────────────────
104
+
105
+run_tests_for_model() {
106
+ local model_name=$1
107
+ local kv_config=$2
108
+ local label=$3
109
+
110
+ echo "================================================================"
111
+ echo "Testing model: $model_name ($label)"
112
+ echo "KV config: $kv_config"
113
+ echo "================================================================"
114
+
115
+ local PREFILL_PORT=8100
116
+ local DECODE_PORT=8200
117
+ local PREFILL_GPU=0
118
+ local DECODE_GPU=1
119
+ local PREFILL_SIDE_CHANNEL_PORT=5559
120
+ local DECODE_SIDE_CHANNEL_PORT=5659
121
+
122
+ # ── Start prefill instance ──
123
+ echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
124
+ BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \
125
+ VLLM_KV_CACHE_LAYOUT='HND' \
126
+ UCX_NET_DEVICES=all \
127
+ VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
128
+ vllm serve $model_name \
129
+ --port $PREFILL_PORT \
130
+ --enforce-eager \
131
+ --block-size ${BLOCK_SIZE} \
132
+ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
133
+ --tensor-parallel-size 1 \
134
+ --kv-transfer-config '$kv_config'"
135
+
136
+ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
137
+ IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
138
+ for arg in "${extra_args[@]}"; do
139
+ BASE_CMD="${BASE_CMD} $arg"
140
+ done
141
+ fi
142
+ eval "$BASE_CMD &"
143
+
144
+ # ── Start decode instance ──
145
+ echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
146
+ BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \
147
+ VLLM_KV_CACHE_LAYOUT='HND' \
148
+ UCX_NET_DEVICES=all \
149
+ VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
150
+ vllm serve $model_name \
151
+ --port $DECODE_PORT \
152
+ --enforce-eager \
153
+ --block-size ${BLOCK_SIZE} \
154
+ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
155
+ --tensor-parallel-size 1 \
156
+ --kv-transfer-config '$kv_config'"
157
+
158
+ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
159
+ IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
160
+ for arg in "${extra_args[@]}"; do
161
+ BASE_CMD="${BASE_CMD} $arg"
162
+ done
163
+ fi
164
+ eval "$BASE_CMD &"
165
+
166
+ # ── Wait for servers ──
167
+ echo "Waiting for prefill instance on port $PREFILL_PORT to start..."
168
+ wait_for_server "$PREFILL_PORT"
169
+ echo "Waiting for decode instance on port $DECODE_PORT to start..."
170
+ wait_for_server "$DECODE_PORT"
171
+
172
+ # ── Start proxy ──
173
+ PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port 8192"
174
+ PROXY_CMD+=" --prefiller-hosts localhost"
175
+ PROXY_CMD+=" --prefiller-ports $PREFILL_PORT"
176
+ PROXY_CMD+=" --decoder-hosts localhost"
177
+ PROXY_CMD+=" --decoder-ports $DECODE_PORT"
178
+
179
+ echo "Starting proxy server with command: $PROXY_CMD"
180
+ $PROXY_CMD &
181
+ sleep 5
182
+
183
+ # ── Run accuracy test ──
184
+ echo "Running accuracy tests for $model_name ($label)"
185
+ TEST_MODEL=$model_name python3 -m pytest -s -x \
186
+ "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_accuracy.py
187
+
188
+ # ── Cleanup ──
189
+ cleanup_instances
190
+ sleep 3
191
+}
192
+
193
+# ── Main ─────────────────────────────────────────────────────────────────
194
+
195
+for model in "${MODELS[@]}"; do
196
+ if [[ -z "${SKIP_NORMAL_LAYOUT:-}" ]]; then
197
+ run_tests_for_model "$model" "$KV_CONFIG_NORMAL" "MultiConnector normal layout"
198
+ fi
199
+
200
+ if [[ -z "${SKIP_CROSS_LAYERS:-}" ]]; then
201
+ run_tests_for_model "$model" "$KV_CONFIG_CROSS_LAYERS" "MultiConnector cross-layer layout"
202
+ fi
203
+done
204
+
205
+echo "All MultiConnector accuracy tests passed!"