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 7, 01:30 UTC
[XPU] Add test script of PD disaggregation (#36434)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
liuzhenwei committed
Mar 9, 2026 at 13:50 UTC
1bc9c77f6d324bf7b9253b0c78626fbc50286bfb
1 file changed
+174
tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh
new
+174
@@ -0,0 +1,174 @@
1
+#!/bin/bash
2
+set -e
3
+
4
+# Hosts / ports
5
+PREFILL_HOST=${PREFILL_HOST:-"localhost"}
6
+PREFILL_PORT=${PREFILL_PORT:-8100}
7
+PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577}
8
+DECODE_HOST=${DECODE_HOST:-"localhost"}
9
+DECODE_PORT=${DECODE_PORT:-8200}
10
+PROXY_HOST=${PROXY_HOST:-"localhost"}
11
+PROXY_PORT=${PROXY_PORT:-8192}
12
+BASELINE_HOST=${BASELINE_HOST:-"localhost"}
13
+BASELINE_PORT=${BASELINE_PORT:-9290}
14
+
15
+# Model to run.
16
+MODEL_NAME=${MODEL_NAME:-"Qwen/Qwen3-0.6B"}
17
+MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024}
18
+BLOCK_SIZE=${BLOCK_SIZE:-64}
19
+PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
20
+DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
21
+KV_BUFFER_DEVICE=${KV_BUFFER_DEVICE:-"xpu"}
22
+GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8}
23
+
24
+generate_affinity_mask() {
25
+ local count=$1
26
+ local start=${2:-0}
27
+ local mask=""
28
+ local i
29
+
30
+ for ((i=0; i<count; i++)); do
31
+ local device=$((start + i))
32
+ if [[ -z "${mask}" ]]; then
33
+ mask="${device}"
34
+ else
35
+ mask="${mask},${device}"
36
+ fi
37
+ done
38
+
39
+ echo "${mask}"
40
+}
41
+
42
+PREFILLER_ZE_AFFINITY_MASK=${PREFILLER_ZE_AFFINITY_MASK:-$(generate_affinity_mask "${PREFILLER_TP_SIZE}" 0)}
43
+DECODER_ZE_AFFINITY_MASK=${DECODER_ZE_AFFINITY_MASK:-$(generate_affinity_mask "${DECODER_TP_SIZE}" "${PREFILLER_TP_SIZE}")}
44
+
45
+
46
+# execution env
47
+GIT_ROOT=$(git rev-parse --show-toplevel)
48
+EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
49
+
50
+OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.xpu_accuracy_test_outputs.txt"}
51
+
52
+# Trap the SIGINT signal (triggered by Ctrl+C)
53
+trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
54
+
55
+cleanup() {
56
+ echo "Cleaning up any running vLLM instances..."
57
+ pkill -f "vllm serve" || true
58
+ sleep 2
59
+}
60
+
61
+wait_for_server() {
62
+ local host=$1
63
+ local port=$2
64
+ timeout 1200 bash -c "
65
+ until curl -s ${host}:${port}/v1/completions > /dev/null; do
66
+ sleep 1
67
+ done" && return 0 || return 1
68
+}
69
+
70
+launch_baseline() {
71
+ BASELINE_BASE_CMD="
72
+ ZE_AFFINITY_MASK=0 \
73
+ VLLM_WORKER_MULTIPROC_METHOD=spawn \
74
+ VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \
75
+ --host ${BASELINE_HOST} \
76
+ --port ${BASELINE_PORT} \
77
+ --max-model-len ${MAX_MODEL_LEN}\
78
+ --seed 42 \
79
+ -tp 1 \
80
+ --block-size ${BLOCK_SIZE} \
81
+ --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \
82
+ --dtype float16 \
83
+ --enforce-eager"
84
+ echo ${BASELINE_BASE_CMD}
85
+ bash -c "${BASELINE_BASE_CMD}" &
86
+ sleep 10
87
+ wait_for_server ${BASELINE_HOST} ${BASELINE_PORT}
88
+}
89
+
90
+launch_pd() {
91
+ PREFILL_BASE_CMD="
92
+ ZE_AFFINITY_MASK=${PREFILLER_ZE_AFFINITY_MASK} \
93
+ VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
94
+ VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \
95
+ VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \
96
+ VLLM_WORKER_MULTIPROC_METHOD=spawn \
97
+ VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \
98
+ --host ${PREFILL_HOST} \
99
+ --port ${PREFILL_PORT} \
100
+ --max-model-len ${MAX_MODEL_LEN}\
101
+ --seed 42 \
102
+ --block-size ${BLOCK_SIZE} \
103
+ --enforce-eager \
104
+ --dtype float16 \
105
+ -tp ${PREFILLER_TP_SIZE} \
106
+ --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \
107
+ --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}'"
108
+
109
+
110
+ DECODE_BASE_CMD="
111
+ ZE_AFFINITY_MASK=${DECODER_ZE_AFFINITY_MASK} \
112
+ VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
113
+ VLLM_WORKER_MULTIPROC_METHOD=spawn \
114
+ VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \
115
+ --host ${DECODE_HOST} \
116
+ --port ${DECODE_PORT} \
117
+ --max-model-len ${MAX_MODEL_LEN}\
118
+ --seed 42 \
119
+ --block-size ${BLOCK_SIZE} \
120
+ --enforce-eager \
121
+ -tp ${DECODER_TP_SIZE} \
122
+ --dtype float16 \
123
+ --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \
124
+ --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}'"
125
+
126
+ echo ${PREFILL_BASE_CMD}
127
+ echo ${DECODE_BASE_CMD}
128
+ sleep 2
129
+
130
+ # execute on hosts
131
+ bash -c "${PREFILL_BASE_CMD}" &
132
+ bash -c "${DECODE_BASE_CMD}" &
133
+ sleep 1
134
+ wait_for_server ${PREFILL_HOST} ${PREFILL_PORT}
135
+ sleep 1
136
+ wait_for_server ${DECODE_HOST} ${DECODE_PORT}
137
+ sleep 1
138
+}
139
+
140
+launch_pd_proxy(){
141
+ PROXY_BASE_CMD="
142
+ python3 ${EXP_ROOT}/toy_proxy_server.py \
143
+ --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \
144
+ --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \
145
+ --host=${PROXY_HOST} --port ${PROXY_PORT}"
146
+ echo ${PROXY_BASE_CMD}
147
+ bash -c "${PROXY_BASE_CMD}" &
148
+ sleep 2
149
+}
150
+
151
+run_tests(){
152
+ local service_url=$1
153
+ local mode=$2
154
+ python3 ${EXP_ROOT}/test_disagg_accuracy.py --service_url=${service_url} --model_name=${MODEL_NAME} --mode=${mode} --file_name=${OUTPUT_FILE}
155
+}
156
+
157
+
158
+# run non-disagg. baseline & save outputs
159
+launch_baseline
160
+run_tests "http://${BASELINE_HOST}:${BASELINE_PORT}" "baseline"
161
+cleanup
162
+sleep 10
163
+
164
+
165
+# run disagg. & do exact-match with the outputs from baseline
166
+launch_pd
167
+launch_pd_proxy
168
+run_tests "http://${PROXY_HOST}:${PROXY_PORT}" "disagg"
169
+echo "-----P/D success----"
170
+
171
+rm ${OUTPUT_FILE}
172
+cleanup
173
+
174
+exit 0